feat(companion): add waitForIdle runtime drain helper

README.md

@@ -1190,7 +1190,7 @@ Methods:
 - `system.capabilities` returns gateway protocol and node policy snapshot.
 
 Companion runtime helper:
-- `src/companion/runtimeClient.ts` provides a typed Node/WebSocket client for companion runtimes (macOS/iOS/Android workers) with wrappers for `node.register`, `node.capabilities.get`, `node.location.set/get`, `node.status.set`, `node.push_token.set`, `system.capabilities`, `system.nodes`, and canvas artifact RPCs (`canvas.put/get/list/delete/clear`), plus convenience helpers (`bootstrapNode`, optional `autoConnect`, `dispose()`) and event helpers (`subscribeEvents()`, `subscribeEvent()`, `subscribeAgentStream()`, `subscribeAgentTyping()`, `subscribeContextWarning()`, `waitForEvent()` with timeout/predicate/abort support, non-empty event-name validation, and deterministic teardown cancellation (including socket-close rejection), `waitForAnyEvent()` (with non-empty event-name-list validation), `waitForAgentStream()`, `waitForAgentTyping()`, `waitForContextWarning()`, `clearEventSubscriptions()`, `listKnownEventNames()`, `eventSubscriptionCount`) plus in-flight observability via `pendingRequestCount`, `pendingEventWaitCount`, and `hasPendingWork`.
+- `src/companion/runtimeClient.ts` provides a typed Node/WebSocket client for companion runtimes (macOS/iOS/Android workers) with wrappers for `node.register`, `node.capabilities.get`, `node.location.set/get`, `node.status.set`, `node.push_token.set`, `system.capabilities`, `system.nodes`, and canvas artifact RPCs (`canvas.put/get/list/delete/clear`), plus convenience helpers (`bootstrapNode`, optional `autoConnect`, `dispose()`, `waitForIdle()` for pending-work drain synchronization) and event helpers (`subscribeEvents()`, `subscribeEvent()`, `subscribeAgentStream()`, `subscribeAgentTyping()`, `subscribeContextWarning()`, `waitForEvent()` with timeout/predicate/abort support, non-empty event-name validation, and deterministic teardown cancellation (including socket-close rejection), `waitForAnyEvent()` (with non-empty event-name-list validation), `waitForAgentStream()`, `waitForAgentTyping()`, `waitForContextWarning()`, `clearEventSubscriptions()`, `listKnownEventNames()`, `eventSubscriptionCount`) plus in-flight observability via `pendingRequestCount`, `pendingEventWaitCount`, and `hasPendingWork`.
 - `src/companion/platformClients.ts` provides platform-focused wrappers:
   - `MacOSCompanionClient` (`platform: "macos"`, APNs push registration)
   - `IOSCompanionClient` (`platform: "ios"`, APNs push registration)
@@ -1201,7 +1201,7 @@ Companion runtime helper:
   - optional `defaultSessionId` for canvas helper calls so `sessionId` can be omitted per call
   - lifecycle passthroughs for connection state/teardown (`connected`, `dispose(code?, reason?)`)
   - stream passthrough helpers (`subscribeEvents`, `subscribeEvent`, `clearEventSubscriptions`, `listKnownEventNames`, `eventSubscriptionCount`, `subscribeAgentStream/Typing/ContextWarning`, `waitForEvent`, `waitForAnyEvent`, `waitForAgentStream/Typing/ContextWarning`)
-  - runtime observability passthroughs (`pendingRequestCount`, `pendingEventWaitCount`, `hasPendingWork`, `connected`)
+  - runtime observability/control passthroughs (`pendingRequestCount`, `pendingEventWaitCount`, `hasPendingWork`, `connected`, `waitForIdle()`)
 - `src/companion/heartbeatLoop.ts` provides `CompanionHeartbeatLoop` for periodic heartbeat scheduling (`publishHeartbeat`) with start/stop safety, optional interval jitter (`jitterRatio`) to spread load, `tickNow()` for manual sends, success/error hooks, loop observability (`successCount`, `lastSuccessAt`, `failureCount`, `lastFailure`, `getState()`), and optional auto-stop after repeated failures.
 
 ## Canvas / A2UI Foundation

docs/plans/state.json

@@ -900,6 +900,22 @@
       ],
       "test_status": "pnpm test:run src/companion/runtimeClient.test.ts src/companion/platformClients.test.ts src/companion/heartbeatLoop.test.ts src/companion/platformClients.integration.test.ts + pnpm typecheck passing"
     },
