runa/vscode/src/extension.ts

import * as vscode from "vscode";
import { ImageClient, parseEndpoint } from "./client";
import { ActorSnapshot, formatSnapshot } from "./protocol";

class ActorTree implements vscode.TreeDataProvider<ActorSnapshot> {
  private readonly changes = new vscode.EventEmitter<void>();
  readonly onDidChangeTreeData = this.changes.event;
  actors: ActorSnapshot[] = [];

  replace(actors: ActorSnapshot[]): void {
    this.actors = actors;
    this.changes.fire();
  }

  getChildren(): ActorSnapshot[] {
    return this.actors;
  }

  getTreeItem(actor: ActorSnapshot): vscode.TreeItem {
    const item = new vscode.TreeItem(`@${actor.actor_ref}`);
    item.id = actor.actor_ref;
    item.contextValue = "runa.actor";
    item.description = `${actor.status.status}, queued ${actor.queued.length}${actor.reserved ? ", reserved 1" : ""}`;
    item.tooltip = [
      `Actor: @${actor.actor_ref}`,
      `Status: ${actor.status.status}`,
      ...(actor.status.message ? [`Failure: ${actor.status.message}`] : []),
      `Incarnation: ${actor.incarnation}`,
      `Implementation: ${actor.implementation_id}`,
    ].join("\n");
    item.iconPath = new vscode.ThemeIcon(
      actor.status.status === "failed" ? "error" : "symbol-interface",
    );
    item.command = {
      command: "runa.inspect",
      title: "Inspect Actor",
      arguments: [actor],
    };
    return item;
  }

  dispose(): void {
    this.changes.dispose();
  }
}

interface DocumentSnapshot {
  session: number;
  endpoint: string;
  actor: ActorSnapshot;
  capturedAt: string;
  connected: boolean;
}

class SnapshotDocuments implements vscode.TextDocumentContentProvider {
  private readonly changes = new vscode.EventEmitter<vscode.Uri>();
  readonly onDidChange = this.changes.event;
  private readonly documents = new Map<string, DocumentSnapshot>();

  set(session: number, endpoint: string, actor: ActorSnapshot): vscode.Uri {
    const uri = vscode.Uri.from({
      scheme: "runa",
      authority: String(session),
      path: `/actors/${actor.actor_ref}.json`,
    });
    this.documents.set(uri.toString(), {
      session,
      endpoint,
      actor,
      capturedAt: new Date().toISOString(),
      connected: true,
    });
    this.changes.fire(uri);
    return uri;
  }

  refresh(session: number, endpoint: string, actors: ActorSnapshot[]): void {
    const byRef = new Map(actors.map((actor) => [actor.actor_ref, actor]));
    for (const snapshot of this.documents.values()) {
      const actor = byRef.get(snapshot.actor.actor_ref);
      if (snapshot.session === session && actor) this.set(session, endpoint, actor);
    }
  }

  disconnect(session: number): void {
    for (const [key, snapshot] of this.documents) {
      if (snapshot.session === session) {
        snapshot.connected = false;
        this.changes.fire(vscode.Uri.parse(key));
      }
    }
  }

  provideTextDocumentContent(uri: vscode.Uri): string {
    const snapshot = this.documents.get(uri.toString());
    if (!snapshot) throw new Error("Snapshot is no longer available; inspect the actor again");
    return formatSnapshot(
      snapshot.actor,
      snapshot.endpoint,
      snapshot.capturedAt,
      snapshot.connected,
    );
  }

  remove(uri: vscode.Uri): void {
    this.documents.delete(uri.toString());
  }

  dispose(): void {
    this.changes.dispose();
    this.documents.clear();
  }
}

