runa/vscode/src/protocol.ts

import { parse, stringify } from "lossless-json";

export const PROTOCOL_VERSION = 3;
export const MAX_MESSAGE_BYTES = 1024 * 1024;
export const WEBSOCKET_SUBPROTOCOL = "runa-debug";
const ACTOR_REF = /^[A-Za-z0-9_-]{21}$/;

export type Value =
  | { type: "null" }
  | { type: "bool"; data: boolean }
  | { type: "integer"; data: bigint }
  | { type: "string" | "actor_ref"; data: string }
  | { type: "list"; data: Value[] }
  | { type: "record"; data: Record<string, Value> };

export interface ActorSnapshot {
  actor_ref: string;
  mailbox_id: bigint;
  implementation_id: string;
  incarnation: bigint;
  state: Value;
  status: {
    status: "idle" | "running" | "awaiting" | "paused" | "stopped" | "failed";
    message?: string;
  };
  execution: {
    vm_steps: bigint;
    instruction_pointer: bigint;
    call_depth: bigint;
  } | null;
  queued: Value[];
  reserved: Value | null;
}

export type ReadCommand =
  { operation: "list_actors" } | { operation: "inspect_actor"; parameters: { actor_ref: string } };

export class ProtocolError extends Error {
  constructor(
    public readonly code: string,
    message: string,
  ) {
    super(`${code}: ${message}`);
  }
}

export function isActorRef(value: unknown): value is string {
  return typeof value === "string" && ACTOR_REF.test(value);
}

