runa/vscode/src/client.ts

import { isIP } from "node:net";
import { TextDecoder } from "node:util";
import WebSocket, { RawData } from "ws";
import { stringify } from "lossless-json";
import {
  ActorSnapshot,
  Behavior,
  decodeNotification,
  decodeResponse,
  decodeValue,
  isActorRef,
  MAX_MESSAGE_BYTES,
  PROTOCOL_VERSION,
  ProtocolError,
  Command,
  Reply,
  LifecycleOperation,
  Value,
  WEBSOCKET_SUBPROTOCOL,
} from "./protocol";

export function parseEndpoint(endpoint: string): { host: string; port: number } {
  const match = /^(?:\[([^\]]+)\]|([^:]+)):(\d+)$/.exec(endpoint);
  if (!match) throw new Error("Use a numeric loopback address, e.g. 127.0.0.1:3737 or [::1]:3737");
  const host = match[1] ?? match[2];
  if (
    host === undefined ||
    !((isIP(host) === 4 && host.startsWith("127.")) || (isIP(host) === 6 && host === "::1"))
  ) {
    throw new Error("Only numeric IPv4 loopback or [::1] endpoints are supported");
  }
  const port = Number(match[3]);
  if (!Number.isInteger(port) || port < 1 || port > 65535)
    throw new Error("Invalid WebSocket port");
  return { host, port };
}

interface Pending {
  requestId: string;
  command: Command;
  resolve: (reply: Reply) => void;
  reject: (error: Error) => void;
  timer: NodeJS.Timeout;
}

export class ImageClient {
  private pending?: Pending;
  private queue: Promise<unknown> = Promise.resolve();
  private sequence = 0;
  private revision = 0n;
  private closed = false;
  private ready = false;

  private constructor(
    private readonly socket: WebSocket,
    private readonly onDisconnect: (error?: Error) => void,
    private readonly onImageChanged: (revision: bigint) => void,
    private readonly requestTimeout: number,
  ) {
    socket.on("message", (data, binary) => this.receive(data, binary));
    socket.on("error", (error) => this.close(error));
    socket.on("close", () => this.close(new Error("Image WebSocket closed")));
  }

  static async connect(
    endpoint: string,
    onDisconnect: (error?: Error) => void,
    onImageChanged: (revision: bigint) => void = () => {},
    timeouts = { connect: 5000, request: 30000 },
  ): Promise<ImageClient> {
    parseEndpoint(endpoint);
    const socket = new WebSocket(`ws://${endpoint}/`, WEBSOCKET_SUBPROTOCOL, {
      handshakeTimeout: timeouts.connect,
      maxPayload: MAX_MESSAGE_BYTES,
      perMessageDeflate: false,
      followRedirects: false,
    });
    const client = new ImageClient(socket, onDisconnect, onImageChanged, timeouts.request);
    try {
      await new Promise<void>((resolve, reject) => {
        const fail = (error: Error) => {
          socket.off("open", opened);
          reject(error);
        };
        const opened = () => {
          socket.off("error", fail);
          if (socket.protocol !== WEBSOCKET_SUBPROTOCOL) {
            reject(new Error("Server did not negotiate runa-debug"));
          } else {
            client.ready = true;
            resolve();
          }
        };
        socket.once("error", fail);
        socket.once("open", opened);
      });
    } catch (error) {
      client.close();
      throw error;
    }
    return client;
  }

  async listActors(): Promise<ActorSnapshot[]> {
    const reply = await this.request({ operation: "list_actors" });
    if (reply.kind !== "actors") throw new Error("Expected actor listing");
    return reply.data;
  }

  async inspectActor(actorRef: string): Promise<ActorSnapshot> {
    if (!isActorRef(actorRef)) throw new Error("Invalid actor reference");
    const reply = await this.request({
      operation: "inspect_actor",
      parameters: { actor_ref: actorRef },
    });
    if (reply.kind !== "actor") throw new Error("Expected actor snapshot");
    return reply.data;
  }

  async behaviors(): Promise<Behavior[]> {
    const reply = await this.request({ operation: "get_behaviors" });
    if (reply.kind !== "behaviors") throw new Error("Expected behavior listing");
    return reply.data;
  }

