@johnhenry/wsh 0.16.0 → 0.17.0

package/src/control-listener.mjs

@@ -0,0 +1,130 @@
+/**
+ * Control-message subscription for the helper classes that compose their
+ * own wire messages — `WshMcpBridge` and `WshFileTransfer`.
+ *
+ * Both accept "a WshClient, or any object exposing sendControl() plus a
+ * way to observe inbound control messages", and both used to inline the
+ * same three-way registration. The third branch of it — the fallback for
+ * an object whose only hook is a settable `onControl` property, which is
+ * what a bare `WshTransport` is — wrapped that property and **never
+ * unwrapped it**:
+ *
+ * ```js
+ * const prev = client.onControl;
+ * client.onControl = (msg) => { prev?.(msg); listener(msg); };
+ * ```
+ *
+ * The matching `cleanup()` handled `removeControlListener` and
+ * `_controlListeners` and did nothing at all for this branch. So every
+ * `discover()` / `call()` / `list()` left another wrapper permanently in
+ * the chain. Measured against a real `WshTransport`, the number of frames
+ * an inbound control message traverses to reach the connection's own
+ * handler grew one-for-one with the number of operations — 1, 5, 20 and
+ * 50 operations gave 5, 9, 24 and 54 frames — with no ceiling.
+ *
+ * Attaching once per subscriber and detaching on cleanup keeps that flat.
+ */
+
+/**
+ * Subscribe `listener` to a client's inbound control messages.
+ *
+ * Supports, in order of preference:
+ *  1. `addControlListener()` / `removeControlListener()` — what `WshClient`
+ *     exposes, and the only shape that can deregister cleanly.
+ *  2. A `_controlListeners` array.
+ *  3. A settable `onControl` property (a bare `WshTransport`), wrapped
+ *     once and restored on detach.
+ *
+ * @param {object} client
+ * @param {function(object): void} listener
+ * @returns {function(): void} detach — idempotent; safe to call twice.
+ */
+export function attachControlListener(client, listener) {
+  if (typeof client.addControlListener === 'function') {
+    client.addControlListener(listener);
+    let detached = false;
+    return () => {
+      if (detached) return;
+      detached = true;
+      client.removeControlListener?.(listener);
+    };
+  }
+
+  if (Array.isArray(client._controlListeners)) {
+    client._controlListeners.push(listener);
+    let detached = false;
+    return () => {
+      if (detached) return;
+      detached = true;
+      const idx = client._controlListeners.indexOf(listener);
+      if (idx !== -1) client._controlListeners.splice(idx, 1);
+    };
+  }
+
+  // Fallback: wrap `onControl`. Restoring `prev` is only correct while our
+  // wrapper is still the installed handler — if something else has wrapped
+  // it since, putting `prev` back would drop that handler. In that case the
+  // wrapper stays but is made inert, so the chain stops growing per
+  // operation even though this one link cannot be spliced out.
+  const prev = client.onControl;
+  let live = true;
+  const wrapper = (msg) => {
+    prev?.(msg);
+    if (live) listener(msg);
+  };
+  client.onControl = wrapper;
+
+  let detached = false;
+  return () => {
+    if (detached) return;
+    detached = true;
+    live = false;
+    if (client.onControl === wrapper) client.onControl = prev;
+  };
+}
+
+/**
+ * Wait for the first inbound control message matching `predicate`.
+ *
+ * @param {object} client
+ * @param {function(object): boolean} predicate
+ * @param {number} timeoutMs
+ * @param {string} label - Used in the timeout message.
+ * @returns {Promise<object>}
+ */
+export function waitForControlMessage(client, predicate, timeoutMs, label) {
+  return new Promise((resolve, reject) => {
+    let settled = false;
+    let detach = () => {};
+    let timer = null;
+
+    const cleanup = () => {
+      settled = true;
+      if (timer !== null) clearTimeout(timer);
+      detach();
+    };
+
+    // `detach` is assigned after attach() returns, but a client that
+    // dispatches synchronously from inside attach() would fire the
+    // listener first — hence the `settled` guard rather than relying on
+    // detach() having been wired up.
+    detach = attachControlListener(client, (msg) => {
+      if (settled) return;
+      let matches = false;
+      try {
+        matches = predicate(msg);
+      } catch {
+        matches = false;
+      }
+      if (!matches) return;
+      cleanup();
+      resolve(msg);
+    });
+
+    timer = setTimeout(() => {
+      if (settled) return;
+      cleanup();
+      reject(new Error(`${label} timed out after ${timeoutMs}ms`));
+    }, timeoutMs);
+  });
+}