function isObject(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function object(value: unknown): Record<string, unknown> {
  if (!isObject(value)) {
    throw new Error("Expected a protocol object");
  }
  return value;
}

function text(value: unknown): string {
  if (typeof value !== "string") throw new Error("Expected a protocol string");
  return value;
}

function uint(value: unknown): bigint {
  if (typeof value !== "bigint" || value < 0n || value > 18446744073709551615n) {
    throw new Error("Expected an unsigned 64-bit protocol integer");
  }
  return value;
}

function array(value: unknown): unknown[] {
  if (!Array.isArray(value)) throw new Error("Expected a protocol array");
  return value;
}

export function decodeValue(input: unknown): Value {
  let nodes = 0;
  let bytes = 0;
  const visit = (input: unknown, depth: number): Value => {
    if (++nodes > 4096 || depth > 32) throw new Error("Value exceeds node/depth limits");
    const value = object(input);
    if (Object.keys(value).some((key) => key !== "type" && key !== "data")) {
      throw new Error("Unknown portable value field");
    }
    const type = value.type;
    switch (type) {
      case "null":
        if ("data" in value) throw new Error("Null must not carry data");
        return { type };
      case "bool":
        if (typeof value.data !== "boolean") throw new Error("Invalid boolean value");
        return { type, data: value.data };
      case "integer":
        if (
          typeof value.data !== "bigint" ||
          value.data < -9223372036854775808n ||
          value.data > 9223372036854775807n
        ) {
          throw new Error("Invalid signed 64-bit value");
        }
        return { type, data: value.data };
      case "string": {
        const data = text(value.data);
        bytes += Buffer.byteLength(data);
        if (bytes > MAX_MESSAGE_BYTES) throw new Error("Value exceeds string/key byte limit");
        return { type, data };
      }
      case "actor_ref":
        if (!isActorRef(value.data)) throw new Error("Invalid actor reference value");
        return { type, data: value.data };
      case "list":
        return { type, data: array(value.data).map((child) => visit(child, depth + 1)) };
      case "record": {
        const entries = Object.entries(object(value.data)).map(([key, child]) => {
          bytes += Buffer.byteLength(key);
          if (bytes > MAX_MESSAGE_BYTES) throw new Error("Value exceeds string/key byte limit");
          return [key, visit(child, depth + 1)];
        });
        return { type, data: Object.fromEntries(entries) };
      }
      default:
        throw new Error(`Unsupported portable value type: ${String(type)}`);
    }
  };
  return visit(input, 0);
}

function decodeActor(input: unknown): ActorSnapshot {
  const actor = object(input);
  if (!isActorRef(actor.actor_ref)) throw new Error("Invalid actor reference in snapshot");
  const status = object(actor.status);
  const name = status.status;
  if (
    name !== "idle" &&
    name !== "running" &&
    name !== "awaiting" &&
    name !== "paused" &&
    name !== "stopped" &&
    name !== "failed"
  ) {
    throw new Error("Unknown actor status");
  }
  const execution = actor.execution === null ? null : object(actor.execution);
  return {
    actor_ref: actor.actor_ref,
    mailbox_id: uint(actor.mailbox_id),
    implementation_id: text(actor.implementation_id),
    incarnation: uint(actor.incarnation),
    state: decodeValue(actor.state),
    status: name === "failed" ? { status: name, message: text(status.message) } : { status: name },
    execution: execution && {
      vm_steps: uint(execution.vm_steps),
      instruction_pointer: uint(execution.instruction_pointer),
      call_depth: uint(execution.call_depth),
    },
    queued: array(actor.queued).map(decodeValue),
    reserved: actor.reserved === null ? null : decodeValue(actor.reserved),
  };
}

export type ReadReply =
  { kind: "actors"; data: ActorSnapshot[] } | { kind: "actor"; data: ActorSnapshot };

function decodeEnvelope(message: string): Record<string, unknown> {
  const envelope = object(
    parse(message, undefined, (literal) => {
      if (!/^-?(0|[1-9]\d*)$/.test(literal)) throw new Error("Non-integer protocol number");
      return BigInt(literal);
    }),
  );
  if (envelope.version !== BigInt(PROTOCOL_VERSION)) throw new Error("Protocol version mismatch");
  return envelope;
}

export function decodeNotification(message: string): bigint | undefined {
  const envelope = decodeEnvelope(message);
  if (!Object.hasOwn(envelope, "event")) return undefined;
  if (Object.keys(envelope).length !== 2) throw new Error("Invalid notification envelope");
  const event = object(envelope.event);
  if (event.event !== "image_changed" || Object.keys(event).length !== 2) {
    throw new Error("Unsupported server event");
  }
  const data = object(event.data);
  const revision = uint(data.revision);
  if (Object.keys(data).length !== 1 || revision === 0n) {
    throw new Error("Invalid image revision");
  }
  return revision;
}

export function decodeResponse(line: string, requestId: string, command: ReadCommand): ReadReply {
  const response = decodeEnvelope(line);
  if (response.request_id !== requestId) throw new Error("Response request ID mismatch");
  const result = object(response.result);
  if (Object.keys(result).length !== 1) throw new Error("Invalid response result");
  if (Object.hasOwn(result, "Err")) {
    const error = object(result.Err);
    const code = text(error.code);
    if (
      ![
        "invalid_request",
        "unsupported_version",
        "not_found",
        "conflict",
        "capacity_exceeded",
        "runtime_failure",
      ].includes(code)
    ) {
      throw new Error("Unknown protocol error code");
    }
    throw new ProtocolError(code, text(error.message));
  }
  const reply = object(result.Ok);
  if (command.operation === "list_actors" && reply.kind === "actors") {
    const data = array(reply.data).map(decodeActor);
    if (new Set(data.map((actor) => actor.actor_ref)).size !== data.length) {
      throw new Error("Duplicate actor references in snapshot");
    }
    return { kind: "actors", data };
  }
  if (command.operation === "inspect_actor" && reply.kind === "actor") {
    const data = decodeActor(reply.data);
    if (data.actor_ref !== command.parameters.actor_ref) throw new Error("Snapshot actor mismatch");
    return { kind: "actor", data };
  }
  throw new Error("Unexpected reply kind");
}

export function formatSnapshot(
  actor: ActorSnapshot,
  endpoint: string,
  capturedAt: string,
  connected: boolean,
): string {
  const result = stringify(
    { connection: { endpoint, connected }, captured_at: capturedAt, actor },
    null,
    2,
  );
  if (result === undefined) throw new Error("Could not format actor snapshot");
  return result + "\n";
}