export function activate(context: vscode.ExtensionContext): void {
  const tree = new ActorTree();
  const documents = new SnapshotDocuments();
  const view = vscode.window.createTreeView("runa.actors", { treeDataProvider: tree });
  const output = vscode.window.createOutputChannel("Runa");
  const status = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left);
  const reported = new WeakSet<Error>();
  let client: ImageClient | undefined;
  let session = 0;
  let endpoint = "";
  let disposed = false;

  const report = (error: unknown) => {
    const failure = error instanceof Error ? error : new Error(String(error));
    if (disposed || reported.has(failure)) return;
    reported.add(failure);
    output.appendLine(`Error: ${failure.message}`);
    void vscode.window.showErrorMessage(`Runa: ${failure.message}`);
  };

  const connectionStatus = (connected: boolean) => {
    void vscode.commands.executeCommand("setContext", "runa.connected", connected);
    status.text = connected ? "$(plug) Runa" : "$(debug-disconnect) Runa";
    status.tooltip = connected
      ? `Connected to ${endpoint}. Click to refresh actors.`
      : "Connect to Runa";
    status.command = connected ? "runa.refresh" : "runa.connect";
    view.description = connected ? endpoint : "Disconnected";
    view.message = connected ? undefined : "Not connected to an image.";
  };

  const disconnect = () => {
    const previous = client;
    client = undefined;
    documents.disconnect(session);
    session++;
    previous?.close();
    tree.replace([]);
    connectionStatus(false);
  };

  const refresh = async () => {
    const current = client;
    const currentSession = session;
    if (!current) throw new Error("Connect to a local image first");
    const actors = await current.listActors();
    if (current !== client || currentSession !== session) return;
    tree.replace(actors);
    documents.refresh(session, endpoint, actors);
    view.message = actors.length
      ? undefined
      : "No actors. Spawn one with the Runa REPL, then refresh.";
  };

  const connect = async (target?: string) => {
    const address =
      target ??
      (await vscode.window.showInputBox({
        title: "Connect to a trusted local Runa image (no authentication)",
        value: vscode.workspace.getConfiguration("runa").get<string>("address", "127.0.0.1:3737"),
        prompt:
          "Start the server with mise run serve-network. The extension does not start an image.",
        validateInput: (value) => {
          try {
            parseEndpoint(value);
            return undefined;
          } catch (error) {
            return error instanceof Error ? error.message : String(error);
          }
        },
      }));
    if (address === undefined || disposed) return;
    parseEndpoint(address);
    disconnect();
    endpoint = address;
    const currentSession = session;
    view.message = `Connecting to ${address}...`;
    let connected: ImageClient;
    try {
      connected = await ImageClient.connect(address, (error) => {
        if (currentSession !== session || disposed) return;
        client = undefined;
        tree.replace([]);
        documents.disconnect(currentSession);
        connectionStatus(false);
        output.appendLine(`Disconnected from ${address}`);
        if (error) report(error);
      });
    } catch (error) {
      if (currentSession === session) connectionStatus(false);
      throw error;
    }
    if (currentSession !== session || disposed) {
      connected.close();
      return;
    }
    client = connected;
    try {
      await refresh();
    } catch (error) {
      if (currentSession === session) disconnect();
      throw error;
    }
    if (connected !== client) return;
    connectionStatus(true);
    view.message = tree.actors.length
      ? undefined
      : "No actors. Spawn one with the Runa REPL, then refresh.";
    output.appendLine(`Connected to ${address} (protocol v3)`);
  };

  const inspect = async (selected?: ActorSnapshot) => {
    const current = client;
    const currentSession = session;
    if (!current) throw new Error("Connect to a local image first");
    let actorRef = selected?.actor_ref;
    if (!actorRef) {
      const choice = await vscode.window.showQuickPick(
        tree.actors.map((actor) => ({
          label: `@${actor.actor_ref}`,
          description: actor.status.status,
          actorRef: actor.actor_ref,
        })),
        { placeHolder: "Choose an actor to inspect (refresh to discover newly spawned actors)" },
      );
      if (!choice) return;
      actorRef = choice.actorRef;
    }
    if (current !== client || currentSession !== session) return;
    const actor = await current.inspectActor(actorRef);
    if (current !== client || currentSession !== session) return;
    const uri = documents.set(session, endpoint, actor);
    const document = await vscode.workspace.openTextDocument(uri);
    await vscode.languages.setTextDocumentLanguage(document, "json");
    await vscode.window.showTextDocument(document, { preview: true });
  };

  const command = <Args extends unknown[]>(
    name: string,
    action: (...args: Args) => Promise<void> | void,
  ) =>
    vscode.commands.registerCommand(name, async (...args: Args) => {
      try {
        await action(...args);
      } catch (error) {
        report(error);
      }
    });

  context.subscriptions.push(
    tree,
    documents,
    view,
    output,
    status,
    vscode.workspace.registerTextDocumentContentProvider("runa", documents),
    vscode.workspace.onDidCloseTextDocument((document) => documents.remove(document.uri)),
    command("runa.connect", connect),
    command("runa.disconnect", disconnect),
    command("runa.refresh", refresh),
    command("runa.inspect", inspect),
    {
      dispose: () => {
        disposed = true;
        client?.close();
      },
    },
  );
  connectionStatus(false);
  status.show();
}