+    "companion-runtime-wait-for-idle-helper": {
+      "status": "completed",
+      "date": "2026-02-17",
+      "updated": "2026-02-17",
+      "summary": "Added runtime `waitForIdle()` helper (timeout/poll/abort aware) to await drain of pending RPC/event wait work, with platform passthrough methods and validation/lifecycle tests.",
+      "files_modified": [
+        "src/companion/runtimeClient.ts",
+        "src/companion/runtimeClient.test.ts",
+        "src/companion/platformClients.ts",
+        "src/companion/platformClients.test.ts",
+        "src/companion/index.ts",
+        "README.md",
+        "docs/plans/state.json"
+      ],
+      "test_status": "pnpm test:run src/companion/runtimeClient.test.ts src/companion/platformClients.test.ts src/companion/heartbeatLoop.test.ts src/companion/platformClients.integration.test.ts + pnpm typecheck passing"
+    },
     "browser-tools-activation-clarity": {
       "status": "completed",
       "date": "2026-02-17",

src/companion/index.ts

@@ -12,6 +12,7 @@ export { CompanionHeartbeatLoop } from './heartbeatLoop.js';
 
 export type {
   CompanionRuntimeClientOptions,
+  WaitForIdleOptions,
   CompanionEventHandler,
   CompanionTypedEventHandler,
   CompanionEventName,

src/companion/platformClients.test.ts

@@ -33,6 +33,7 @@ function createRuntimeMock(): {
   clearEventSubscriptions: ReturnType<typeof vi.fn>;
   listKnownEventNames: ReturnType<typeof vi.fn>;
   waitForEvent: ReturnType<typeof vi.fn>;
+  waitForIdle: ReturnType<typeof vi.fn>;
   waitForAgentStream: ReturnType<typeof vi.fn>;
   waitForAgentTyping: ReturnType<typeof vi.fn>;
   waitForContextWarning: ReturnType<typeof vi.fn>;
@@ -75,6 +76,7 @@ function createRuntimeMock(): {
   const clearEventSubscriptions = vi.fn(() => undefined);
   const listKnownEventNames = vi.fn(() => ['agent.stream', 'agent.typing', 'context_warning']);
   const waitForEvent = vi.fn(async () => ({ token: 'evented' }));
+  const waitForIdle = vi.fn(async () => undefined);
   const waitForAgentStream = vi.fn(async () => ({ token: 'streamed' }));
   const waitForAgentTyping = vi.fn(async () => ({ active: true }));
   const waitForContextWarning = vi.fn(async () => ({ thresholdPct: 75, estimatedPct: 90 }));
@@ -111,6 +113,7 @@ function createRuntimeMock(): {
     clearEventSubscriptions,
     listKnownEventNames,
     waitForEvent,
+    waitForIdle,
     waitForAgentStream,
     waitForAgentTyping,
     waitForContextWarning,
@@ -159,6 +162,7 @@ function createRuntimeMock(): {
     clearEventSubscriptions,
     listKnownEventNames,
     waitForEvent,
+    waitForIdle,
     waitForAgentStream,
     waitForAgentTyping,
     waitForContextWarning,
@@ -333,6 +337,15 @@ describe('platform companion clients', () => {
     expect(client.hasPendingWork).toBe(mock.hasPendingWork);
   });
 
+  it('platform waitForIdle forwards options to runtime client', async () => {
+    const mock = createRuntimeMock();
+    const client = new IOSCompanionClient({ runtime: mock.runtime, nodeId: 'ios-node' });
+
+    await client.waitForIdle({ timeoutMs: 250, pollIntervalMs: 10 });
+
+    expect(mock.waitForIdle).toHaveBeenCalledWith({ timeoutMs: 250, pollIntervalMs: 10 });
+  });
+
   it('macOS client forwards canvas methods to runtime client', async () => {
     const mock = createRuntimeMock();
     const client = new MacOSCompanionClient({ runtime: mock.runtime, nodeId: 'mac-node' });

src/companion/platformClients.ts

@@ -22,6 +22,7 @@ import type {
   SetNodeLocationInput,
   SystemCapabilitiesResult,
   SystemNodesResult,
+  WaitForIdleOptions,
 } from './runtimeClient.js';
 import {
   CompanionHeartbeatLoop,
@@ -325,6 +326,10 @@ export class MacOSCompanionClient {
     return this.runtime.waitForContextWarning(options);
   }
 
+  waitForIdle(options?: WaitForIdleOptions): Promise<void> {
+    return this.runtime.waitForIdle(options);
+  }
+
   private resolveSessionId(sessionId?: string): string {
     const resolved = sessionId ?? this.defaultSessionId;
     if (!resolved) {
@@ -579,6 +584,10 @@ export class IOSCompanionClient {
     return this.runtime.waitForContextWarning(options);
   }
 
+  waitForIdle(options?: WaitForIdleOptions): Promise<void> {
+    return this.runtime.waitForIdle(options);
+  }
+
   private resolveSessionId(sessionId?: string): string {
     const resolved = sessionId ?? this.defaultSessionId;
     if (!resolved) {
@@ -831,6 +840,10 @@ export class AndroidCompanionClient {
     return this.runtime.waitForContextWarning(options);
   }
 
+  waitForIdle(options?: WaitForIdleOptions): Promise<void> {
+    return this.runtime.waitForIdle(options);
+  }
+
   private resolveSessionId(sessionId?: string): string {
     const resolved = sessionId ?? this.defaultSessionId;
     if (!resolved) {

src/companion/runtimeClient.test.ts

@@ -643,6 +643,56 @@ describe('CompanionRuntimeClient', () => {
     expect(client.hasPendingWork).toBe(false);
   });
 
+  it('waitForIdle resolves immediately when no work is pending', async () => {
+    const client = new CompanionRuntimeClient({
+      url: 'ws://127.0.0.1:1',
+    });
+
+    await expect(client.waitForIdle()).resolves.toBeUndefined();
+  });
+
+  it('waitForIdle validates pollIntervalMs option', () => {
+    const client = new CompanionRuntimeClient({
+      url: 'ws://127.0.0.1:1',
+    });
+
+    expect(() => client.waitForIdle({ pollIntervalMs: 0 })).toThrow(
+      'pollIntervalMs must be a positive number',
+    );
+  });
+
+  it('waitForIdle resolves after pending event waiters are cleared', async () => {
+    const client = new CompanionRuntimeClient({
+      url: 'ws://127.0.0.1:1',
+    });
+    const pendingWait = client.waitForEvent('agent.stream', { timeoutMs: 10_000 }).catch(() => undefined);
+    expect(client.pendingEventWaitCount).toBe(1);
+
+    const idle = client.waitForIdle({ timeoutMs: 1_000, pollIntervalMs: 5 });
+    setTimeout(() => {
+      client.clearEventSubscriptions();
+    }, 20);
+
+    await expect(idle).resolves.toBeUndefined();
+    await pendingWait;
+    expect(client.pendingEventWaitCount).toBe(0);
+    expect(client.hasPendingWork).toBe(false);
+  });
+
+  it('waitForIdle rejects on timeout while work remains pending', async () => {
+    const client = new CompanionRuntimeClient({
+      url: 'ws://127.0.0.1:1',
+    });
+    const pendingWait = client.waitForEvent('agent.stream', { timeoutMs: 10_000 }).catch(() => undefined);
+
+    await expect(
+      client.waitForIdle({ timeoutMs: 40, pollIntervalMs: 5 }),
+    ).rejects.toThrow('Timed out waiting for runtime idle state');
+
+    client.clearEventSubscriptions();
+    await pendingWait;
+  });
+
   it('connects and performs node registration + capability discovery', async () => {
     if (!LISTEN_ALLOWED) {
       return;

src/companion/runtimeClient.ts

@@ -41,6 +41,12 @@ export interface CompanionRuntimeClientOptions {
   websocketFactory?: (url: string) => WebSocket;
 }
 
+export interface WaitForIdleOptions {
+  timeoutMs?: number;
+  pollIntervalMs?: number;
+  signal?: AbortSignal;
+}
+
 export type CompanionEventHandler = (event: string, data: unknown) => void;
 export type CompanionTypedEventHandler<TData = unknown> = (data: TData) => void;
 export type CompanionEventPredicate<TData = unknown> = (data: TData) => boolean;
@@ -607,6 +613,74 @@ export class CompanionRuntimeClient {
     return this.waitForEvent<TData>(COMPANION_EVENT_NAMES.contextWarning, options);
   }
 
+  waitForIdle(options?: WaitForIdleOptions): Promise<void> {
+    const pollIntervalMs = options?.pollIntervalMs ?? 25;
+    if (!Number.isFinite(pollIntervalMs) || pollIntervalMs <= 0) {
+      throw new Error('pollIntervalMs must be a positive number');
+    }
+    if (!this.hasPendingWork) {
+      return Promise.resolve();
+    }
+    const timeoutMs = options?.timeoutMs ?? this.requestTimeoutMs;
+    const signal = options?.signal;
+
+    return new Promise<void>((resolve, reject) => {
+      let settled = false;
+      let timeout: NodeJS.Timeout | null = null;
+      let poll: NodeJS.Timeout | null = null;
+      let abortCleanup: (() => void) | null = null;
+
+      const cleanup = () => {
+        if (timeout) {
+          clearTimeout(timeout);
+          timeout = null;
+        }
+        if (poll) {
+          clearInterval(poll);
+          poll = null;
+        }
+        if (abortCleanup) {
+          abortCleanup();
+          abortCleanup = null;
+        }
+      };
+
+      const finish = (fn: () => void) => {
+        if (settled) {
+          return;
+        }
+        settled = true;
+        cleanup();
+        fn();
+      };
+
+      const check = () => {
+        if (!this.hasPendingWork) {
+          finish(() => resolve());
+        }
+      };
+
+      timeout = setTimeout(() => {
+        finish(() => reject(new Error('Timed out waiting for runtime idle state')));
+      }, timeoutMs);
+      poll = setInterval(check, pollIntervalMs);
+      check();
+
+      if (signal) {
+        const onAbort = () => {
+          finish(() => reject(new Error('Aborted while waiting for runtime idle state')));
+        };
+        signal.addEventListener('abort', onAbort, { once: true });
+        abortCleanup = () => {
+          signal.removeEventListener('abort', onAbort);
+        };
+        if (signal.aborted) {
+          onAbort();
+        }
+      }
+    });
+  }
+
   listKnownEventNames(): CompanionEventName[] {
     return Object.values(COMPANION_EVENT_NAMES);
   }