ai/gateway/src/clients/zmq-protocol.ts

/**
 * ZMQ Protocol encoding/decoding using Protobuf
 *
 * Protocol format (as defined in protobuf/ingestor.proto):
 * Frame 1: [1 byte: protocol version]
 * Frame 2: [1 byte: message type ID][N bytes: protobuf message]
 *
 * For PUB/SUB: [topic frame][version frame][message frame]
 */

import protobuf from 'protobufjs';
import { readFileSync } from 'fs';
import { fileURLToPath } from 'url';
import { dirname, join } from 'path';
import type {
  SubmitHistoricalRequest,
  SubmitResponse,
  HistoryReadyNotification,
  SubmitStatus,
  NotificationStatus,
} from '../types/ohlc.js';

const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);

/**
 * Protocol constants
 */
export const PROTOCOL_VERSION = 0x01;

export enum MessageType {
  SUBMIT_HISTORICAL_REQUEST = 0x10,
  SUBMIT_RESPONSE = 0x11,
  HISTORY_READY_NOTIFICATION = 0x12,
}

// Load protobuf types at runtime (same pattern as ingestor)
// Proto files are copied to /app/protobuf/ in the Docker image
const protoDir = join(__dirname, '../..', 'protobuf');
const root = new protobuf.Root();

// Load proto file and parse it
const ingestorProto = readFileSync(join(protoDir, 'ingestor.proto'), 'utf8');
protobuf.parse(ingestorProto, root);

// Export message types
const SubmitHistoricalRequestType = root.lookupType('SubmitHistoricalRequest');
const SubmitResponseType = root.lookupType('SubmitResponse');
const HistoryReadyNotificationType = root.lookupType('HistoryReadyNotification');

/**
 * Encode SubmitHistoricalRequest to ZMQ frames
 *
 * Returns: [version_frame, message_frame]
 */
export function encodeSubmitHistoricalRequest(req: SubmitHistoricalRequest): Buffer[] {
  const versionFrame = Buffer.from([PROTOCOL_VERSION]);

  // Convert to protobuf-compatible format (pbjs uses camelCase)
  // Note: protobufjs handles bigint/number conversion automatically for uint64
  const protoMessage = {
    requestId: req.request_id,
    ticker: req.ticker,
    startTime: Number(req.start_time),  // Convert bigint to number for protobuf
    endTime: Number(req.end_time),
    periodSeconds: req.period_seconds,
    limit: req.limit,
    clientId: req.client_id,
  };

  // Encode as protobuf
  const message = SubmitHistoricalRequestType.create(protoMessage);
  const payloadBuffer = SubmitHistoricalRequestType.encode(message).finish();

  const messageFrame = Buffer.concat([
    Buffer.from([MessageType.SUBMIT_HISTORICAL_REQUEST]),
    Buffer.from(payloadBuffer),
  ]);

  return [versionFrame, messageFrame];
}

/**
 * Decode SubmitResponse from ZMQ frames
 *
 * Input: [version_frame, message_frame]
 */
export function decodeSubmitResponse(frames: Buffer[]): SubmitResponse {
  try {
    if (frames.length < 2) {
      throw new Error(`Expected 2 frames, got ${frames.length}`);
    }

    const versionFrame = frames[0];
    const messageFrame = frames[1];

    // Validate version
    if (versionFrame[0] !== PROTOCOL_VERSION) {
      throw new Error(`Unsupported protocol version: ${versionFrame[0]}`);
    }

    // Validate message type
    const messageType = messageFrame[0];
    if (messageType !== MessageType.SUBMIT_RESPONSE) {
      throw new Error(`Expected SUBMIT_RESPONSE (0x11), got 0x${messageType.toString(16)}`);
    }

    // Decode protobuf payload
    const payloadBuffer = messageFrame.slice(1);
    const decoded = SubmitResponseType.decode(payloadBuffer);
    const payload = SubmitResponseType.toObject(decoded, {
      longs: String,
      enums: Number,  // Keep enums as numbers for comparison
      defaults: true,
    });

    return {
      request_id: payload.requestId,
      status: payload.status as SubmitStatus,
      error_message: payload.errorMessage || undefined,
      notification_topic: payload.notificationTopic,
    };
  } catch (error) {
    console.error('Error decoding SubmitResponse:', error);
    console.error('Frame count:', frames.length);
    if (frames.length >= 2) {
      console.error('Version frame:', frames[0].toString('hex'));
      console.error('Message frame (first 100 bytes):', frames[1].slice(0, 100).toString('hex'));
    }
    throw error;
  }
}

/**
 * Decode HistoryReadyNotification from ZMQ frames
 *
 * Input: [topic_frame, version_frame, message_frame] (for SUB socket)
 */
export function decodeHistoryReadyNotification(frames: Buffer[]): HistoryReadyNotification {
  if (frames.length < 3) {
    throw new Error(`Expected 3 frames (topic, version, message), got ${frames.length}`);
  }

  const versionFrame = frames[1];
  const messageFrame = frames[2];

  // Validate version
  if (versionFrame[0] !== PROTOCOL_VERSION) {
    throw new Error(`Unsupported protocol version: ${versionFrame[0]}`);
  }

  // Validate message type
  const messageType = messageFrame[0];
  if (messageType !== MessageType.HISTORY_READY_NOTIFICATION) {
    throw new Error(`Expected HISTORY_READY_NOTIFICATION (0x12), got 0x${messageType.toString(16)}`);
  }

  // Decode protobuf payload
  const payloadBuffer = messageFrame.slice(1);
  const decoded = HistoryReadyNotificationType.decode(payloadBuffer);
  const payload = HistoryReadyNotificationType.toObject(decoded, {
    longs: String,
    enums: Number,  // Keep enums as numbers for comparison
    defaults: true,
  });

  return {
    request_id: payload.requestId,
    ticker: payload.ticker,
    period_seconds: payload.periodSeconds,
    start_time: BigInt(payload.startTime),
    end_time: BigInt(payload.endTime),
    status: payload.status as NotificationStatus,
    error_message: payload.errorMessage || undefined,
    iceberg_namespace: payload.icebergNamespace,
    iceberg_table: payload.icebergTable,
    row_count: payload.rowCount,
    completed_at: BigInt(payload.completedAt),
  };
}