xo-cli/src/utils/syncing/ws/client.ts

import { fromExtendedJson, toExtendedJson } from '@xo-cash/utils';
import type { PrivateKey } from '@xo-cash/primitives';

import { WsMessageSchema, WsSuccessResponseSchema, WsErrorResponseSchema, type WsMessage } from './types.js';
import { SyncClient } from '../shared/client.js';

/** A request sent over the sync server's WebSocket connection. */
type WsRequest = {
  id?: string;
  path: string;
  body?: unknown;
  headers?: Record<string, string>;
};

type PendingRequest = {
  resolve: (body: unknown) => void;
  reject: (error: Error) => void;
};

export type WsClientOptions = {
  /** Called for server-pushed events and uncorrelated server errors. */
  onMessage: (message: WsMessage) => void;

  /** Called for transport errors and malformed server messages. */
  onError: (error: Error) => void;
};

/**
 * Minimal WebSocket client for the XO sync server.
 *
 * This class intentionally owns only the basic WebSocket protocol:
 *
 * - one connection is opened at `/ws`;
 * - request IDs correlate concurrent read/write/unsubscribe calls;
 * - subscriptions stay on that connection and receive pushed events; and
 * - a dropped connection fails pending requests but is not reconnected.
 *
 * A production client would normally wrap this class with retry, heartbeat,
 * resubscription, and request-timeout behavior. Those concerns are omitted
 * here to keep the transport demo easy to follow.
 */
export class WsClient extends SyncClient {
  private readonly url: string;
  private readonly privateKey: PrivateKey;
  private socket: WebSocket | undefined;
  private readonly pendingRequests = new Map<string, PendingRequest>();
  private readonly messageListeners = new Set<(message: WsMessage) => void>();
  private readonly errorListeners = new Set<(error: Error) => void>();

  constructor(
    url: string,
    privateKey: PrivateKey,
    options: Partial<WsClientOptions> = {},
  ) {
    super();

    this.url = url;
    this.privateKey = privateKey;
    this.messageListeners.add(options.onMessage ?? (() => {}));
    this.errorListeners.add(options.onError ?? (() => {}));
  }

  /** Open the socket and resolve once the WebSocket handshake completes. */
  async connect(): Promise<void> {
    if (this.socket?.readyState === WebSocket.OPEN) {
      return;
    }

    await new Promise<void>((resolve, reject) => {
      const socket = new WebSocket(WsClient.httpToWsUrl(this.url));
      this.socket = socket;

      socket.onopen = () => resolve();

      socket.onmessage = (event) => {
        this.handleMessage(String(event.data));
      };

      socket.onerror = () => {
        const error = new Error('WebSocket connection failed');
        this.emitError(error);
        reject(error);
      };

      socket.onclose = () => {
        if (this.socket === socket) {
          this.socket = undefined;
        }

        this.rejectPendingRequests(new Error('WebSocket connection closed'));
      };
    });
  }

  /**
   * Close the socket.
   *
   * Closing a WebSocket automatically removes all of its server-side topics,
   * so disconnect does not need to send an unsubscribe request first.
   */
  async disconnect(): Promise<void> {
    const socket = this.socket;
    this.socket = undefined;

    this.rejectPendingRequests(new Error('WebSocket client disconnected'));

    if (!socket || socket.readyState === WebSocket.CLOSED) {
      return;
    }

    // Wait for `close` so callers know the underlying connection is gone.
    await new Promise<void>((resolve) => {
      socket.onclose = () => resolve();
      socket.close();
    });
  }

  /** Read every stored instance of one resource. */
  async read(resourceId: string): Promise<unknown> {
    return this.request('/data/get', { resourceId: [resourceId] });
  }

  /**
   * Sign and write one resource, using the same wire format as the SSE client.
   */
  async write(resourceId: string, value: Record<string, unknown>): Promise<unknown> {
    const valueBytes = new TextEncoder().encode(toExtendedJson(value));
    const resource = SyncClient.signWriteRequest(this.privateKey.toBytes(), resourceId, valueBytes);

    return this.request('/data/write', { resources: [resource] });
  }

