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,
  ExecutionInspection,
  SourceBreakpoint,
  DebugStopReason,
  STEP_MODES,
  StepMode,
  decodeNotification,
  decodeResponse,
  decodeValue,
  isActorRef,
  isSourcePath,
  isSourceRevision,
  MAX_MESSAGE_BYTES,
  PROTOCOL_VERSION,
  ProtocolError,
  Command,
  Reply,
  SourceCommit,
  SourceFile,
  SourceDiagnostic,
  SystemMessage,
  Value,
  WEBSOCKET_SUBPROTOCOL,
  LogLevel,
  LogEntry,
  LOG_LEVELS,
} 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 };
}

function sourcePath(path: string): string {
  if (!isSourcePath(path)) {
    throw new Error("Source paths look like actors/counter (letters, digits, _ and -; no .rn)");
  }
  return path;
}

function sourceRevision(revision: string): string {
  if (!isSourceRevision(revision)) throw new Error("Invalid source revision");
  return revision;
}

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 closing: Promise<void> = Promise.resolve();

  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 checkSource(source: string): Promise<SourceDiagnostic[]> {
    const reply = await this.request({ operation: "check_source", parameters: { source } });
    if (reply.kind !== "source_checked") throw new Error("Expected source diagnostics");
    return reply.data.diagnostics;
  }

  async actorBehavior(actorRef: string): Promise<{ source: string; source_id: string }> {
    const reply = await this.request({
      operation: "actor_behavior",
      parameters: { actor_ref: actorRef },
    });
    if (reply.kind !== "actor_behavior") throw new Error("Expected actor behavior");
    return reply.data;
  }

  async writeActorBehavior(
    actorRef: string,
    expected: string,
    source: string,
  ): Promise<{ source: string; source_id: string }> {
    const reply = await this.request({
      operation: "write_actor_behavior",
      parameters: { actor_ref: actorRef, expected_source_id: expected, source },
    });
    if (reply.kind !== "actor_behavior") throw new Error("Expected actor behavior");
    return reply.data;
  }

  async updateBehaviorSource(
    actorRef: string,
    expected: string,
    temporary: string,
    path: string,
  ): Promise<SourceFile> {
    const reply = await this.request({
      operation: "update_behavior_source",
      parameters: {
        actor_ref: actorRef,
        expected_source_id: expected,
        temporary_source_id: temporary,
        path: sourcePath(path),
      },
    });
    if (reply.kind !== "source") throw new Error("Expected updated behavior source");
    return reply.data;
  }

  async renameSource(
    path: string,
    target: string,
    expected: string,
    updateActors: boolean,
  ): Promise<SourceFile> {
    const reply = await this.request({
      operation: "rename_source",
      parameters: {
        path: sourcePath(path),
        target: sourcePath(target),
        expected_source_id: expected,
        update_actors: updateActors,
      },
    });
    if (reply.kind !== "source") throw new Error("Expected renamed source");
    return reply.data;
  }

  async deleteSource(path: string, expected: string): Promise<void> {
    await this.request({
      operation: "delete_source",
      parameters: { path: sourcePath(path), expected_source_id: expected },
    });
  }

  async logs(actorRef: string): Promise<LogEntry[]> {
    const entries: LogEntry[] = [];
    let after = 0n;
    for (;;) {
      const reply = await this.request({
        operation: "logs",
        parameters: { actors: [actorRef], follow: false, after },
      });
      if (reply.kind !== "logs") throw new Error("Expected actor logs");
      entries.push(...reply.data.entries);
      if (!reply.data.has_more) return entries;
      if (reply.data.next_sequence <= after) throw new Error("Non-increasing log cursor");
      after = reply.data.next_sequence;
    }
  }

  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 inspectExecution(actorRef: string): Promise<ExecutionInspection> {
    if (!isActorRef(actorRef)) throw new Error("Invalid actor reference");
    const reply = await this.request({
      operation: "inspect_execution",
      parameters: { actor_ref: actorRef },
    });
    if (reply.kind !== "execution_inspected") throw new Error("Expected execution inspection");
    return reply.data;
  }

  async setBreakpoints(
    actorRef: string,
    implementationId: string,
    lines: number[],
  ): Promise<SourceBreakpoint[]> {
    if (
      !isActorRef(actorRef) ||
      !implementationId ||
      lines.length > 128 ||
      lines.some((line) => !Number.isSafeInteger(line) || line < 0 || line > 4294967295)
    )
      throw new Error("Invalid actor, implementation or breakpoint lines");
    const reply = await this.request({
      operation: "set_breakpoints",
      parameters: { actor_ref: actorRef, implementation_id: implementationId, lines },
    });
    if (reply.kind !== "breakpoints_set") throw new Error("Expected breakpoint verification");
    return reply.data.breakpoints;
  }

  async stepSource(
    actorRef: string,
    implementationId: string,
    budget = 1000,
    mode: StepMode = "into",
  ): Promise<DebugStopReason> {
    if (
      !isActorRef(actorRef) ||
      !implementationId ||
      !Number.isSafeInteger(budget) ||
      budget < 1 ||
      budget > 10000 ||
      !STEP_MODES.includes(mode)
    )
      throw new Error("Invalid actor, implementation or source step budget");
    const reply = await this.request({
      operation: "step_source",
      parameters: {
        actor_ref: actorRef,
        implementation_id: implementationId,
        instruction_budget: budget,
        mode,
      },
    });
    if (reply.kind !== "debug_step_completed") throw new Error("Expected source step result");
    return reply.data.reason;
  }
  async listSources(): Promise<SourceFile[]> {
    const reply = await this.request({ operation: "list_sources" });
    if (reply.kind !== "sources") throw new Error("Expected source listing");
    return reply.data;
  }

  async readSource(path: string): Promise<SourceFile> {
    const reply = await this.request({
      operation: "read_source",
      parameters: { path: sourcePath(path) },
    });
    if (reply.kind !== "source") throw new Error("Expected source file");
    return reply.data;
  }

  /** Saves working source only if the image still holds `expectedSourceId` (null creates). */
  async writeSource(
    path: string,
    expectedSourceId: string | null,
    source: string,
  ): Promise<SourceFile> {
    if (expectedSourceId !== null && (typeof expectedSourceId !== "string" || !expectedSourceId)) {
      throw new Error("Expected source identity must be a non-empty string or null");
    }
    if (typeof source !== "string") throw new Error("Source must be text");
    const reply = await this.request({
      operation: "write_source",
      parameters: { path: sourcePath(path), expected_source_id: expectedSourceId, source },
    });
    if (reply.kind !== "source") throw new Error("Expected source file");
    return reply.data;
  }

  async commitSource(path: string): Promise<SourceCommit> {
    const reply = await this.request({
      operation: "commit_source",
      parameters: { path: sourcePath(path) },
    });
    if (reply.kind !== "source_committed") throw new Error("Expected source commit");
    return reply.data;
  }

  async sourceHistory(path: string): Promise<SourceCommit[]> {
    const reply = await this.request({
      operation: "source_history",
      parameters: { path: sourcePath(path) },
    });
    if (reply.kind !== "source_history") throw new Error("Expected source history");
    return reply.data;
  }

  async revisionSource(path: string, revision: string): Promise<string> {
    const reply = await this.request({
      operation: "revision_source",
      parameters: { path: sourcePath(path), revision: sourceRevision(revision) },
    });
    if (reply.kind !== "behavior_source") throw new Error("Expected revision source");
    return reply.data.source;
  }

  async rollbackSource(path: string, revision: string): Promise<SourceFile> {
    const reply = await this.request({
      operation: "rollback_source",
      parameters: { path: sourcePath(path), revision: sourceRevision(revision) },
    });
    if (reply.kind !== "source") throw new Error("Expected source file");
    return reply.data;
  }

  async spawn(path: string, state: Value, capacity = 64): Promise<string> {
    if (!Number.isSafeInteger(capacity) || capacity < 1)
      throw new Error("Mailbox capacity must be positive");
    const reply = await this.request({
      operation: "spawn",
      parameters: { path: sourcePath(path), 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 setLogLevel(actorRef: string, level: LogLevel): Promise<void> {
    if (!isActorRef(actorRef) || !LOG_LEVELS.includes(level))
      throw new Error("Invalid actor or log level");
    await this.request({ operation: "set_log_level", parameters: { actor_ref: actorRef, level } });
  }

  async editActor(
    actorRef: string,
    state: Value | null,
    messages: { reserved: Value | null; queued: Value[] } | null,
  ): Promise<void> {
    if (!isActorRef(actorRef)) throw new Error("Invalid actor reference");
    await this.request({
      operation: "edit_actor",
      parameters: {
        actor_ref: actorRef,
        state: state === null ? null : decodeValue(state),
        messages:
          messages === null
            ? null
            : {
                reserved: messages.reserved === null ? null : decodeValue(messages.reserved),
                queued: messages.queued.map(decodeValue),
              },
      },
    });
  }

  async system(actorRef: string, message: SystemMessage): Promise<void> {
    if (!isActorRef(actorRef)) throw new Error("Invalid actor reference");
    const kinds = ["stop", "kill", "restart", "pause", "resume", "reload", "become"];
    if (typeof message !== "object" || message === null || !kinds.includes(message.kind)) {
      throw new Error("Unsupported system message");
    }
    const checked: SystemMessage =
      message.kind === "become"
        ? { kind: "become", path: sourcePath(message.path) }
        : { kind: message.kind };
    await this.request({
      operation: "system",
      parameters: { actor_ref: actorRef, message: checked },
    });
  }

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

  async saveImage(file: string): Promise<void> {
    await this.snapshotFile("save_image", file);
  }

  async restoreImage(file: string): Promise<void> {
    await this.snapshotFile("restore_image", file);
  }

  private async snapshotFile(
    operation: "save_image" | "restore_image",
    file: string,
  ): Promise<void> {
    if (!file.trim() || Buffer.byteLength(file, "utf8") > 4096 || file.includes("\0"))
      throw new Error("Snapshot file must be a nonempty server-local path of at most 4096 bytes");
    await this.request({ operation, parameters: { file } });
  }

  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): Promise<void> {
    if (this.closed) return this.closing;
    this.closed = true;
    this.closing = new Promise<void>((resolve) => {
      if (this.socket.readyState === WebSocket.CLOSED) {
        resolve();
        return;
      }
      const timer = setTimeout(() => this.socket.terminate(), 5000);
      timer.unref();
      this.socket.once("close", () => {
        clearTimeout(timer);
        resolve();
      });
      if (error || this.socket.readyState === WebSocket.CONNECTING) {
        this.socket.terminate();
      } else {
        this.socket.close(1000);
      }
    });
    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);
    return this.closing;
  }
}