runa/vscode/src/client.ts

import { createConnection, isIP, Socket } from "node:net";
import { TextDecoder } from "node:util";
import {
  ActorSnapshot,
  decodeResponse,
  isActorRef,
  MAX_LINE_BYTES,
  PROTOCOL_VERSION,
  ProtocolError,
  ReadCommand,
  ReadReply,
} 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 TCP port");
  return { host, port };
}

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

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

  private constructor(
    private readonly socket: Socket,
    private readonly onDisconnect: (error?: Error) => void,
    private readonly requestTimeout: number,
  ) {
    socket.on("data", (chunk: Buffer) => this.receive(chunk));
    socket.on("error", (error) => this.close(error));
    socket.on("close", () => this.close(new Error("Image connection closed")));
  }

  static async connect(
    endpoint: string,
    onDisconnect: (error?: Error) => void,
    timeouts = { connect: 5000, request: 30000 },
  ): Promise<ImageClient> {
    const address = parseEndpoint(endpoint);
    const socket = createConnection(address);
    await new Promise<void>((resolve, reject) => {
      const cleanup = () => {
        clearTimeout(timer);
        socket.off("error", fail);
        socket.off("connect", connected);
      };
      const fail = (error: Error) => {
        cleanup();
        socket.destroy();
        reject(error);
      };
      const connected = () => {
        cleanup();
        resolve();
      };
      const timer = setTimeout(
        () => fail(new Error("Image connection timed out")),
        timeouts.connect,
      );
      socket.once("error", fail);
      socket.once("connect", connected);
    });
    return new ImageClient(socket, onDisconnect, timeouts.request);
  }

  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;
  }

  private request(command: ReadCommand): Promise<ReadReply> {
    const request = this.queue.then(
      () =>
        new Promise<ReadReply>((resolve, reject) => {
          if (this.closed) {
            reject(new Error("Not connected to an image"));
            return;
          }
          const requestId = `runa-vscode-${++this.sequence}`;
          const line = JSON.stringify({
            version: PROTOCOL_VERSION,
            request_id: requestId,
            command,
          });
          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.write(line + "\n", (error) => {
            if (error) this.close(error);
          });
        }),
    );
    // A failed request must not prevent later requests; each caller retains its own rejection.
    this.queue = request.then(
      () => undefined,
      () => undefined,
    );
    return request;
  }

  private receive(chunk: Buffer): void {
    if (this.closed) return;
    let offset = 0;
    while (offset < chunk.length && !this.closed) {
      if (!this.pending) {
        this.close(new Error("Unexpected unsolicited image response"));
        return;
      }
      const newline = chunk.indexOf(10, offset);
      const end = newline === -1 ? chunk.length : newline;
      if (this.buffer.length + end - offset > MAX_LINE_BYTES) {
        this.close(new Error("Image response exceeds the 1 MiB line limit"));
        return;
      }
      this.buffer = Buffer.concat([this.buffer, chunk.subarray(offset, end)]);
      if (newline === -1) return;
      const frame = this.buffer;
      this.buffer = Buffer.alloc(0);
      offset = newline + 1;
      const pending = this.pending;
      this.pending = undefined;
      clearTimeout(pending.timer);
      try {
        const line = new TextDecoder("utf-8", { fatal: true }).decode(frame);
        pending.resolve(decodeResponse(line, pending.requestId, pending.command));
      } catch (error) {
        const failure = error instanceof Error ? error : new Error(String(error));
        pending.reject(failure);
        if (!(failure instanceof ProtocolError)) this.close(failure);
      }
    }
  }

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