Add files using upload-large-folder tool

4e23b01 verified about 2 months ago

8.84 kB

/**
 * WS layer for `DaemonClient` — owns the socket, queues incoming frames so
 * fast tests don't race the first `server_hello`, exposes a `waitForFrame`
 * with timeouts, and routes control-message acks back to the original sender
 * via `id` correlation.
 *
 * Frame shape is the union of WS.md §2 envelopes:
 *   - `event` envelope     : `{type, seq, session_id, timestamp, payload}`
 *   - `ack`                : `{type:'ack', id, code, msg, payload}`
 *   - `server_hello`/`ping`/`resync_required`/`error`: each carries `timestamp`
 *
 * We don't Zod-validate frames here — preserving forward-compat ("unknown
 * fields pass through") and avoiding double-work since the server emits the
 * shapes already.
 */
import { WebSocket as WsWebSocket } from 'ws';

import { recordReportEvent } from './report.js';

/** Wire frame shape — kept loose because the server adds new event types. */
export interface AnyFrame {
  readonly type: string;
  readonly seq?: number;
  readonly session_id?: string;
  readonly timestamp?: string;
  readonly id?: string;
  readonly code?: number;
  readonly msg?: string;
  readonly payload?: unknown;
}

export interface WsClientOptions {
  url: string;
  wsImpl: typeof WsWebSocket;
  logger: (level: 'info' | 'warn' | 'error' | 'debug', msg: string, meta?: unknown) => void;
  reportDir?: string;
}

type FrameWaiter = (frame: AnyFrame) => boolean;

interface PendingWaiter {
  match: FrameWaiter;
  resolve: (frame: AnyFrame) => void;
  reject: (err: Error) => void;
  timer?: NodeJS.Timeout;
}

/**
 * Thin WS wrapper. Two-tier delivery:
 *   - All frames also fan out to subscribers added with `onFrame()`.
 *   - `waitForFrame(predicate)` consumes the *first matching* frame; matching
 *     frames already in `_queue` are dispatched immediately.
 *
 * Both queue and waiters are needed because the server's first `server_hello`
 * can land in the same tick as `open`, before the test has a chance to
 * register its first waiter (see `server/test/ws-handshake.e2e.test.ts:88-117`
 * for the pattern this is ported from).
 */
export class WsClient {
  private ws: WsWebSocket | null = null;
  private readonly _queue: AnyFrame[] = [];
  private readonly _waiters: PendingWaiter[] = [];
  private readonly _subscribers = new Set<(f: AnyFrame) => void>();
  private _closed = false;
  private _closeReason: { code: number; reason: string } | null = null;
  private _closeWaiters: Array<(v: { code: number; reason: string }) => void> = [];

  constructor(private readonly opts: WsClientOptions) {}

  /** Open the socket; resolves once `open` fires. */
  async open(): Promise<void> {
    if (this.ws) return;
    await new Promise<void>((resolve, reject) => {
      const ws = new this.opts.wsImpl(this.opts.url);
      this.ws = ws;
      ws.once('open', () => {
        recordReportEvent(
          { kind: 'ws', direction: 'lifecycle', url: this.opts.url, message: 'open' },
          { reportDir: this.opts.reportDir },
        );
        resolve();
      });
      ws.once('error', (err) => {
        if (this._closed) return;
        recordReportEvent(
          {
            kind: 'ws',
            direction: 'lifecycle',
            url: this.opts.url,
            message: 'error',
            error: errorForReport(err),
          },
          { reportDir: this.opts.reportDir },
        );
        reject(err as Error);
      });
      ws.on('message', (data) => this._onMessage(data));
      ws.on('close', (code, reason) => this._onClose(code, String(reason ?? '')));
    });
  }

  /** JSON-stringifies and sends a frame. */
  send(frame: object): void {
    if (!this.ws) throw new Error('ws not open');
    this.ws.send(JSON.stringify(frame));
    recordReportEvent(
      { kind: 'ws', direction: 'out', url: this.opts.url, frame },
      { reportDir: this.opts.reportDir },
    );
  }

  /** Register a frame subscriber. Returns an unsubscribe handle. */
  onFrame(handler: (f: AnyFrame) => void): () => void {
    this._subscribers.add(handler);
    return () => {
      this._subscribers.delete(handler);
    };
  }

