@johnhenry/wsh 0.16.1 → 0.17.1

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

@@ -145,6 +145,10 @@ export const MSG = Object.freeze({
 
   // Terminal
   TERMINAL_CONFIG:   0x9e,
+
+  // Keys
+  AUTHORIZED_KEY_ADD: 0x9f,
+  AUTHORIZED_KEY_RESULT: 0xa0,
 });
 
 // Reverse lookup: number → name
@@ -187,13 +191,10 @@ export function hello({ username, features = [], authMethod = AUTH_METHOD.PUBKEY
   };
 }
 
-export function serverHello({ sessionId, features = [], fingerprints = [] } = {}) {
-  return {
-    type: MSG.SERVER_HELLO,
-    session_id: sessionId,
-    features,
-    fingerprints,
-  };
+export function serverHello({ sessionId, features = [], fingerprints = [], hostFingerprint } = {}) {
+  const msg = { type: MSG.SERVER_HELLO, session_id: sessionId, features, fingerprints };
+  if (hostFingerprint !== undefined) msg.host_fingerprint = hostFingerprint;
+  return msg;
 }
 
 export function challenge({ nonce, sessionId } = {}) {
@@ -445,15 +446,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 } = {}) {
@@ -863,8 +865,8 @@ export function fileOp({ channelId, op, path, offset, length } = {}) {
   return msg;
 }
 
-export function fileResult({ channelId, success, metadata = {}, errorMessage } = {}) {
-  const msg = { type: MSG.FILE_RESULT, channel_id: channelId, success, metadata };
+export function fileResult({ channelId, success, metadata = {}, entries = [], errorMessage } = {}) {
+  const msg = { type: MSG.FILE_RESULT, channel_id: channelId, success, metadata, entries };
   if (errorMessage !== undefined) msg.error_message = errorMessage;
   return msg;
 }
@@ -914,6 +916,18 @@ export function terminalConfig({ channelId, frontend, options = {} } = {}) {
   };
 }
 
+export function authorizedKeyAdd({ publicKey, comment } = {}) {
+  const msg = { type: MSG.AUTHORIZED_KEY_ADD, public_key: publicKey };
+  if (comment !== undefined) msg.comment = comment;
+  return msg;
+}
+
+export function authorizedKeyResult({ success, added = false, errorMessage } = {}) {
+  const msg = { type: MSG.AUTHORIZED_KEY_RESULT, success, added };
+  if (errorMessage !== undefined) msg.error_message = errorMessage;
+  return msg;
+}
+
 // ── Relay forwarding ─────────────────────────────────────────────────
 
 // Message types a relay may forward inside a RelayForward wrapper.

package/src/mlkem.mjs

@@ -66,7 +66,34 @@ function bytesEqual(a, b) {
   return diff === 0;
 }
 
+/**
+ * Which implementation serves ML-KEM-768, with an override for testing.
+ *
+ * The noble branch is the ONLY path a browser can take -- no browser
+ * implements the experimental WebCrypto ML-KEM draft -- and until this
+ * override existed no test had ever executed a line of it. `getBackend()`
+ * memoised a probe that always succeeds on the repo's target runtime, so
+ * every assertion ran against native, and every noble call
+ * (`ml_kem768.keygen`, `.encapsulate`, `.decapsulate`) was untouched. A
+ * change in noble's argument order or seed semantics -- it is pinned only to
+ * `^0.7.0` and called through three positional APIs -- would break hybrid PQ
+ * key exchange for every browser consumer while this repo stayed green.
+ *
+ * `WSH_MLKEM_BACKEND=noble` forces it. Read defensively: this module runs in
+ * browsers, where `process` does not exist.
+ */
+function backendOverride() {
+  try {
+    const value = globalThis.process?.env?.WSH_MLKEM_BACKEND;
+    return value === 'noble' || value === 'native' ? value : null;
+  } catch {
+    return null;
+  }
+}
+
 async function getBackend() {
+  const forced = backendOverride();
+  if (forced) return forced;
   if (!backendPromise) {
     backendPromise = probeNativeSupport().then((supported) => (supported ? 'native' : 'noble'));
   }

package/src/session.mjs

@@ -153,6 +153,8 @@ export class WshSession {
    * @type {ReturnType<typeof setTimeout>|null}
    */
   #closeGraceTimer = null;
+  /** Why the session closed, when it closed for a reason worth naming. */
+  #closeReason = null;
 
   // ── Callbacks ───────────────────────────────────────────────────────
 
@@ -475,7 +477,7 @@ export class WshSession {
     for (const { nonce, ciphertext } of wireChunks) {
       const counter = this.#e2eRecvCounter++;
       try {
-        const plaintext = await openFrame(this.#e2eKey, this.#sessionId, counter, { nonce, ciphertext });
+        const plaintext = await openFrame(this.#e2eKey, this.#sessionId, counter, { nonce, ciphertext, expectedRoleTag: this.#e2eRecvRoleTag });
         plaintexts.push(plaintext);
       } catch (err) {
         console.error(
@@ -682,7 +684,7 @@ export class WshSession {
         // delivery of *this* frame's plaintext to onData is deferred by
         // one microtask/macrotask relative to synchronous cases.
         const counter = this.#e2eRecvCounter++;
-        openFrame(this.#e2eKey, this.#sessionId, counter, { nonce: msg.nonce, ciphertext: msg.ciphertext })
+        openFrame(this.#e2eKey, this.#sessionId, counter, { nonce: msg.nonce, ciphertext: msg.ciphertext, expectedRoleTag: this.#e2eRecvRoleTag })
           .then((plaintext) => {
             this.#virtualBackend?.pushData(plaintext);
             try {
@@ -724,14 +726,7 @@ export class WshSession {
         // any pending data-EOF grace timer, regardless of which arrived
         // first.
         this.#clearCloseGraceTimer();
-        if (this.#state !== STATE_CLOSED) {
-          this.#state = STATE_CLOSED;
-          this.#abort.abort();
-          this.#virtualBackend?.close();
-          this.#releaseStreams();
-          this.#resolvePendingFileChunk(null);
-          this.#emitClose();
-        }
+        this.#closeNow();
         break;
       }
 
@@ -874,11 +869,32 @@ export class WshSession {
           if (this.#e2eKey !== null && this.#streamAccumulator) {
             const deliverable = await this.#openStreamChunks(value);
             if (deliverable === null) {
-              // Authentication failure (or corrupt framing) -- matches
-              // the strict "no skip-and-continue" philosophy of
-              // virtual-mode's counter check: tear the session down
-              // rather than deliver anything from this point on.
-              break;
+              /*
+               * Authentication failure (or corrupt framing) -- the strict
+               * "no skip-and-continue" philosophy of virtual-mode's counter
+               * check: tear the session down rather than deliver anything
+               * from this point on.
+               *
+               * `break` alone did not tear anything down. It left the loop,
+               * and control fell through to the data-EOF grace timer, so the
+               * session stayed `active` and writable for the full 300ms and
+               * then closed as an ordinary EOF: onClose with no error,
+               * onExit never fired, exitCode still null. An application could
+               * not tell "the peer forged data at me" from "the process
+               * exited and the server forgot to send CLOSE" -- the two were
+               * byte-identical to observe.
+               *
+               * Closed immediately, and the reason recorded so the two are
+               * distinguishable. `StreamAuthenticationError` was already
+               * built here and handed only to console.error; it is now the
+               * close reason as well.
+               */
+              this.#closeReason = new StreamAuthenticationError(
+                'wsh: stream authentication failed -- session torn down'
+              );
+              this.#clearCloseGraceTimer();
+              this.#closeNow();
+              return;
             }
             for (const plaintext of deliverable) {
               try {
@@ -918,11 +934,7 @@ export class WshSession {
     if (this.#state !== STATE_CLOSED && this.#closeGraceTimer === null) {
       this.#closeGraceTimer = setTimeout(() => {
         this.#closeGraceTimer = null;
-        if (this.#state !== STATE_CLOSED) {
-          this.#state = STATE_CLOSED;
-          this.#releaseStreams();
-          this.#emitClose();
-        }
+        this.#closeNow();
       }, DATA_EOF_CLOSE_GRACE_MS);
       this.#closeGraceTimer?.unref?.();
     }
@@ -935,6 +947,36 @@ export class WshSession {
    * whether or not a timer is currently scheduled.
    * @private
    */
+  /**
+   * Close the session once, releasing everything a close has to release.
+   *
+   * There were two ways a stream-mode session reached `closed` and they did
+   * not agree. The CLOSE control handler aborted, closed the virtual backend,
+   * released the streams AND settled any parked file chunk. The data-EOF
+   * grace fallback -- the path taken when a server ends the data stream
+   * without sending CLOSE, which the comment above it says is real -- only
+   * released the streams and fired onClose.
+   *
+   * So `download()`, parked on `_readFileChunk()`, was never settled on that
+   * path. It documents returning `null` for a truncated transfer and
+   * `download()` turns that into a thrown error; via the fallback it returned
+   * nothing at all and the download hung forever, while `onClose` had already
+   * fired and the session read `closed`.
+   *
+   * One method rather than two call sites kept in step by hand: the drift is
+   * the bug, and the next field added to a close would have drifted the same
+   * way.
+   */
+  #closeNow() {
+    if (this.#state === STATE_CLOSED) return;
+    this.#state = STATE_CLOSED;
+    this.#abort.abort();
+    this.#virtualBackend?.close();
+    this.#releaseStreams();
+    this.#resolvePendingFileChunk(null);
+    this.#emitClose();
+  }
+
   #clearCloseGraceTimer() {
     if (this.#closeGraceTimer !== null) {
       clearTimeout(this.#closeGraceTimer);
@@ -961,9 +1003,19 @@ export class WshSession {
    * Emit the onClose callback exactly once.
    * @private
    */
+  /**
+   * Why the session closed, or null for an ordinary close.
+   *
+   * A torn-down session and a clean exit used to be indistinguishable from
+   * the outside; this is what tells them apart.
+   */
+  get closeReason() {
+    return this.#closeReason;
+  }
+
   #emitClose() {
     try {
-      this.onClose?.();
+      this.onClose?.(this.#closeReason);
     } catch (err) {
       console.error('[wsh:session] onClose handler error:', err);
     }

package/src/transport-ws.mjs

@@ -214,8 +214,22 @@ export class WebSocketTransport extends WshTransport {
     // own stream objects -- without this, any reader still waiting on a
     // stream's readable side would hang forever, since nothing else
     // errors or closes it once the peer can no longer respond.
-    this.#qmux?.close();
-    this.#qmux?.destroy(new Error('Transport closed'));
+    //
+    // Each step is guarded separately. The courtesy CONNECTION_CLOSE is
+    // the one that can realistically throw -- it writes to the socket,
+    // and a socket that is already broken is the single most likely
+    // reason we are closing at all. Letting that throw would skip both
+    // the local `destroy()` (hanging every waiting reader, the exact
+    // thing above) and the socket close (leaking the connection), which
+    // makes the failure mode strictly worse than the failure that
+    // triggered it.
+    try {
+      this.#qmux?.close();
+    } catch { /* peer is unreachable; the local teardown below still matters */ }
+
+    try {
+      this.#qmux?.destroy(new Error('Transport closed'));
+    } catch { /* a stream's error callback threw; keep tearing down */ }
 
     if (this.#ws) {
       try {