  /**
   * Register one resource topic on the current connection.
   *
   * The subscribe route is intentionally long-running and therefore does not
   * send an acknowledgement. This method resolves after the request frame has
   * been handed to the socket; future updates arrive through `onMessage`.
   */
  async subscribe(resourceId: string): Promise<void> {
    this.send({
      path: '/data/subscribe',
      body: { resourceId: [resourceId] },
    });
  }

  /** Remove one resource topic without closing the shared connection. */
  async unsubscribe(resourceId: string): Promise<unknown> {
    return this.request('/data/unsubscribe', { resourceId: [resourceId] });
  }

  /** Send a request and wait for the response carrying the same ID. */
  private async request(path: string, body?: unknown): Promise<unknown> {
    const id = crypto.randomUUID();

    const response = new Promise<unknown>((resolve, reject) => {
      this.pendingRequests.set(id, { resolve, reject });
    });

    // Sign the request body
    const headers = SyncClient.signRequest(this.privateKey.toBytes(), path, toExtendedJson(body));

    try {
      this.send({ id, path, body, headers });
    } catch (error) {
      // Avoid leaving a promise in the map when the socket was not open.
      this.pendingRequests.delete(id);
      throw error;
    }

    return response;
  }

  /** Encode one complete request envelope using the shared Extended JSON codec. */
  private send(request: WsRequest): void {
    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
      throw new Error('WebSocket is not connected');
    }

    this.socket.send(toExtendedJson(request));
  }

  /** Decode and route one frame received from the sync server. */
  private handleMessage(raw: string): void {
    const { success, data: message } = WsMessageSchema.safeParse(fromExtendedJson(raw));
    if (!success) {
      this.emitError(new Error('Invalid WebSocket message'));
      return;
    }

    // Only response/error messages with an outstanding ID are RPC replies.
    // Application events are left for message listeners, even if they have IDs.
    const pending = message.id ? this.pendingRequests.get(message.id) : undefined;

    // If the message has no outstanding ID, it is a subscription event so we will just notify the listeners
    if (!pending) {
      for (const listener of this.messageListeners) {
        listener(message);
      }

      return;
    }

    // Delete the pending request
    this.pendingRequests.delete(message.id!);

    // If the message is an error, reject the pending request
    if (message.type === 'error') {
      const error = WsErrorResponseSchema.parse(message);
      pending.reject(new Error(`${error.error} (${error.statusCode})`));
      return;
    }

    // If the message is a response, resolve the pending request
    if (message.type === 'response') {
      const response = WsSuccessResponseSchema.parse(message);

      // If the response is not successful, reject the pending request
      if (response.statusCode < 200 || response.statusCode >= 300) {
        pending.reject(new Error(`Request failed (${response.statusCode})`));
        return;
      }

      // If the response is successful, resolve the pending request
      pending.resolve(response.body);
    }
  }

  // Helper to emit an error to all error listeners
  private emitError(error: Error): void {
    for (const listener of this.errorListeners) {
      listener(error);
    }
  }

  // Helper for when the client disconnects
  private rejectPendingRequests(error: Error): void {
    for (const pending of this.pendingRequests.values()) {
      pending.reject(error);
    }

    this.pendingRequests.clear();
  }

  // Helper to convert an HTTP URL to a WebSocket URL
  static httpToWsUrl(httpUrl: string): string {
    const url = new URL(httpUrl);

    if (url.protocol === 'http:') {
      url.protocol = 'ws:';
    } else if (url.protocol === 'https:') {
      url.protocol = 'wss:';
    } else if (url.protocol !== 'ws:' && url.protocol !== 'wss:') {
      throw new Error(`Unsupported sync server protocol: ${url.protocol}`);
    }

    // The constructor accepts either the server root or the complete `/ws` URL.
    const pathname = url.pathname.replace(/\/+$/, '');
    url.pathname = pathname.endsWith('/ws') ? pathname : `${pathname}/ws`;

    return url.toString();
  }
}