package/src/e2e-frame.mjs

@@ -138,13 +138,44 @@ export async function sealFrame(key, sessionId, roleTag, counter, plaintext) {
  * @param {{nonce: Uint8Array, ciphertext: Uint8Array}} frame
  * @returns {Promise<Uint8Array>} plaintext
  */
-export async function openFrame(key, sessionId, expectedCounter, { nonce, ciphertext }) {
+export async function openFrame(key, sessionId, expectedCounter, { nonce, ciphertext, expectedRoleTag }) {
   if (!key || typeof key !== 'object') {
     throw new Error('e2e-frame: openFrame requires a CryptoKey');
   }
   if (!(nonce instanceof Uint8Array) || nonce.length !== NONCE_LENGTH) {
     throw new Error(`e2e-frame: nonce must be a ${NONCE_LENGTH}-byte Uint8Array`);
   }
+  /*
+   * The role tag is REQUIRED, and a caller that omits it gets an error rather
+   * than a skipped check.
+   *
+   * The tag exists so the two directions of one session are structurally
+   * distinguishable -- it is the reason a colliding nonce is impossible when
+   * both directions share one AES-GCM key. WshSession computed the peer's
+   * expected tag and stored it, and nothing ever read it: `grep -rn
+   * e2eRecvRoleTag` returned the declaration and the assignment and no third
+   * line. So the receiver accepted a frame sealed with its OWN tag, and the
+   * relay this layer exists to distrust could echo a peer's sealed frame
+   * straight back at it. Both counters start at zero, so the first frame of a
+   * session reflects cleanly.
+   *
+   * Optional would have reproduced the defect: a security check nobody is
+   * forced to pass is the same as no check, and this file already had one --
+   * see TimestampProof.verify's ignored verifyFn for the same shape.
+   */
+  if (!(expectedRoleTag instanceof Uint8Array) || expectedRoleTag.length !== ROLE_TAG_LENGTH) {
+    throw new Error(
+      `e2e-frame: openFrame requires expectedRoleTag, a ${ROLE_TAG_LENGTH}-byte Uint8Array`
+    );
+  }
+  const actualRoleTag = nonce.subarray(COUNTER_LENGTH);
+  for (let i = 0; i < ROLE_TAG_LENGTH; i += 1) {
+    if (actualRoleTag[i] !== expectedRoleTag[i]) {
+      throw new Error(
+        'e2e-frame: role tag mismatch -- this frame was sealed by the wrong side of the session'
+      );
+    }
+  }
   const actualCounter = Number(new DataView(nonce.buffer, nonce.byteOffset, nonce.byteLength).getBigUint64(0, false));
   if (actualCounter !== expectedCounter) {
     throw new Error(

package/src/file-transfer.mjs

@@ -8,6 +8,7 @@
  */
 
 import { MSG, open, CHANNEL_KIND } from './messages.mjs';
+import { waitForControlMessage } from './control-listener.mjs';
 
 /** Default timeout for waiting on control messages (30 seconds). */
 const RESPONSE_TIMEOUT_MS = 30_000;
@@ -30,6 +31,10 @@ export class WshFileTransfer {
    *   - openStream(): open a new bidirectional stream
    *   - onControl: settable callback for incoming control messages
    *     (or a method to add a listener — see _waitForMessage)
+   *
+   *   Only `list()` needs the control-channel members, so they are checked
+   *   there rather than here: uploading and downloading through a client
+   *   that has `upload()`/`download()` and nothing else stays valid.
    */
   constructor(client) {
     if (!client) throw new Error('WshFileTransfer requires a client');
@@ -104,6 +109,11 @@ export class WshFileTransfer {
     if (!remotePath || typeof remotePath !== 'string') {
       throw new Error('remotePath is required');
     }
+    if (typeof this.#client.sendControl !== 'function' || typeof this.#client.openStream !== 'function') {
+      throw new TypeError(
+        'WshFileTransfer.list() requires a client exposing sendControl() and openStream() (e.g. WshClient)'
+      );
+    }
 
     // Use an exec channel to run ls
     const openMsg = open({
@@ -154,51 +164,7 @@ export class WshFileTransfer {
    * @returns {Promise<object>}
    */
   _waitForMessage(predicate, timeoutMs) {
-    return new Promise((resolve, reject) => {
-      let timer = null;
-      let settled = false;
-
-      const cleanup = () => {
-        settled = true;
-        if (timer !== null) clearTimeout(timer);
-        // Remove listener
-        if (this.#client.removeControlListener) {
-          this.#client.removeControlListener(listener);
-        } else if (this.#client._controlListeners) {
-          const idx = this.#client._controlListeners.indexOf(listener);
-          if (idx !== -1) this.#client._controlListeners.splice(idx, 1);
-        }
-      };
-
-      const listener = (msg) => {
-        if (settled) return;
-        if (predicate(msg)) {
-          cleanup();
-          resolve(msg);
-        }
-      };
-
-      // Register listener on the client
-      if (this.#client.addControlListener) {
-        this.#client.addControlListener(listener);
-      } else if (this.#client._controlListeners) {
-        this.#client._controlListeners.push(listener);
-      } else {
-        // Fallback: wrap existing onControl
-        const prev = this.#client.onControl;
-        this.#client.onControl = (msg) => {
-          prev?.(msg);
-          listener(msg);
-        };
-      }
-
-      timer = setTimeout(() => {
-        if (!settled) {
-          cleanup();
-          reject(new Error(`Timeout waiting for response (${timeoutMs}ms)`));
-        }
-      }, timeoutMs);
-    });
+    return waitForControlMessage(this.#client, predicate, timeoutMs, 'File transfer response');
   }
 
   /**

package/src/index.d.ts

@@ -228,6 +228,16 @@ export const AUTH_METHOD: {
 /** Protocol version string. */
 export const PROTOCOL_VERSION: 'wsh-v1';
 
+/**
+ * ServerHello feature advertising that McpResult echoes McpCall's call_id.
+ *
+ * Compare against `WshClient.hasFeature()` before setting `callId`:
+ * McpCallPayload is `deny_unknown_fields` server-side, so sending the field
+ * to a server predating it makes the server reject the call rather than
+ * ignore the field.
+ */
+export const MCP_CALL_ID_FEATURE: 'mcp-call-id';
+
 /** A wsh protocol control message (all messages have a numeric `type`). */
 export interface WshMessage {
   type: number;
@@ -396,10 +406,14 @@ export function mcpTools(opts?: {
 export function mcpCall(opts?: {
   tool?: string;
   arguments?: Record<string, unknown>;
+  /** Correlates this call with its McpResult. Omitted from the message entirely when undefined. */
+  callId?: string;
 }): WshMessage;
 
 export function mcpResult(opts?: {
   result?: unknown;
+  /** Echoed verbatim from the McpCall this answers. */
+  callId?: string;
 }): WshMessage;
 
 // -- Message constructors (reverse) --
@@ -817,7 +831,19 @@ export function isValidMessage(msg: unknown): boolean;
 // ============================================================================
 
 /**
- * Generate a new Ed25519 key pair.
+ * Whether this runtime's WebCrypto implements Ed25519, measured by
+ * generating a key rather than sniffed from a version string.
+ *
+ * There is no pure-JS fallback: one would need the private scalar as
+ * ordinary bytes, giving up the non-extractable `CryptoKey` property
+ * `WshKeyStore` is built around. Check this before committing to pubkey
+ * auth and fall back to password auth where it returns false.
+ */
+export function isEd25519Supported(): Promise<boolean>;
+
+/**
+ * Generate a new Ed25519 key pair. Throws with an actionable message
+ * where WebCrypto has no Ed25519 -- see `isEd25519Supported()`.
  * @param extractable - Whether private key can be exported (default false)
  */
 export function generateKeyPair(extractable?: boolean): Promise<CryptoKeyPair>;
@@ -1017,7 +1043,7 @@ export class WshTransport {
   readonly state: 'disconnected' | 'connecting' | 'connected' | 'closed';
 
   /** Connect to a wsh server. */
-  connect(url: string): Promise<void>;
+  connect(url: string, options?: object): Promise<void>;
 
   /** Gracefully close the transport. */
   close(): Promise<void>;
@@ -1044,7 +1070,7 @@ export class WshTransport {
   protected _emitError(err: Error): void;
 
   /** @protected Must be overridden by subclasses. */
-  protected _doConnect(url: string): Promise<void>;
+  protected _doConnect(url: string, options?: object): Promise<void>;
 
   /** @protected Must be overridden by subclasses. */
   protected _doClose(): Promise<void>;
@@ -1064,7 +1090,65 @@ export class WshTransport {
  * - Subsequent streams carry raw byte data (no framing).
  * - Server-initiated streams are surfaced via onStreamOpen.
  */
-export class WebTransportTransport extends WshTransport {}
+export class WebTransportTransport extends WshTransport {
+  /**
+   * @param options Forwarded to the `WebTransport` constructor. Options
+   *   passed to `connect()` override these per call.
+   */
+  constructor(options?: WshWebTransportOptions);
+
+  connect(url: string, options?: WshWebTransportOptions): Promise<void>;
+}
+
+/**
+ * One entry of `serverCertificateHashes`. The digest may be given as raw
+ * bytes, as hex (with or without `:` separators -- what
+ * `openssl x509 -fingerprint -sha256` prints), or as base64/base64url.
+ */
+export interface WshCertificateHash {
+  /** Defaults to `'sha-256'`, the only algorithm the platform accepts today. */
+  algorithm?: string;
+  value: BufferSource | string;
+}
+
+/**
+ * Options forwarded to the `WebTransport` constructor.
+ *
+ * `serverCertificateHashes` is the interesting one: it pins a specific
+ * self-signed certificate by digest, which is the only mechanism in the
+ * web platform that lets page JavaScript reach a server whose certificate
+ * no certificate authority signed. The platform's own constraints apply --
+ * `https:` URL, HTTP/3 only, ECDSA P-256 key, at most 14 days validity,
+ * no connection pooling.
+ *
+ * Every other key is forwarded verbatim, so platform options not listed
+ * here still work.
+ */
+export interface WshWebTransportOptions {
+  serverCertificateHashes?: Array<WshCertificateHash | BufferSource | string>;
+  allowPooling?: boolean;
+  requireUnreliable?: boolean;
+  congestionControl?: 'default' | 'throughput' | 'low-latency';
+  anticipatedConcurrentIncomingUnidirectionalStreams?: number | null;
+  anticipatedConcurrentIncomingBidirectionalStreams?: number | null;
+  protocols?: string[];
+  [key: string]: unknown;
+}
+
+/**
+ * Decode one `serverCertificateHashes` digest into bytes. Accepts a
+ * `BufferSource`, hex (separators and an `... Fingerprint=` prefix are
+ * stripped), or base64/base64url. Throws a `TypeError` on an
+ * unrecognised shape and a `RangeError` on a digest of the wrong length
+ * for the algorithm.
+ */
+export function parseCertificateHash(value: BufferSource | string, algorithm?: string): Uint8Array;
+
+/**
+ * Normalise a caller-supplied options bag into a `WebTransportOptions`
+ * dictionary, or `undefined` when there is nothing to pass.
+ */
+export function normalizeWebTransportOptions(options?: WshWebTransportOptions): object | undefined;
 
 /**
  * Dispatch a fixed, already-available batch of items one at a time,
@@ -1307,8 +1391,21 @@ export class WshSession {
   /** Called when the remote process exits. */
   onExit: ((code: number) => void) | null;
 
-  /** Called when the session is fully closed. */
-  onClose: (() => void) | null;
+  /**
+   * Called when the session is fully closed.
+   *
+   * Receives the same value as `closeReason`: an Error for a torn-down
+   * session, null for an ordinary close.
+   */
+  onClose: ((reason: Error | null) => void) | null;
+
+  /**
+   * Why the session closed, or null for an ordinary close.
+   *
+   * A torn-down session and a clean exit are otherwise indistinguishable
+   * from the outside.
+   */
+  readonly closeReason: Error | null;
 
   /** Called when the remote peer acknowledges speculative local echo. */
   onEchoAck: ((payload: Record<string, unknown>) => void) | null;
@@ -1427,6 +1524,12 @@ export interface WshConnectOptions {
   keyPair?: CryptoKeyPair;
   password?: string;
   transport?: 'wt' | 'ws';
+  /**
+   * Forwarded to the `WebTransport` constructor when the WebTransport
+   * rung of the transport ladder is tried; ignored by the WebSocket rung,
+   * which has no `serverCertificateHashes` equivalent.
+   */
+  webTransport?: WshWebTransportOptions;
   timeout?: number;
 }
 
@@ -1466,6 +1569,8 @@ export interface WshConnectReverseOptions {
   username: string;
   keyPair?: CryptoKeyPair;
   password?: string;
+  /** Forwarded to the underlying `connect()` -- see `WshConnectOptions`. */
+  webTransport?: WshWebTransportOptions;
   expose?: {
     shell?: boolean;
     exec?: boolean;
@@ -1697,6 +1802,31 @@ export class WshClient {
    */
   sendRelayControl(msg: WshMessage): Promise<void>;
 
+  /**
+   * Send a raw control message on the authenticated connection.
+   *
+   * Prefer the typed methods. This is the hook `WshMcpBridge` and
+   * `WshFileTransfer` compose their own messages through.
+   */
+  sendControl(msg: WshMessage): Promise<void>;
+
+  /**
+   * Open a raw bidirectional stream on the underlying transport.
+   */
+  openStream(): Promise<WshStream>;
+
+  /**
+   * Observe every inbound control message this client handles. Listeners
+   * are called past the RelayForward trust gate, so a message from an
+   * unaccepted peer never reaches one.
+   */
+  addControlListener(fn: (msg: WshMessage) => void): void;
+
+  /**
+   * Stop observing inbound control messages.
+   */
+  removeControlListener(fn: (msg: WshMessage) => void): void;
+
   /**
    * Initiate end-to-end encryption for a session: `'X25519'` (default,
    * classical ECDH) or `'X25519+ML-KEM-768'` (hybrid, with automatic

package/src/index.mjs

@@ -42,6 +42,7 @@ export {
 
 // Auth + crypto
 export {
+  isEd25519Supported,
   generateKeyPair, exportPublicKeyRaw, exportPublicKeySSH,
   importPublicKeyRaw, exportPrivateKeyPKCS8, importPrivateKeyPKCS8,
   sign, verify, buildTranscript, signChallenge, verifyChallenge,
@@ -51,7 +52,10 @@ export {
 } from './auth.mjs';
 
 // Transport
-export { WshTransport, WebTransportTransport, dispatchSerially, SerialQueue } from './transport.mjs';
+export {
+  WshTransport, WebTransportTransport, dispatchSerially, SerialQueue,
+  parseCertificateHash, normalizeWebTransportOptions,
+} from './transport.mjs';
 export { WebSocketTransport, WS_FRAME_TYPE } from './transport-ws.mjs';
 
 // QMux (QUIC-v1 frames over an ordered byte stream, draft-ietf-quic-qmux-02)
@@ -68,7 +72,7 @@ export {
 // Session + Client
 export { WshSession } from './session.mjs';
 export { WshVirtualSessionBackend, normalizeSessionData } from './virtual-session.mjs';
-export { WshClient } from './client.mjs';
+export { WshClient, MCP_CALL_ID_FEATURE } from './client.mjs';
 
 // Key storage
 export { WshKeyStore } from './keystore.mjs';

package/src/mcp-bridge.mjs

@@ -7,6 +7,8 @@
  */
 
 import { MSG, mcpDiscover, mcpCall } from './messages.mjs';
+import { waitForControlMessage } from './control-listener.mjs';
+import { MCP_CALL_ID_FEATURE } from './client.mjs';
 
 /** Default timeout for MCP operations (15 seconds). */
 const MCP_TIMEOUT_MS = 15_000;
@@ -18,14 +20,35 @@ export class WshMcpBridge {
   /** @type {Map<string, object>} Cached tool specs: name -> { name, description, parameters } */
   #tools = new Map();
 
+  /** @type {number} Monotonically increasing counter for call_id. */
+  #callCounter = 0;
+
+  /**
+   * Tail of the serialised call chain, used only against servers that do not
+   * advertise `mcp-call-id`. Never rejects.
+   * @type {Promise<void>}
+   */
+  #queue = Promise.resolve();
+
   /**
-   * @param {object} client - A WshClient or transport exposing:
+   * @param {object} client - A `WshClient`, a `WshTransport`, or any object
+   *   exposing:
    *   - sendControl(msg): send a control message
    *   - addControlListener(fn) / removeControlListener(fn): message listeners
-   *     (or _controlListeners array, or onControl callback)
+   *     (or a `_controlListeners` array, or a settable `onControl` callback)
+   *
+   *   `sendControl` is checked here rather than at first use: without it
+   *   `discover()` failed with `TypeError: this.#client.sendControl is not
+   *   a function` from inside a promise, naming a private field and not
+   *   the argument that was actually wrong.
    */
   constructor(client) {
     if (!client) throw new Error('WshMcpBridge requires a client');
+    if (typeof client.sendControl !== 'function') {
+      throw new TypeError(
+        'WshMcpBridge requires a client exposing sendControl() — pass a WshClient or a WshTransport'
+      );
+    }
     this.#client = client;
   }
 
@@ -73,6 +96,13 @@ export class WshMcpBridge {
    *
    * Sends an MCP_CALL message and waits for the MCP_RESULT response.
    *
+   * Safe to call concurrently. Against a server advertising `mcp-call-id`
+   * each call carries a correlation id and takes only its own reply;
+   * otherwise calls are queued one at a time, because without an id on the
+   * wire the replies are indistinguishable and `{ success, output }`
+   * normalisation strips whatever the payload might have identified itself
+   * by.
+   *
    * @param {string} toolName - Name of the tool to invoke
    * @param {object} [args={}] - Arguments to pass to the tool
    * @param {object} [opts]
@@ -91,10 +121,40 @@ export class WshMcpBridge {
       );
     }
 
-    await this.#client.sendControl(mcpCall({ tool: toolName, arguments: args }));
+    if (this.#correlates()) return this.#callOnce(toolName, args, timeout);
+
+    const run = this.#queue.then(
+      () => this.#callOnce(toolName, args, timeout),
+      () => this.#callOnce(toolName, args, timeout),
+    );
+    this.#queue = run.then(() => undefined, () => undefined);
+    return run;
+  }
+
+  /** True when the connected server echoes call_id. @private */
+  #correlates() {
+    return this.#client?.hasFeature?.(MCP_CALL_ID_FEATURE) === true;
+  }
+
+  /** @private */
+  async #callOnce(toolName, args, timeout) {
+    const callId = this.#correlates()
+      ? `bridge-${++this.#callCounter}-${Date.now()}`
+      : undefined;
+
+    await this.#client.sendControl(
+      mcpCall({ tool: toolName, arguments: args, callId })
+    );
 
     const response = await this._waitForMessage(
-      (msg) => msg.type === MSG.MCP_RESULT || msg.type === MSG.ERROR,
+      (msg) => {
+        // ERROR carries no call_id and cannot be attributed to one call;
+        // it is taken by whichever call is waiting.
+        if (msg.type === MSG.ERROR) return true;
+        if (msg.type !== MSG.MCP_RESULT) return false;
+        if (callId === undefined) return true;
+        return msg.call_id === undefined || msg.call_id === callId;
+      },
       timeout
     );
 
@@ -171,58 +231,17 @@ export class WshMcpBridge {
   /**
    * Wait for a control message matching a predicate.
    *
-   * Temporarily hooks into the client's control message flow and resolves
-   * when a matching message arrives (or rejects on timeout).
+   * Subscribes to the client's control message flow and resolves when a
+   * matching message arrives (or rejects on timeout). Deregistration is
+   * `control-listener.mjs`'s job — the version inlined here leaked one
+   * permanent `onControl` wrapper per operation against any client whose
+   * only hook is that property.
    *
    * @param {function(object): boolean} predicate
    * @param {number} timeoutMs
    * @returns {Promise<object>}
    */
   _waitForMessage(predicate, timeoutMs) {
-    return new Promise((resolve, reject) => {
-      let timer = null;
-      let settled = false;
-
-      const cleanup = () => {
-        settled = true;
-        if (timer !== null) clearTimeout(timer);
-        // Remove listener
-        if (this.#client.removeControlListener) {
-          this.#client.removeControlListener(listener);
-        } else if (this.#client._controlListeners) {
-          const idx = this.#client._controlListeners.indexOf(listener);
-          if (idx !== -1) this.#client._controlListeners.splice(idx, 1);
-        }
-      };
-
-      const listener = (msg) => {
-        if (settled) return;
-        if (predicate(msg)) {
-          cleanup();
-          resolve(msg);
-        }
-      };
-
-      // Register listener on the client
-      if (this.#client.addControlListener) {
-        this.#client.addControlListener(listener);
-      } else if (this.#client._controlListeners) {
-        this.#client._controlListeners.push(listener);
-      } else {
-        // Fallback: wrap existing onControl
-        const prev = this.#client.onControl;
-        this.#client.onControl = (msg) => {
-          prev?.(msg);
-          listener(msg);
-        };
-      }
-
-      timer = setTimeout(() => {
-        if (!settled) {
-          cleanup();
-          reject(new Error(`MCP operation timed out after ${timeoutMs}ms`));
-        }
-      }, timeoutMs);
-    });
+    return waitForControlMessage(this.#client, predicate, timeoutMs, 'MCP operation');
   }
 }

package/src/messages.gen.mjs

@@ -445,15 +445,16 @@ export function mcpTools({ tools } = {}) {
   };
 }
 
-export function mcpCall({ tool, arguments: args } = {}) {
-  return { type: MSG.MCP_CALL, tool, arguments: args };
+export function mcpCall({ tool, arguments: args, callId } = {}) {
+  const m = { type: MSG.MCP_CALL, tool, arguments: args };
+  if (callId !== undefined) m.call_id = callId;
+  return m;
 }
 
-export function mcpResult({ result } = {}) {
-  return {
-    type: MSG.MCP_RESULT,
-    result,
-  };
+export function mcpResult({ result, callId } = {}) {
+  const msg = { type: MSG.MCP_RESULT, result };
+  if (callId !== undefined) msg.call_id = callId;
+  return msg;
 }
 
 export function reverseRegister({ username, capabilities = [], peerType = "host", shellBackend = "pty", supportsAttach = false, supportsReplay = false, supportsEcho = false, supportsTermSync = false, publicKey, seq, recordSignature } = {}) {