@btatum5/codex-bridge 0.1.0 → 1.3.2

package/src/bridge.js

@@ -0,0 +1,620 @@
+// FILE: bridge.js
+// Purpose: Runs Codex locally, bridges relay traffic, and coordinates desktop refreshes for Codex.app.
+// Layer: CLI service
+// Exports: startBridge
+// Depends on: ws, crypto, os, ./qr, ./codex-desktop-refresher, ./codex-transport, ./rollout-watch
+
+const WebSocket = require("ws");
+const { randomBytes } = require("crypto");
+const os = require("os");
+const {
+  CodexDesktopRefresher,
+  readBridgeConfig,
+} = require("./codex-desktop-refresher");
+const { createCodexTransport } = require("./codex-transport");
+const { createThreadRolloutActivityWatcher } = require("./rollout-watch");
+const { printQR } = require("./qr");
+const { rememberActiveThread } = require("./session-state");
+const { handleDesktopRequest } = require("./desktop-handler");
+const { handleGitRequest } = require("./git-handler");
+const { handleThreadContextRequest } = require("./thread-context-handler");
+const { handleWorkspaceRequest } = require("./workspace-handler");
+const { createNotificationsHandler } = require("./notifications-handler");
+const { createPushNotificationServiceClient } = require("./push-notification-service-client");
+const { createPushNotificationTracker } = require("./push-notification-tracker");
+const {
+  loadOrCreateBridgeDeviceState,
+  resolveBridgeRelaySession,
+} = require("./secure-device-state");
+const { createBridgeSecureTransport } = require("./secure-transport");
+const { createRolloutLiveMirrorController } = require("./rollout-live-mirror");
+
+function startBridge({
+  config: explicitConfig = null,
+  printPairingQr = true,
+  onPairingPayload = null,
+  onBridgeStatus = null,
+} = {}) {
+  const config = explicitConfig || readBridgeConfig();
+  const relayBaseUrl = config.relayUrl.replace(/\/+$/, "");
+  if (!relayBaseUrl) {
+    console.error("[codex-bridge] No relay URL configured.");
+    console.error("[codex-bridge] In a source checkout, run ./run-local-codex.sh or set REMODEX_RELAY.");
+    process.exit(1);
+  }
+
+  let deviceState;
+  try {
+    deviceState = loadOrCreateBridgeDeviceState();
+  } catch (error) {
+    console.error(`[codex-bridge] ${(error && error.message) || "Failed to load the saved bridge pairing state."}`);
+    process.exit(1);
+  }
+  const relaySession = resolveBridgeRelaySession(deviceState);
+  deviceState = relaySession.deviceState;
+  const sessionId = relaySession.sessionId;
+  const relaySessionUrl = `${relayBaseUrl}/${sessionId}`;
+  const notificationSecret = randomBytes(24).toString("hex");
+  const desktopRefresher = new CodexDesktopRefresher({
+    enabled: config.refreshEnabled,
+    debounceMs: config.refreshDebounceMs,
+    refreshCommand: config.refreshCommand,
+    bundleId: config.codexBundleId,
+    appPath: config.codexAppPath,
+  });
+  const pushServiceClient = createPushNotificationServiceClient({
+    baseUrl: config.pushServiceUrl,
+    sessionId,
+    notificationSecret,
+  });
+  const notificationsHandler = createNotificationsHandler({
+    pushServiceClient,
+  });
+  const pushNotificationTracker = createPushNotificationTracker({
+    sessionId,
+    pushServiceClient,
+    previewMaxChars: config.pushPreviewMaxChars,
+  });
+
+  // Keep the local Codex runtime alive across transient relay disconnects.
+  let socket = null;
+  let isShuttingDown = false;
+  let reconnectAttempt = 0;
+  let reconnectTimer = null;
+  let lastConnectionStatus = null;
+  let codexHandshakeState = config.codexEndpoint ? "warm" : "cold";
+  const forwardedInitializeRequestIds = new Set();
+  const secureTransport = createBridgeSecureTransport({
+    sessionId,
+    relayUrl: relayBaseUrl,
+    deviceState,
+    onTrustedPhoneUpdate(nextDeviceState) {
+      deviceState = nextDeviceState;
+      sendRelayRegistrationUpdate(nextDeviceState);
+    },
+  });
+  // Keeps one stable sender identity across reconnects so buffered replay state
+  // reflects what actually made it onto the current relay socket.
+  function sendRelayWireMessage(wireMessage) {
+    if (socket?.readyState !== WebSocket.OPEN) {
+      return false;
+    }
+
+    socket.send(wireMessage);
+    return true;
+  }
+  // Only the spawned local runtime needs rollout mirroring; a real endpoint
+  // already provides the authoritative live stream for resumed threads.
+  const rolloutLiveMirror = !config.codexEndpoint
+    ? createRolloutLiveMirrorController({
+      sendApplicationResponse,
+    })
+    : null;
+  let contextUsageWatcher = null;
+  let watchedContextUsageKey = null;
+
+  const codex = createCodexTransport({
+    endpoint: config.codexEndpoint,
+    env: process.env,
+    logPrefix: "[codex-bridge]",
+  });
+  publishBridgeStatus({
+    state: "starting",
+    connectionStatus: "starting",
+    pid: process.pid,
+    lastError: "",
+  });
+
+  codex.onError((error) => {
+    publishBridgeStatus({
+      state: "error",
+      connectionStatus: "error",
+      pid: process.pid,
+      lastError: error.message,
+    });
+    if (config.codexEndpoint) {
+      console.error(`[codex-bridge] Failed to connect to Codex endpoint: ${config.codexEndpoint}`);
+    } else {
+      console.error("[codex-bridge] Failed to start `codex app-server`.");
+      console.error(`[codex-bridge] Launch command: ${codex.describe()}`);
+      console.error("[codex-bridge] Make sure the Codex CLI is installed and that the launcher works on this OS.");
+    }
+    console.error(error.message);
+    process.exit(1);
+  });
+
+  function clearReconnectTimer() {
+    if (!reconnectTimer) {
+      return;
+    }
+
+    clearTimeout(reconnectTimer);
+    reconnectTimer = null;
+  }
+
+  // Keeps npm start output compact by emitting only high-signal connection states.
+  function logConnectionStatus(status) {
+    if (lastConnectionStatus === status) {
+      return;
+    }
+
+    lastConnectionStatus = status;
+    publishBridgeStatus({
+      state: "running",
+      connectionStatus: status,
+      pid: process.pid,
+      lastError: "",
+    });
+    console.log(`[codex-bridge] ${status}`);
+  }
+
+  // Retries the relay socket while preserving the active Codex process and session id.
+  function scheduleRelayReconnect(closeCode) {
+    if (isShuttingDown) {
+      return;
+    }
+
+    if (closeCode === 4000 || closeCode === 4001) {
+      logConnectionStatus("disconnected");
+      shutdown(codex, () => socket, () => {
+        isShuttingDown = true;
+        clearReconnectTimer();
+      });
+      return;
+    }
+
+    if (reconnectTimer) {
+      return;
+    }
+
+    reconnectAttempt += 1;
+    const delayMs = Math.min(1_000 * reconnectAttempt, 5_000);
+    logConnectionStatus("connecting");
+    reconnectTimer = setTimeout(() => {
+      reconnectTimer = null;
+      connectRelay();
+    }, delayMs);
+  }
+
+  function connectRelay() {
+    if (isShuttingDown) {
+      return;
+    }
+
+    logConnectionStatus("connecting");
+    const nextSocket = new WebSocket(relaySessionUrl, {
+      // The relay uses this per-session secret to authenticate the first push registration.
+      headers: {
+        "x-role": "mac",
+        "x-notification-secret": notificationSecret,
+        ...buildMacRegistrationHeaders(deviceState),
+      },
+    });
+    socket = nextSocket;
+
+    nextSocket.on("open", () => {
+      clearReconnectTimer();
+      reconnectAttempt = 0;
+      logConnectionStatus("connected");
+      secureTransport.bindLiveSendWireMessage(sendRelayWireMessage);
+      sendRelayRegistrationUpdate(deviceState);
+    });
+
+    nextSocket.on("message", (data) => {
+      const message = typeof data === "string" ? data : data.toString("utf8");
+      if (secureTransport.handleIncomingWireMessage(message, {
+        sendControlMessage(controlMessage) {
+          if (nextSocket.readyState === WebSocket.OPEN) {
+            nextSocket.send(JSON.stringify(controlMessage));
+          }
+        },
+        onApplicationMessage(plaintextMessage) {
+          handleApplicationMessage(plaintextMessage);
+        },
+      })) {
+        return;
+      }
+    });
+
+    nextSocket.on("close", (code) => {
+      logConnectionStatus("disconnected");
+      if (socket === nextSocket) {
+        socket = null;
+      }
+      stopContextUsageWatcher();
+      rolloutLiveMirror?.stopAll();
+      desktopRefresher.handleTransportReset();
+      scheduleRelayReconnect(code);
+    });
+
+    nextSocket.on("error", () => {
+      logConnectionStatus("disconnected");
+    });
+  }
+
+  const pairingPayload = secureTransport.createPairingPayload();
+  onPairingPayload?.(pairingPayload);
+  if (printPairingQr) {
+    printQR(pairingPayload);
+  }
+  pushServiceClient.logUnavailable();
+  connectRelay();
+
+  codex.onMessage((message) => {
+    trackCodexHandshakeState(message);
+    desktopRefresher.handleOutbound(message);
+    pushNotificationTracker.handleOutbound(message);
+    rememberThreadFromMessage("codex", message);
+    secureTransport.queueOutboundApplicationMessage(message, sendRelayWireMessage);
+  });
+
+  codex.onClose(() => {
+    logConnectionStatus("disconnected");
+    publishBridgeStatus({
+      state: "stopped",
+      connectionStatus: "disconnected",
+      pid: process.pid,
+      lastError: "",
+    });
+    isShuttingDown = true;
+    clearReconnectTimer();
+    stopContextUsageWatcher();
+    rolloutLiveMirror?.stopAll();
+    desktopRefresher.handleTransportReset();
+    if (socket?.readyState === WebSocket.OPEN || socket?.readyState === WebSocket.CONNECTING) {
+      socket.close();
+    }
+  });
+
+  process.on("SIGINT", () => shutdown(codex, () => socket, () => {
+    isShuttingDown = true;
+    clearReconnectTimer();
+  }));
+  process.on("SIGTERM", () => shutdown(codex, () => socket, () => {
+    isShuttingDown = true;
+    clearReconnectTimer();
+  }));
+
+  // Routes decrypted app payloads through the same bridge handlers as before.
+  function handleApplicationMessage(rawMessage) {
+    if (handleBridgeManagedHandshakeMessage(rawMessage)) {
+      return;
+    }
+    if (handleThreadContextRequest(rawMessage, sendApplicationResponse)) {
+      return;
+    }
+    if (handleWorkspaceRequest(rawMessage, sendApplicationResponse)) {
+      return;
+    }
+    if (notificationsHandler.handleNotificationsRequest(rawMessage, sendApplicationResponse)) {
+      return;
+    }
+    if (handleDesktopRequest(rawMessage, sendApplicationResponse, {
+      bundleId: config.codexBundleId,
+      appPath: config.codexAppPath,
+    })) {
+      return;
+    }
+    if (handleGitRequest(rawMessage, sendApplicationResponse)) {
+      return;
+    }
+    desktopRefresher.handleInbound(rawMessage);
+    rolloutLiveMirror?.observeInbound(rawMessage);
+    rememberThreadFromMessage("phone", rawMessage);
+    codex.send(rawMessage);
+  }
+
+  // Encrypts bridge-generated responses instead of letting the relay see plaintext.
+  function sendApplicationResponse(rawMessage) {
+    secureTransport.queueOutboundApplicationMessage(rawMessage, sendRelayWireMessage);
+  }
+
+  function rememberThreadFromMessage(source, rawMessage) {
+    const context = extractBridgeMessageContext(rawMessage);
+    if (!context.threadId) {
+      return;
+    }
+
+    rememberActiveThread(context.threadId, source);
+    if (shouldStartContextUsageWatcher(context)) {
+      ensureContextUsageWatcher(context);
+    }
+  }
+
+  // Mirrors CodexMonitor's persisted token_count fallback so the phone keeps
+  // receiving context-window usage even when the runtime omits live thread usage.
+  function ensureContextUsageWatcher({ threadId, turnId }) {
+    const normalizedThreadId = readString(threadId);
+    const normalizedTurnId = readString(turnId);
+    if (!normalizedThreadId) {
+      return;
+    }
+
+    const nextWatcherKey = `${normalizedThreadId}|${normalizedTurnId || "pending-turn"}`;
+    if (watchedContextUsageKey === nextWatcherKey && contextUsageWatcher) {
+      return;
+    }
+
+    stopContextUsageWatcher();
+    watchedContextUsageKey = nextWatcherKey;
+    contextUsageWatcher = createThreadRolloutActivityWatcher({
+      threadId: normalizedThreadId,
+      turnId: normalizedTurnId,
+      onUsage: ({ threadId: usageThreadId, usage }) => {
+        sendContextUsageNotification(usageThreadId, usage);
+      },
+      onIdle: () => {
+        if (watchedContextUsageKey === nextWatcherKey) {
+          stopContextUsageWatcher();
+        }
+      },
+      onTimeout: () => {
+        if (watchedContextUsageKey === nextWatcherKey) {
+          stopContextUsageWatcher();
+        }
+      },
+      onError: () => {
+        if (watchedContextUsageKey === nextWatcherKey) {
+          stopContextUsageWatcher();
+        }
+      },
+    });
+  }
+
+  function stopContextUsageWatcher() {
+    if (contextUsageWatcher) {
+      contextUsageWatcher.stop();
+    }
+
+    contextUsageWatcher = null;
+    watchedContextUsageKey = null;
+  }
+
+  function sendContextUsageNotification(threadId, usage) {
+    if (!threadId || !usage) {
+      return;
+    }
+
+    sendApplicationResponse(JSON.stringify({
+      method: "thread/tokenUsage/updated",
+      params: {
+        threadId,
+        usage,
+      },
+    }));
+  }
+
+  // The spawned/shared Codex app-server stays warm across phone reconnects.
+  // When iPhone reconnects it sends initialize again, but forwarding that to the
+  // already-initialized Codex transport only produces "Already initialized".
+  function handleBridgeManagedHandshakeMessage(rawMessage) {
+    let parsed = null;
+    try {
+      parsed = JSON.parse(rawMessage);
+    } catch {
+      return false;
+    }
+
+    const method = typeof parsed?.method === "string" ? parsed.method.trim() : "";
+    if (!method) {
+      return false;
+    }
+
+    if (method === "initialize" && parsed.id != null) {
+      if (codexHandshakeState !== "warm") {
+        forwardedInitializeRequestIds.add(String(parsed.id));
+        return false;
+      }
+
+      sendApplicationResponse(JSON.stringify({
+        id: parsed.id,
+        result: {
+          bridgeManaged: true,
+        },
+      }));
+      return true;
+    }
+
+    if (method === "initialized") {
+      return codexHandshakeState === "warm";
+    }
+
+    return false;
+  }
+
+  // Learns whether the underlying Codex transport has already completed its own MCP handshake.
+  function trackCodexHandshakeState(rawMessage) {
+    let parsed = null;
+    try {
+      parsed = JSON.parse(rawMessage);
+    } catch {
+      return;
+    }
+
+    const responseId = parsed?.id;
+    if (responseId == null) {
+      return;
+    }
+
+    const responseKey = String(responseId);
+    if (!forwardedInitializeRequestIds.has(responseKey)) {
+      return;
+    }
+
+    forwardedInitializeRequestIds.delete(responseKey);
+
+    if (parsed?.result != null) {
+      codexHandshakeState = "warm";
+      return;
+    }
+
+    const errorMessage = typeof parsed?.error?.message === "string"
+      ? parsed.error.message.toLowerCase()
+      : "";
+    if (errorMessage.includes("already initialized")) {
+      codexHandshakeState = "warm";
+    }
+  }
+
+  function publishBridgeStatus(status) {
+    onBridgeStatus?.(status);
+  }
+
+  // Refreshes the relay's trusted-mac index after the QR bootstrap locks in a phone identity.
+  function sendRelayRegistrationUpdate(nextDeviceState) {
+    deviceState = nextDeviceState;
+    if (socket?.readyState !== WebSocket.OPEN) {
+      return;
+    }
+
+    socket.send(JSON.stringify({
+      kind: "relayMacRegistration",
+      registration: buildMacRegistration(nextDeviceState),
+    }));
+  }
+}
+
+// Registers the canonical Mac identity and the one trusted iPhone allowed for auto-resolve.
+function buildMacRegistrationHeaders(deviceState) {
+  const registration = buildMacRegistration(deviceState);
+  const headers = {
+    "x-mac-device-id": registration.macDeviceId,
+    "x-mac-identity-public-key": registration.macIdentityPublicKey,
+    "x-machine-name": registration.displayName,
+  };
+  if (registration.trustedPhoneDeviceId && registration.trustedPhonePublicKey) {
+    headers["x-trusted-phone-device-id"] = registration.trustedPhoneDeviceId;
+    headers["x-trusted-phone-public-key"] = registration.trustedPhonePublicKey;
+  }
+  return headers;
+}
+
+function buildMacRegistration(deviceState) {
+  const trustedPhoneEntry = Object.entries(deviceState?.trustedPhones || {})[0] || null;
+  return {
+    macDeviceId: normalizeNonEmptyString(deviceState?.macDeviceId),
+    macIdentityPublicKey: normalizeNonEmptyString(deviceState?.macIdentityPublicKey),
+    displayName: normalizeNonEmptyString(os.hostname()),
+    trustedPhoneDeviceId: normalizeNonEmptyString(trustedPhoneEntry?.[0]),
+    trustedPhonePublicKey: normalizeNonEmptyString(trustedPhoneEntry?.[1]),
+  };
+}
+
+function shutdown(codex, getSocket, beforeExit = () => {}) {
+  beforeExit();
+
+  const socket = getSocket();
+  if (socket?.readyState === WebSocket.OPEN || socket?.readyState === WebSocket.CONNECTING) {
+    socket.close();
+  }
+
+  codex.shutdown();
+
+  setTimeout(() => process.exit(0), 100);
+}
+
+function extractBridgeMessageContext(rawMessage) {
+  let parsed = null;
+  try {
+    parsed = JSON.parse(rawMessage);
+  } catch {
+    return { method: "", threadId: null, turnId: null };
+  }
+
+  const method = parsed?.method;
+  const params = parsed?.params;
+  const threadId = extractThreadId(method, params);
+  const turnId = extractTurnId(method, params);
+
+  return {
+    method: typeof method === "string" ? method : "",
+    threadId,
+    turnId,
+  };
+}
+
+function shouldStartContextUsageWatcher(context) {
+  if (!context?.threadId) {
+    return false;
+  }
+
+  return context.method === "turn/start"
+    || context.method === "turn/started";
+}
+
+function extractThreadId(method, params) {
+  if (method === "turn/start" || method === "turn/started") {
+    return (
+      readString(params?.threadId)
+      || readString(params?.thread_id)
+      || readString(params?.turn?.threadId)
+      || readString(params?.turn?.thread_id)
+    );
+  }
+
+  if (method === "thread/start" || method === "thread/started") {
+    return (
+      readString(params?.threadId)
+      || readString(params?.thread_id)
+      || readString(params?.thread?.id)
+      || readString(params?.thread?.threadId)
+      || readString(params?.thread?.thread_id)
+    );
+  }
+
+  if (method === "turn/completed") {
+    return (
+      readString(params?.threadId)
+      || readString(params?.thread_id)
+      || readString(params?.turn?.threadId)
+      || readString(params?.turn?.thread_id)
+    );
+  }
+
+  return null;
+}
+
+function extractTurnId(method, params) {
+  if (method === "turn/started" || method === "turn/completed") {
+    return (
+      readString(params?.turnId)
+      || readString(params?.turn_id)
+      || readString(params?.id)
+      || readString(params?.turn?.id)
+      || readString(params?.turn?.turnId)
+      || readString(params?.turn?.turn_id)
+    );
+  }
+
+  return null;
+}
+
+function readString(value) {
+  return typeof value === "string" && value ? value : null;
+}
+
+function normalizeNonEmptyString(value) {
+  return typeof value === "string" && value.trim() ? value.trim() : "";
+}
+
+module.exports = { startBridge };