  async spawn(behavior: string, state: Value, capacity = 64): Promise<string> {
    if (!behavior.trim()) throw new Error("Specify a behavior to spawn");
    if (!Number.isSafeInteger(capacity) || capacity < 1)
      throw new Error("Mailbox capacity must be positive");
    const reply = await this.request({
      operation: "spawn",
      parameters: {
        behavior_name: behavior,
        state: decodeValue(state),
        mailbox_capacity: capacity,
      },
    });
    if (reply.kind !== "actor_spawned") throw new Error("Expected spawned actor reference");
    return reply.data.actor_ref;
  }

  async send(actorRef: string, message: Value): Promise<void> {
    if (!isActorRef(actorRef)) throw new Error("Invalid actor reference");
    await this.request({
      operation: "send",
      parameters: { actor_ref: actorRef, message: decodeValue(message) },
    });
  }

  async lifecycle(operation: LifecycleOperation, actorRef: string): Promise<void> {
    if (!isActorRef(actorRef)) throw new Error("Invalid actor reference");
    await this.request({ operation, parameters: { actor_ref: actorRef } });
  }

  async tick(ticks: number, budget = 32): Promise<void> {
    if (!Number.isSafeInteger(ticks) || ticks < 1 || !Number.isSafeInteger(budget) || budget < 1) {
      throw new Error("Ticks and instruction budget must be positive safe integers");
    }
    await this.request({ operation: "tick", parameters: { ticks, instruction_budget: budget } });
  }

  private request(command: Command): Promise<Reply> {
    const request = this.queue.then(
      () =>
        new Promise<Reply>((resolve, reject) => {
          if (this.closed) {
            reject(new Error("Not connected to an image"));
            return;
          }
          const requestId = `runa-vscode-${++this.sequence}`;
          const message = stringify({
            version: PROTOCOL_VERSION,
            request_id: requestId,
            command,
          });
          if (message === undefined || Buffer.byteLength(message) > MAX_MESSAGE_BYTES) {
            reject(new Error("Request exceeds 1 MiB"));
            return;
          }
          const timer = setTimeout(
            () => this.close(new Error("Image request timed out; connection closed without retry")),
            this.requestTimeout,
          );
          this.pending = { requestId, command, resolve, reject, timer };
          this.socket.send(message, (error) => {
            if (error) this.close(error);
          });
        }),
    );
    // Each caller retains its own rejection, without poisoning subsequent requests.
    this.queue = request.then(
      () => undefined,
      () => undefined,
    );
    return request;
  }

  private receive(data: RawData, binary: boolean): void {
    if (this.closed) return;
    try {
      if (binary) throw new Error("Only WebSocket text protocol messages are accepted");
      const bytes = Array.isArray(data)
        ? Buffer.concat(data)
        : data instanceof ArrayBuffer
          ? Buffer.from(data)
          : data;
      if (bytes.length > MAX_MESSAGE_BYTES) throw new Error("Image response exceeds 1 MiB");
      const text = new TextDecoder("utf-8", { fatal: true }).decode(bytes);
      const revision = decodeNotification(text);
      if (revision !== undefined) {
        if (revision <= this.revision)
          throw new Error("Non-increasing image notification revision");
        this.revision = revision;
        this.onImageChanged(revision);
        return;
      }
      const pending = this.pending;
      if (!pending) throw new Error("Unexpected unsolicited image response");
      try {
        pending.resolve(decodeResponse(text, pending.requestId, pending.command));
      } catch (error) {
        if (!(error instanceof ProtocolError)) throw error;
        pending.reject(error);
      }
      this.pending = undefined;
      clearTimeout(pending.timer);
    } catch (error) {
      this.close(error instanceof Error ? error : new Error(String(error)));
    }
  }

  close(error?: Error): void {
    if (this.closed) return;
    this.closed = true;
    this.socket.terminate();
    if (this.pending) {
      clearTimeout(this.pending.timer);
      this.pending.reject(error ?? new Error("Disconnected from image"));
      this.pending = undefined;
    }
    if (this.ready) this.onDisconnect(error);
  }
}