VS Code ExtensionのWebviewをComlinkでEnjoyableにする

VS Code ExtensionのWebviewのめんどくさいところ

VS Codeでは、Webviewを使うことで自由度の高い、リッチな体験を実現できます。
しかし、ExtensionとWebviewはコンテキストが分離されているため、postMessageを使ってやり取りをする必要があります。

// Send a message to our webview.
// You can send any JSON serializable data.
currentPanel.webview.postMessage({ command: 'refactor' });

これ、1つや2つならいいのですが、量が多くなってくるとこれがしんどいんですよね...

Comlinkを使ってEnjoyableにする

Comlinkとは

そこで、今回はComlinkというライブラリを使ってWebview開発をEnjoyableにします。

ComlinkはWeb Workerとメインスレッドの通信を簡単にしてくれるライブラリです。使用例を見るとわかりやすいです。

/*
　Worker内
*/
const obj = {
  counter: 0,
  inc() {
    this.counter++;
  },
};

// Worker内からオブジェクトを公開すると...
Comlink.expose(obj);

/*
　メインスレッド
*/
const worker = new Worker(worker.js);
const obj = Comlink.wrap(worker);

// Worker内のcounter変数にアクセスできる!
// アクセスは非同期
alert(`Counter: ${await obj.counter}`);

// Worker内の関数も呼び出せる!
// 呼び出しは非同期
await obj.inc();

objは実際にはProxyになっていて、プロパティアクセスや関数呼び出し時にトラップ内でいい感じにWorkerとpostMessageでやり取りしてくれます。この仕組みによって、開発者はpostMessageやaddEventListnerを書かずにメインスレッドとWorkerのコミュニケーションを実装できます。

ここで、

const obj = Comlink.wrap(worker);

に注目なのですが、実はこのwrapの引数は必ずしもWorkerである必要はなく、Endpointインターフェースを満たしていればokという抽象化がされています。

export interface Endpoint extends EventSource {
  postMessage(message: any, transfer?: Transferable[]): void;
  start?: () => void;
}

ということは、VS CodeのExtensionとWebview間で通信できるようにこのEndpointを実装すれば、Comlinkの機能を使ってEnjoyできるというわけです！

Endpointの実装(Webview側)

さっそく、Webview側から実装していきます。こっちは簡単です。

• postMessageは、acquireVsCodeApiで取得できるapiをそのまま使えます
• addEventListener, removeEventListenerについては、普通にwindowに生えているものを使います

import { type Endpoint } from "comlink";

type VSCodeApi = ReturnType<typeof acquireVsCodeApi>;

class WebviewEndpoint implements Endpoint {
  constructor(private readonly vscode: VSCodeApi) {}

  postMessage(msg: any, transferables: Transferable[]): void {
    if (transferables) {
      throw new Error("transferables are not supported");
    }
    this.vscode.postMessage(msg);
  }

  addEventListener(
    type: string,
    listener: EventListenerOrEventListenerObject,
    options?: {}
  ): void {
    if (type !== "message") {
      return;
    }

    window.addEventListener("message", listener, options);
  }

  removeEventListener(
    type: string,
    listener: EventListenerOrEventListenerObject,
    options?: {}
  ): void {
    if (type !== "message") {
      return;
    }

    window.removeEventListener("message", listener, options);
  }
}

export const createWebviewEndpoint = (
  vscode: VSCodeApi
) => new WebviewEndpoint(vscode);

transferablesについて

Transferableはコンテキスト間でオブジェクトを転送する際に使うものですが、vscodeのapiは対応していないので、無視しています。

Endpointの実装(Extension側)

Extension側は一工夫必要です。

• postMessageは、webview.postMessage()を通じてwebview側に送信します
• イベントリスナの登録は、this.webview.onDidReceiveMessageで行います。あとでクリーンアップできるように、MapにハンドラとDisposableを保管しておきます

import { type Disposable, type Webview } from "vscode";
import { type Endpoint } from "comlink";

class ExtensionEndpoint implements Endpoint {
  private readonly listenerDisposables: Map<unknown, Disposable> = new Map();

  constructor(private readonly webview: Webview) {}

  postMessage(msg: any, transferables: Transferable[]): void {
    if (transferables) {
      throw new Error("transferables are not supported");
    }
    this.webview.postMessage(msg);
  }

  addEventListener(
    type: string,
    listener: EventListenerOrEventListenerObject,
    options?: {}
  ): void {
    if (type !== "message") {
      return;
    }

    if (typeof listener === "function") {
      const disposable = this.webview.onDidReceiveMessage((msg: unknown) =>
        listener(new MessageEvent("message", { data: msg }))
      );
      this.listenerDisposables.set(listener, disposable);
    } else {
      const disposable = this.webview.onDidReceiveMessage((msg) =>
        listener.handleEvent(new MessageEvent("message", { data: msg }))
      );
      this.listenerDisposables.set(listener, disposable);
    }
  }

  removeEventListener(
    type: string,
    listener: EventListenerOrEventListenerObject,
    options?: {}
  ): void {
    if (type !== "message") {
      return;
    }

    const disposable = this.listenerDisposables.get(listener);
    if (disposable) {
      disposable.dispose();
      this.listenerDisposables.delete(listener);
    }
  }
}

export const createExtensionEndpoint = (
  webview: Webview
) => new ExtensionEndpoint(webview);

使い方

これで、ExtensionとWebviewの両方のEndpointが実装できたので、実際に使ってみます。

Extension側でオブジェクトをexposeすると、

Comlink.expose(
  {
    counter: 0,
    inc() {
      this.counter++;
    },
  },
  createExtensionEndpoint(webview)
);

Extension側で使えるようになります！

const vscode = acquireVsCodeApi();

const obj = Comlink.wrap(
  createWebviewEndpoint(vscode)
);

await obj.counter // 0
await obj.inc()
await obj.counter // 1

あまりないとは思いますが、逆(Webview側でexposeしたものをExtension側で使う)もできます。

おわりに

ほとんどComlinkのおかげですが、無事ComlinkでVS CodeのWebviewをEnjoylableできました🎅