  /**
   * Wait for the next frame matching `predicate`. Drains queued frames first;
   * the first matching frame is consumed and returned. Times out cleanly.
   */
  waitForFrame(predicate: FrameWaiter, timeoutMs: number): Promise<AnyFrame> {
    return new Promise((resolve, reject) => {
      // Drain queue.
      for (let i = 0; i < this._queue.length; i++) {
        const frame = this._queue[i];
        if (frame === undefined) continue;
        if (predicate(frame)) {
          this._queue.splice(i, 1);
          resolve(frame);
          return;
        }
      }
      if (this._closed) {
        reject(new Error(`ws closed before matching frame arrived (code=${this._closeReason?.code})`));
        return;
      }
      const waiter: PendingWaiter = {
        match: predicate,
        resolve: (f) => {
          if (waiter.timer) clearTimeout(waiter.timer);
          resolve(f);
        },
        reject: (e) => {
          if (waiter.timer) clearTimeout(waiter.timer);
          reject(e);
        },
      };
      waiter.timer = setTimeout(() => {
        const idx = this._waiters.indexOf(waiter);
        if (idx >= 0) this._waiters.splice(idx, 1);
        reject(new Error(`waitForFrame timed out after ${timeoutMs}ms`));
      }, timeoutMs);
      waiter.timer.unref?.();
      this._waiters.push(waiter);
    });
  }

  /** Send a control message and wait for its `ack` (matched by `id`). */
  async sendAndAwaitAck(frame: { type: string; id: string; payload: unknown }, timeoutMs: number): Promise<AnyFrame> {
    this.send(frame);
    return this.waitForFrame(
      (f) => f.type === 'ack' && f.id === frame.id,
      timeoutMs,
    );
  }

  /** Resolves when the socket closes (or immediately if already closed). */
  closed(): Promise<{ code: number; reason: string }> {
    if (this._closeReason) return Promise.resolve(this._closeReason);
    return new Promise((resolve) => {
      this._closeWaiters.push(resolve);
    });
  }

  /** Initiate close from the client side. */
  async close(): Promise<void> {
    if (!this.ws || this._closed) return;
    this.ws.close();
    await this.closed();
  }

  private _onMessage(data: unknown): void {
    let frame: AnyFrame;
    try {
      const raw = typeof data === 'string' ? data : String(data);
      frame = JSON.parse(raw) as AnyFrame;
    } catch (err) {
      this.opts.logger('warn', 'ws: dropped non-JSON frame', { err: String(err) });
      recordReportEvent(
        {
          kind: 'ws',
          direction: 'in',
          url: this.opts.url,
          message: 'dropped non-JSON frame',
          error: errorForReport(err),
        },
        { reportDir: this.opts.reportDir },
      );
      return;
    }
    recordReportEvent(
      { kind: 'ws', direction: 'in', url: this.opts.url, frame },
      { reportDir: this.opts.reportDir },
    );

    if (frame.type === 'ping') {
      this.send({ type: 'pong', payload: { nonce: pingNonce(frame) } });
    }

    // Dispatch to subscribers first — they observe every frame, regardless of
    // whether a `waitForFrame` consumed it.
    for (const sub of this._subscribers) {
      try {
        sub(frame);
      } catch (err) {
        this.opts.logger('warn', 'ws: subscriber threw', { err: String(err) });
      }
    }

    // Find the FIRST waiter whose predicate matches. A waiter is single-shot.
    for (let i = 0; i < this._waiters.length; i++) {
      const w = this._waiters[i];
      if (w === undefined) continue;
      let matches = false;
      try {
        matches = w.match(frame);
      } catch (err) {
        this.opts.logger('warn', 'ws: waiter predicate threw', { err: String(err) });
      }
      if (matches) {
        this._waiters.splice(i, 1);
        w.resolve(frame);
        return;
      }
    }
    this._queue.push(frame);
  }

  private _onClose(code: number, reason: string): void {
    this._closed = true;
    this._closeReason = { code, reason };
    recordReportEvent(
      {
        kind: 'ws',
        direction: 'lifecycle',
        url: this.opts.url,
        message: 'close',
        frame: { code, reason },
      },
      { reportDir: this.opts.reportDir },
    );
    for (const w of this._waiters.splice(0)) {
      w.reject(new Error(`ws closed (code=${code}) before matching frame arrived`));
    }
    for (const w of this._closeWaiters.splice(0)) w(this._closeReason);
  }
}

function pingNonce(frame: AnyFrame): string {
  const payload = frame.payload as { nonce?: unknown } | undefined;
  return typeof payload?.nonce === 'string' ? payload.nonce : '';
}

