beat-pocket/extension/content/util/ws-proxy.ts

/**
 * WebSocket instrumentation. We wrap WebSocket.prototype.send so we can:
 *  1. Observe outgoing chat messages (log only, no payload capture for V1)
 *  2. Optionally resend synthetic chat messages if the React dispatch
 *     path is broken (fallback for chat-drawer.send).
 *
 * Idempotent — safe to call multiple times.
 */

interface WSInstrumentation {
  restore: () => void;
  sentCount: () => number;
}

const FLAG = '__beatPocketWSInstrumented__';

export function instrumentWS(): WSInstrumentation {
  const proto = WebSocket.prototype as unknown as Record<string, unknown>;
  if (proto[FLAG]) {
    return existingInstrumentation();
  }
  const originalSend = proto.send as (data: string | ArrayBufferLike | Blob | ArrayBufferView) => void;
  let sent = 0;

  proto.send = function patchedSend(data: unknown): void {
    sent += 1;
    try {
      // Lightweight logging only — we never read or store payload contents.
      // Just observe that WS frames flow so we know auth is alive.
      if (typeof data === 'string') {
        window.dispatchEvent(new CustomEvent('bbp:ws-send', { detail: { len: data.length } }));
      }
    } catch {
      /* never throw out of a wrapped primitive */
    }
    return originalSend.call(this, data as never);
  };

  proto[FLAG] = { originalSend, sent: () => sent, restore: () => restore() };

  function restore() {
    if (proto[FLAG] && (proto[FLAG] as { originalSend?: unknown }).originalSend) {
      proto.send = (proto[FLAG] as { originalSend: typeof originalSend }).originalSend;
      delete proto[FLAG];
    }
  }

  return { restore, sentCount: () => sent };
}

function existingInstrumentation(): WSInstrumentation {
  const proto = WebSocket.prototype as unknown as Record<string, unknown>;
  const flag = proto[FLAG] as { sent: () => number; restore: () => void };
  return { restore: flag.restore, sentCount: flag.sent };
}