@jaw.id/cli 0.0.3 → 0.0.6

package/dist/mcp/handlers/daemon.js

@@ -1,10 +1,8 @@
 import * as fs2 from 'fs';
+import * as crypto from 'crypto';
 import * as path from 'path';
-import 'url';
-import { execSync } from 'child_process';
 import * as os from 'os';
-import 'crypto';
-import 'ws';
+import WebSocket from 'ws';
 import { z } from 'zod';
 
 // src/mcp/handlers/daemon.ts
@@ -26,8 +24,7 @@ var PATHS = {
   root: JAW_DIR,
   config: path.join(JAW_DIR, "config.json"),
   session: path.join(JAW_DIR, "session.json"),
-  bridge: path.join(JAW_DIR, "bridge.json"),
-  daemonLog: path.join(JAW_DIR, "daemon.log")
+  relay: path.join(JAW_DIR, "relay.json")
 };
 function loadConfig() {
   if (!fs2.existsSync(PATHS.config)) {
@@ -49,53 +46,459 @@ function redactConfig(config) {
   };
 }
 
-// src/lib/bridge-singleton.ts
-path.join(PATHS.root, "daemon.lock");
-function isDaemonProcess(pid) {
-  if (!Number.isInteger(pid) || pid <= 0 || pid > 4194304) return false;
-  try {
-    process.kill(pid, 0);
-  } catch {
-    return false;
+// src/lib/crypto.ts
+var subtle = globalThis.crypto.subtle;
+async function deriveSharedSecret(privateKey, peerPublicKey) {
+  return subtle.deriveKey(
+    { name: "ECDH", public: peerPublicKey },
+    privateKey,
+    { name: "AES-GCM", length: 256 },
+    false,
+    ["encrypt", "decrypt"]
+  );
+}
+async function encryptMessage(sharedSecret, payload) {
+  const iv = globalThis.crypto.getRandomValues(new Uint8Array(12));
+  const plaintext = new TextEncoder().encode(JSON.stringify(payload));
+  const cipherBuf = await subtle.encrypt(
+    { name: "AES-GCM", iv },
+    sharedSecret,
+    plaintext
+  );
+  return {
+    iv: bufferToBase64(iv),
+    ciphertext: bufferToBase64(new Uint8Array(cipherBuf))
+  };
+}
+async function decryptMessage(sharedSecret, envelope) {
+  const iv = Buffer.from(envelope.iv, "base64");
+  const ciphertext = Buffer.from(envelope.ciphertext, "base64");
+  const plainBuf = await subtle.decrypt(
+    { name: "AES-GCM", iv },
+    sharedSecret,
+    ciphertext
+  );
+  return JSON.parse(new TextDecoder().decode(plainBuf));
+}
+async function importKeyFromHex(type, hex) {
+  const format = type === "private" ? "pkcs8" : "spki";
+  return subtle.importKey(
+    format,
+    Buffer.from(hexToBytes(hex)),
+    { name: "ECDH", namedCurve: "P-256" },
+    true,
+    type === "private" ? ["deriveKey"] : []
+  );
+}
+function hexToBytes(hex) {
+  if (hex.length % 2 !== 0) throw new Error("Invalid hex: odd length");
+  const bytes = new Uint8Array(hex.length / 2);
+  for (let i = 0; i < hex.length; i += 2) {
+    bytes[i / 2] = parseInt(hex.substring(i, i + 2), 16);
+  }
+  return bytes;
+}
+function bufferToBase64(buf) {
+  return Buffer.from(buf).toString("base64");
+}
+
+// src/lib/ws-bridge.ts
+var DEFAULT_TIMEOUT_MS = 12e4;
+var MAX_MESSAGE_BYTES = 5 * 1024 * 1024;
+var BROWSER_REOPEN_COOLDOWN_MS = 5e3;
+var MAX_RECONNECT_ATTEMPTS = 3;
+var RECONNECT_BASE_DELAY_MS = 1e3;
+var WSBridge = class {
+  relayUrl;
+  session;
+  timeout;
+  config;
+  privateKeyHex;
+  publicKeyHex;
+  peerPublicKeyHex;
+  sharedSecret = null;
+  ws = null;
+  disposed = false;
+  // Auto-reopen browser state
+  onBrowserNeeded;
+  onPeerKeyChanged;
+  lastBrowserOpenTime = 0;
+  // Reconnection state
+  reconnectAttempts = 0;
+  /** Updated after key exchange — caller should persist this. */
+  get peerPublicKey() {
+    return this.peerPublicKeyHex;
+  }
+  constructor(options) {
+    this.relayUrl = options.relayUrl;
+    this.session = options.session;
+    this.timeout = options.timeout ?? DEFAULT_TIMEOUT_MS;
+    this.config = options.config;
+    this.privateKeyHex = options.privateKeyHex;
+    this.publicKeyHex = options.publicKeyHex;
+    this.peerPublicKeyHex = options.peerPublicKeyHex;
+  }
+  /**
+   * Connect to the relay and wait for the browser to be ready.
+   *
+   * @param onBrowserNeeded — called when the relay reports no browser connected.
+   * @param onPeerKeyChanged — called when a key_exchange updates the peer key.
+   */
+  async connect(onBrowserNeeded, onPeerKeyChanged) {
+    this.onBrowserNeeded = onBrowserNeeded;
+    this.onPeerKeyChanged = onPeerKeyChanged;
+    if (this.peerPublicKeyHex) {
+      await this.deriveSecret();
+    }
+    return this.connectInternal(onBrowserNeeded, onPeerKeyChanged);
+  }
+  async connectInternal(onBrowserNeeded, onPeerKeyChanged) {
+    return new Promise((resolve, reject) => {
+      const url = `${this.relayUrl}?session=${encodeURIComponent(this.session)}&role=cli`;
+      const ws = new WebSocket(url);
+      let browserOpened = false;
+      let resolved = false;
+      let expectingKeyExchange = !this.peerPublicKeyHex;
+      const timer = setTimeout(() => {
+        ws.close();
+        reject(
+          new Error(
+            "Browser did not connect in time.\nRun `jaw disconnect` then try again."
+          )
+        );
+      }, 3e4);
+      const sendEncryptedInit = async () => {
+        if (!this.sharedSecret) return;
+        const envelope = await encryptMessage(this.sharedSecret, {
+          type: "init",
+          apiKey: this.config.apiKey,
+          chainId: this.config.chainId,
+          ens: this.config.ens,
+          paymasterUrl: this.config.paymasterUrl
+        });
+        this.sendRaw(ws, JSON.stringify({ type: "encrypted", ...envelope }));
+      };
+      const waitForReady = () => {
+        const readyTimer = setTimeout(() => {
+          ws.close();
+          reject(new Error("Browser SDK did not become ready in time."));
+        }, 15e3);
+        const onMsg = async (data) => {
+          const msg = safeParse(data);
+          if (!msg) return;
+          if (msg.type === "encrypted" && this.sharedSecret) {
+            try {
+              const inner = await decryptMessage(
+                this.sharedSecret,
+                msg
+              );
+              if (inner.type === "ready") {
+                clearTimeout(readyTimer);
+                ws.off("message", onMsg);
+                this.reconnectAttempts = 0;
+                resolve();
+              }
+            } catch {
+            }
+          }
+        };
+        ws.on("message", onMsg);
+      };
+      const onBrowserReady = async () => {
+        if (resolved) return;
+        resolved = true;
+        clearTimeout(timer);
+        waitForReady();
+        await sendEncryptedInit();
+      };
+      ws.on("open", () => {
+        this.ws = ws;
+      });
+      ws.on("message", async (data) => {
+        const msg = safeParse(data);
+        if (!msg) return;
+        if (msg.type === "status") {
+          if (msg.browserConnected) {
+            if (this.sharedSecret) {
+              await onBrowserReady();
+            } else {
+              expectingKeyExchange = true;
+            }
+          } else if (!browserOpened && onBrowserNeeded) {
+            browserOpened = true;
+            expectingKeyExchange = true;
+            onBrowserNeeded().catch(() => {
+            });
+          }
+        } else if (msg.type === "browser_connected") {
+          expectingKeyExchange = true;
+        } else if (msg.type === "browser_disconnected") {
+          this.handleBrowserDisconnect();
+        } else if (msg.type === "key_exchange" && expectingKeyExchange) {
+          expectingKeyExchange = false;
+          const peerKey = msg.publicKey;
+          this.peerPublicKeyHex = peerKey;
+          await this.deriveSecret();
+          onPeerKeyChanged?.(peerKey);
+          await onBrowserReady();
+        }
+      });
+      ws.on("error", (err) => {
+        clearTimeout(timer);
+        reject(err);
+      });
+      ws.on("close", () => {
+        clearTimeout(timer);
+        if (!this.disposed) {
+          this.handleRelayDisconnect();
+        }
+      });
+    });
+  }
+  /**
+   * Send an encrypted RPC request through the relay to the browser SDK.
+   */
+  async request(method, params) {
+    const ws = this.ws;
+    if (!ws || ws.readyState !== WebSocket.OPEN) {
+      throw new Error("Not connected to relay");
+    }
+    if (!this.sharedSecret) {
+      throw new Error("No shared secret \u2014 key exchange not completed");
+    }
+    const id = crypto.randomUUID();
+    const envelope = await encryptMessage(this.sharedSecret, {
+      type: "rpc_request",
+      id,
+      method,
+      params
+    });
+    const serialized = JSON.stringify({ type: "encrypted", ...envelope });
+    assertMessageSize(serialized, method);
+    return new Promise((resolve, reject) => {
+      const timer = setTimeout(() => {
+        reject(
+          new Error(
+            `Request timed out after ${this.timeout / 1e3}s. Did you complete the action in the browser?`
+          )
+        );
+        this.close();
+      }, this.timeout);
+      const onMessage = async (data) => {
+        const msg = safeParse(data);
+        if (!msg || msg.type !== "encrypted" || !this.sharedSecret) return;
+        try {
+          const inner = await decryptMessage(
+            this.sharedSecret,
+            msg
+          );
+          if (inner.type === "rpc_response" && inner.id === id) {
+            clearTimeout(timer);
+            ws.off("message", onMessage);
+            if (inner.success) {
+              resolve(inner.data);
+            } else {
+              const err = inner.error;
+              reject(
+                new Error(
+                  err ? `[${err.code}] ${err.message}` : "Request failed"
+                )
+              );
+            }
+          }
+        } catch {
+        }
+      };
+      ws.on("message", onMessage);
+      this.sendRaw(ws, serialized);
+    });
+  }
+  isOpen() {
+    return this.ws?.readyState === WebSocket.OPEN;
+  }
+  async shutdown() {
+    this.disposed = true;
+    if (this.ws?.readyState === WebSocket.OPEN && this.sharedSecret) {
+      try {
+        const envelope = await encryptMessage(this.sharedSecret, {
+          type: "shutdown"
+        });
+        this.sendRaw(
+          this.ws,
+          JSON.stringify({ type: "encrypted", ...envelope })
+        );
+      } catch {
+      }
+    }
+    this.close();
   }
+  /**
+   * Connect to relay and send shutdown directly — no init/ready handshake.
+   * Used by `jaw disconnect` when we just need to tell the browser to close.
+   */
+  async connectAndShutdown() {
+    if (!this.peerPublicKeyHex) {
+      return;
+    }
+    this.disposed = true;
+    await this.deriveSecret();
+    return new Promise((resolve) => {
+      const url = `${this.relayUrl}?session=${encodeURIComponent(this.session)}&role=cli`;
+      const ws = new WebSocket(url);
+      const timer = setTimeout(() => {
+        try {
+          ws.close();
+        } catch {
+        }
+        resolve();
+      }, 3e3);
+      ws.on("open", async () => {
+        this.ws = ws;
+        try {
+          await this.shutdown();
+        } catch {
+        }
+        clearTimeout(timer);
+        resolve();
+      });
+      ws.on("error", () => {
+        clearTimeout(timer);
+        resolve();
+      });
+    });
+  }
+  close() {
+    this.disposed = true;
+    if (this.ws) {
+      try {
+        this.ws.close();
+      } catch {
+      }
+      this.ws = null;
+    }
+  }
+  /**
+   * Auto-reopen browser when browser_disconnected is received from relay.
+   * Respects a cooldown to prevent rapid re-opening.
+   */
+  handleBrowserDisconnect() {
+    if (this.disposed || !this.onBrowserNeeded) return;
+    const now = Date.now();
+    if (now - this.lastBrowserOpenTime < BROWSER_REOPEN_COOLDOWN_MS) {
+      return;
+    }
+    this.lastBrowserOpenTime = now;
+    this.sharedSecret = null;
+    this.peerPublicKeyHex = null;
+    this.onBrowserNeeded().catch(() => {
+    });
+  }
+  /**
+   * Attempt to reconnect to the relay with exponential backoff
+   * when the WebSocket connection drops unexpectedly.
+   */
+  handleRelayDisconnect() {
+    if (this.disposed) return;
+    if (this.reconnectAttempts >= MAX_RECONNECT_ATTEMPTS) return;
+    const delay = RECONNECT_BASE_DELAY_MS * Math.pow(2, this.reconnectAttempts);
+    this.reconnectAttempts++;
+    setTimeout(() => {
+      if (this.disposed) return;
+      this.connectInternal(this.onBrowserNeeded, this.onPeerKeyChanged).catch(
+        () => {
+        }
+      );
+    }, delay);
+  }
+  /** Send a raw string over the WebSocket, enforcing message size limits. */
+  sendRaw(ws, data) {
+    ws.send(data);
+  }
+  async deriveSecret() {
+    if (!this.peerPublicKeyHex) return;
+    const privateKey = await importKeyFromHex("private", this.privateKeyHex);
+    const peerPublicKey = await importKeyFromHex(
+      "public",
+      this.peerPublicKeyHex
+    );
+    this.sharedSecret = await deriveSharedSecret(privateKey, peerPublicKey);
+  }
+};
+function assertMessageSize(serialized, method) {
+  const byteLength = Buffer.byteLength(serialized, "utf-8");
+  if (byteLength > MAX_MESSAGE_BYTES) {
+    const sizeMB = (byteLength / (1024 * 1024)).toFixed(2);
+    throw new Error(
+      `Message for ${method} is too large (${sizeMB} MB, limit ${MAX_MESSAGE_BYTES / (1024 * 1024)} MB). Try reducing the number of calls in your batch.`
+    );
+  }
+}
+function safeParse(data) {
   try {
-    const cmd = execSync(`ps -p ${String(pid)} -o command=`, {
-      encoding: "utf-8",
-      timeout: 3e3
-    }).trim();
-    return cmd.includes("ws-daemon");
+    return JSON.parse(data.toString());
   } catch {
-    return false;
+    return null;
   }
 }
-function loadBridgeInfo() {
+
+// src/lib/bridge-singleton.ts
+function loadRelaySession() {
   try {
-    if (!fs2.existsSync(PATHS.bridge)) return null;
-    const raw = fs2.readFileSync(PATHS.bridge, "utf-8");
-    const info = JSON.parse(raw);
-    if (!isDaemonProcess(info.pid)) {
-      try {
-        fs2.unlinkSync(PATHS.bridge);
-      } catch {
-      }
+    if (!fs2.existsSync(PATHS.relay)) return null;
+    const raw = fs2.readFileSync(PATHS.relay, "utf-8");
+    const parsed = JSON.parse(raw);
+    if (!parsed.session || !parsed.relayUrl || !parsed.privateKey || !parsed.publicKey) {
       return null;
     }
-    return info;
+    return parsed;
   } catch {
     return null;
   }
 }
+function deleteRelaySession() {
+  try {
+    if (fs2.existsSync(PATHS.relay)) fs2.unlinkSync(PATHS.relay);
+  } catch {
+  }
+}
 async function shutdownDaemon() {
-  const info = loadBridgeInfo();
-  if (!info) return;
+  const session = loadRelaySession();
+  if (!session) return;
   try {
-    process.kill(info.pid, "SIGTERM");
+    const bridge = new WSBridge({
+      relayUrl: session.relayUrl,
+      session: session.session,
+      timeout: 5e3,
+      config: { apiKey: "", chainId: 1 },
+      privateKeyHex: session.privateKey,
+      publicKeyHex: session.publicKey,
+      peerPublicKeyHex: session.peerPublicKey
+    });
+    await bridge.connectAndShutdown();
   } catch {
   }
+  deleteRelaySession();
+  const legacyBridge = PATHS.root + "/bridge.json";
+  const legacyLog = PATHS.root + "/daemon.log";
+  const legacyLock = PATHS.root + "/daemon.lock";
   try {
-    if (fs2.existsSync(PATHS.bridge)) fs2.unlinkSync(PATHS.bridge);
+    if (fs2.existsSync(legacyBridge)) {
+      const info = JSON.parse(fs2.readFileSync(legacyBridge, "utf-8"));
+      if (info.pid && Number.isInteger(info.pid) && info.pid > 0) {
+        try {
+          process.kill(info.pid, "SIGTERM");
+        } catch {
+        }
+      }
+    }
   } catch {
   }
+  for (const f of [legacyBridge, legacyLog, legacyLock]) {
+    try {
+      if (fs2.existsSync(f)) fs2.unlinkSync(f);
+    } catch {
+    }
+  }
 }
 ({
   method: z.string().describe(
@@ -125,25 +528,22 @@ function isBridgeCached() {
 function registerDaemonTools(server) {
   server.tool(
     "jaw_status",
-    "Check the current status of the JAW.id bridge \u2014 whether the daemon is running, the bridge connection is active, and what configuration is in use.",
+    "Check the current status of the JAW.id relay bridge \u2014 whether a relay session exists, the bridge connection is active, and what configuration is in use.",
     {},
     async () => {
       try {
-        let daemonRunning = false;
-        let daemonPid = null;
+        let relaySession = false;
         try {
-          if (fs2.existsSync(PATHS.bridge)) {
-            const info = JSON.parse(fs2.readFileSync(PATHS.bridge, "utf-8"));
-            daemonPid = info.pid;
-            process.kill(info.pid, 0);
-            daemonRunning = true;
+          if (fs2.existsSync(PATHS.relay)) {
+            JSON.parse(fs2.readFileSync(PATHS.relay, "utf-8"));
+            relaySession = true;
           }
         } catch {
-          daemonRunning = false;
+          relaySession = false;
         }
         const config = redactConfig(loadConfig());
         const status = {
-          daemon: daemonRunning ? { running: true, pid: daemonPid } : { running: false },
+          relay: relaySession ? { session: true } : { session: false },
           bridgeConnection: isBridgeCached() ? "connected" : "disconnected",
           config
         };
@@ -159,7 +559,7 @@ function registerDaemonTools(server) {
   );
   server.tool(
     "jaw_disconnect",
-    "Stop the background bridge daemon and close the browser session. Call this when you are done making wallet requests to clean up resources.",
+    "Close the relay session and browser tab. Call this when you are done making wallet requests to clean up resources.",
     {},
     async () => {
       try {
@@ -169,7 +569,7 @@ function registerDaemonTools(server) {
           content: [
             {
               type: "text",
-              text: "Bridge daemon stopped and browser session closed."
+              text: "Relay session closed and browser tab dismissed."
             }
           ]
         };