function errorForReport(error: unknown): unknown {
  if (error instanceof Error) {
    return {
      name: error.name,
      message: error.message,
    };
  }
  return error;
}

	/**
	* WS layer for `DaemonClient` — owns the socket, queues incoming frames so
	* fast tests don't race the first `server_hello`, exposes a `waitForFrame`
	* with timeouts, and routes control-message acks back to the original sender
	* via `id` correlation.
	*
	* Frame shape is the union of WS.md §2 envelopes:
	* - `event` envelope : `{type, seq, session_id, timestamp, payload}`
	* - `ack` : `{type:'ack', id, code, msg, payload}`
	* - `server_hello`/`ping`/`resync_required`/`error`: each carries `timestamp`
	*
	* We don't Zod-validate frames here — preserving forward-compat ("unknown
	* fields pass through") and avoiding double-work since the server emits the
	* shapes already.
	*/
	import { WebSocket as WsWebSocket } from 'ws';

	import { recordReportEvent } from './report.js';

	/** Wire frame shape — kept loose because the server adds new event types. */
	export interface AnyFrame {
	readonly type: string;
	readonly seq?: number;
	readonly session_id?: string;
	readonly timestamp?: string;
	readonly id?: string;
	readonly code?: number;
	readonly msg?: string;
	readonly payload?: unknown;
	}

	export interface WsClientOptions {
	url: string;
	wsImpl: typeof WsWebSocket;
	logger: (level: 'info' \| 'warn' \| 'error' \| 'debug', msg: string, meta?: unknown) => void;
	reportDir?: string;
	}

	type FrameWaiter = (frame: AnyFrame) => boolean;

	interface PendingWaiter {
	match: FrameWaiter;
	resolve: (frame: AnyFrame) => void;
	reject: (err: Error) => void;
	timer?: NodeJS.Timeout;
	}

	/**
	* Thin WS wrapper. Two-tier delivery:
	* - All frames also fan out to subscribers added with `onFrame()`.
	* - `waitForFrame(predicate)` consumes the first matching frame; matching
	* frames already in `_queue` are dispatched immediately.
	*
	* Both queue and waiters are needed because the server's first `server_hello`
	* can land in the same tick as `open`, before the test has a chance to
	* register its first waiter (see `server/test/ws-handshake.e2e.test.ts:88-117`
	* for the pattern this is ported from).
	*/
	export class WsClient {
	private ws: WsWebSocket \| null = null;
	private readonly _queue: AnyFrame[] = [];
	private readonly _waiters: PendingWaiter[] = [];
	private readonly _subscribers = new Set<(f: AnyFrame) => void>();
	private _closed = false;
	private _closeReason: { code: number; reason: string } \| null = null;
	private _closeWaiters: Array<(v: { code: number; reason: string }) => void> = [];

	constructor(private readonly opts: WsClientOptions) {}

	/** Open the socket; resolves once `open` fires. */
	async open(): Promise<void> {
	if (this.ws) return;
	await new Promise<void>((resolve, reject) => {
	const ws = new this.opts.wsImpl(this.opts.url);
	this.ws = ws;
	ws.once('open', () => {
	recordReportEvent(
	{ kind: 'ws', direction: 'lifecycle', url: this.opts.url, message: 'open' },
	{ reportDir: this.opts.reportDir },
	);
	resolve();
	});
	ws.once('error', (err) => {
	if (this._closed) return;
	recordReportEvent(
	{
	kind: 'ws',
	direction: 'lifecycle',
	url: this.opts.url,
	message: 'error',
	error: errorForReport(err),
	},
	{ reportDir: this.opts.reportDir },
	);
	reject(err as Error);
	});
	ws.on('message', (data) => this._onMessage(data));
	ws.on('close', (code, reason) => this._onClose(code, String(reason ?? '')));
	});
	}

	/** JSON-stringifies and sends a frame. */
	send(frame: object): void {
	if (!this.ws) throw new Error('ws not open');
	this.ws.send(JSON.stringify(frame));
	recordReportEvent(
	{ kind: 'ws', direction: 'out', url: this.opts.url, frame },
	{ reportDir: this.opts.reportDir },
	);
	}

	/** Register a frame subscriber. Returns an unsubscribe handle. */
	onFrame(handler: (f: AnyFrame) => void): () => void {
	this._subscribers.add(handler);
	return () => {
	this._subscribers.delete(handler);
	};
	}

	/**
	* Wait for the next frame matching `predicate`. Drains queued frames first;
	* the first matching frame is consumed and returned. Times out cleanly.
	*/
	waitForFrame(predicate: FrameWaiter, timeoutMs: number): Promise<AnyFrame> {
	return new Promise((resolve, reject) => {
	// Drain queue.
	for (let i = 0; i < this._queue.length; i++) {
	const frame = this._queue[i];
	if (frame === undefined) continue;
	if (predicate(frame)) {
	this._queue.splice(i, 1);
	resolve(frame);
	return;
	}
	}
	if (this._closed) {
	reject(new Error(`ws closed before matching frame arrived (code=${this._closeReason?.code})`));
	return;
	}
	const waiter: PendingWaiter = {
	match: predicate,
	resolve: (f) => {
	if (waiter.timer) clearTimeout(waiter.timer);
	resolve(f);
	},
	reject: (e) => {
	if (waiter.timer) clearTimeout(waiter.timer);
	reject(e);
	},
	};
	waiter.timer = setTimeout(() => {
	const idx = this._waiters.indexOf(waiter);
	if (idx >= 0) this._waiters.splice(idx, 1);
	reject(new Error(`waitForFrame timed out after ${timeoutMs}ms`));
	}, timeoutMs);
	waiter.timer.unref?.();
	this._waiters.push(waiter);
	});
	}

	/** Send a control message and wait for its `ack` (matched by `id`). */
	async sendAndAwaitAck(frame: { type: string; id: string; payload: unknown }, timeoutMs: number): Promise<AnyFrame> {
	this.send(frame);
	return this.waitForFrame(
	(f) => f.type === 'ack' && f.id === frame.id,
	timeoutMs,
	);
	}

	/** Resolves when the socket closes (or immediately if already closed). */
	closed(): Promise<{ code: number; reason: string }> {
	if (this._closeReason) return Promise.resolve(this._closeReason);
	return new Promise((resolve) => {
	this._closeWaiters.push(resolve);
	});
	}

	/** Initiate close from the client side. */
	async close(): Promise<void> {
	if (!this.ws \|\| this._closed) return;
	this.ws.close();
	await this.closed();
	}

	private _onMessage(data: unknown): void {
	let frame: AnyFrame;
	try {
	const raw = typeof data === 'string' ? data : String(data);
	frame = JSON.parse(raw) as AnyFrame;
	} catch (err) {
	this.opts.logger('warn', 'ws: dropped non-JSON frame', { err: String(err) });
	recordReportEvent(
	{
	kind: 'ws',
	direction: 'in',
	url: this.opts.url,
	message: 'dropped non-JSON frame',
	error: errorForReport(err),
	},
	{ reportDir: this.opts.reportDir },
	);
	return;
	}
	recordReportEvent(
	{ kind: 'ws', direction: 'in', url: this.opts.url, frame },
	{ reportDir: this.opts.reportDir },
	);

	if (frame.type === 'ping') {
	this.send({ type: 'pong', payload: { nonce: pingNonce(frame) } });
	}

	// Dispatch to subscribers first — they observe every frame, regardless of
	// whether a `waitForFrame` consumed it.
	for (const sub of this._subscribers) {
	try {
	sub(frame);
	} catch (err) {
	this.opts.logger('warn', 'ws: subscriber threw', { err: String(err) });
	}
	}

	// Find the FIRST waiter whose predicate matches. A waiter is single-shot.
	for (let i = 0; i < this._waiters.length; i++) {
	const w = this._waiters[i];
	if (w === undefined) continue;
	let matches = false;
	try {
	matches = w.match(frame);
	} catch (err) {
	this.opts.logger('warn', 'ws: waiter predicate threw', { err: String(err) });
	}
	if (matches) {
	this._waiters.splice(i, 1);
	w.resolve(frame);
	return;
	}
	}
	this._queue.push(frame);
	}

	private _onClose(code: number, reason: string): void {
	this._closed = true;
	this._closeReason = { code, reason };
	recordReportEvent(
	{
	kind: 'ws',
	direction: 'lifecycle',
	url: this.opts.url,
	message: 'close',
	frame: { code, reason },
	},
	{ reportDir: this.opts.reportDir },
	);
	for (const w of this._waiters.splice(0)) {
	w.reject(new Error(`ws closed (code=${code}) before matching frame arrived`));
	}
	for (const w of this._closeWaiters.splice(0)) w(this._closeReason);
	}
	}

	function pingNonce(frame: AnyFrame): string {
	const payload = frame.payload as { nonce?: unknown } \| undefined;
	return typeof payload?.nonce === 'string' ? payload.nonce : '';
	}

	function errorForReport(error: unknown): unknown {
	if (error instanceof Error) {
	return {
	name: error.name,
	message: error.message,
	};
	}
	return error;
	}