From ed471072bba733214e15c212809097a108f142f9 Mon Sep 17 00:00:00 2001 From: William Valentin Date: Mon, 16 Feb 2026 20:56:08 -0800 Subject: [PATCH] feat(companion): add waitForIdle runtime drain helper --- README.md | 4 +- docs/plans/state.json | 16 ++++++ src/companion/index.ts | 1 + src/companion/platformClients.test.ts | 13 +++++ src/companion/platformClients.ts | 13 +++++ src/companion/runtimeClient.test.ts | 50 ++++++++++++++++++ src/companion/runtimeClient.ts | 74 +++++++++++++++++++++++++++ 7 files changed, 169 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 55daad0..72539e8 100644 --- a/README.md +++ b/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 diff --git a/docs/plans/state.json b/docs/plans/state.json index 6a2643b..f2d334f 100644 --- a/docs/plans/state.json +++ b/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", diff --git a/src/companion/index.ts b/src/companion/index.ts index 8c9b25a..bb4e1bf 100644 --- a/src/companion/index.ts +++ b/src/companion/index.ts @@ -12,6 +12,7 @@ export { CompanionHeartbeatLoop } from './heartbeatLoop.js'; export type { CompanionRuntimeClientOptions, + WaitForIdleOptions, CompanionEventHandler, CompanionTypedEventHandler, CompanionEventName, diff --git a/src/companion/platformClients.test.ts b/src/companion/platformClients.test.ts index 3c71bb5..deb8323 100644 --- a/src/companion/platformClients.test.ts +++ b/src/companion/platformClients.test.ts @@ -33,6 +33,7 @@ function createRuntimeMock(): { clearEventSubscriptions: ReturnType; listKnownEventNames: ReturnType; waitForEvent: ReturnType; + waitForIdle: ReturnType; waitForAgentStream: ReturnType; waitForAgentTyping: ReturnType; waitForContextWarning: ReturnType; @@ -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' }); diff --git a/src/companion/platformClients.ts b/src/companion/platformClients.ts index 5738df7..0650f13 100644 --- a/src/companion/platformClients.ts +++ b/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 { + 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 { + 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 { + return this.runtime.waitForIdle(options); + } + private resolveSessionId(sessionId?: string): string { const resolved = sessionId ?? this.defaultSessionId; if (!resolved) { diff --git a/src/companion/runtimeClient.test.ts b/src/companion/runtimeClient.test.ts index 2c8b653..145a287 100644 --- a/src/companion/runtimeClient.test.ts +++ b/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; diff --git a/src/companion/runtimeClient.ts b/src/companion/runtimeClient.ts index c1e6e8b..1cf0e11 100644 --- a/src/companion/runtimeClient.ts +++ b/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 = (data: TData) => void; export type CompanionEventPredicate = (data: TData) => boolean; @@ -607,6 +613,74 @@ export class CompanionRuntimeClient { return this.waitForEvent(COMPANION_EVENT_NAMES.contextWarning, options); } + waitForIdle(options?: WaitForIdleOptions): Promise { + 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((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); }