VS Codeのエージェントモードを拡張するには、を試してください!

WebAssembly を使用した拡張機能開発

2024年5月8日 by Dirk Bäumer

Visual Studio Code は WebAssembly Execution Engine 拡張機能を通じて WASM バイナリの実行をサポートしています。主なユースケースは、C/C++ や Rust で書かれたプログラムを WebAssembly にコンパイルし、それらのプログラムを VS Code で直接実行することです。注目すべき例として Visual Studio Code for Education があり、このサポートを利用して Python インタープリターを VS Code for the Web で実行しています。この ブログ記事 では、これがどのように実装されているかについて詳しく解説しています。

2024年1月、Bytecode Alliance は WASI 0.2 preview を発表しました。WASI 0.2 preview の主要技術は Component Model です。WebAssembly Component Model は、インターフェース、データ型、モジュール構成を標準化することで、WebAssembly コンポーネントとそのホスト環境との間の相互作用を効率化します。この標準化は、WIT (WASM Interface Type) ファイルを使用することで促進されます。WIT ファイルは、JavaScript/TypeScript 拡張機能 (ホスト) と、Rust や C/C++ などの別の言語でコーディングされた計算を実行する WebAssembly コンポーネントとの間の相互作用を記述するのに役立ちます。

このブログ記事では、開発者が Component Model を利用して WebAssembly ライブラリを自身の拡張機能に統合する方法を概説します。私たちは3つのユースケースに焦点を当てます。(a) WebAssembly を使用してライブラリを実装し、JavaScript/TypeScript の拡張機能コードからそれを呼び出す、(b) WebAssembly コードから VS Code API を呼び出す、(c) リソースを使用して WebAssembly または TypeScript コードのいずれかでステートフルなオブジェクトをカプセル化し管理する方法を示す。

これらの例では、VS Code と NodeJS に加えて、以下のツールの最新バージョンがインストールされている必要があります: Rust コンパイラツールチェーンwasm-tools、そして wit-bindgen

また、この記事に対する貴重なフィードバックをいただいた Fastly の L. Pereira 氏と Luke Wagner 氏にも感謝の意を表します。

Rust での電卓

最初の例では、開発者が Rust で書かれたライブラリを VS Code 拡張機能に統合する方法を示します。前述の通り、コンポーネントは WIT ファイルを使用して記述されます。私たちの例では、ライブラリは加算、減算、乗算、除算などの単純な操作を実行します。対応する WIT ファイルを以下に示します。

package vscode:example;

interface types {
	record operands {
		left: u32,
		right: u32
	}

	variant operation {
		add(operands),
		sub(operands),
		mul(operands),
		div(operands)
	}
}
world calculator {
	use types.{ operation };

	export calc: func(o: operation) -> u32;
}

Rust ツール wit-bindgen は、電卓の Rust バインディングを生成するために利用されます。このツールを使用するには2つの方法があります。

  • 実装ファイル内で直接バインディングを生成する手続き型マクロとして使用する方法。この方法は標準的ですが、生成されたバインディングコードを検査できないという欠点があります。

  • コマンドラインツールとして使用し、ディスク上にバインディングファイルを作成する方法。このアプローチは、後述するリソースの例のために VS Code 拡張機能サンプルリポジトリにあるコードで実証されています。

wit-bindgen ツールを手続き型マクロとして使用した、対応する Rust ファイルは次のようになります。

// Use a procedural macro to generate bindings for the world we specified in
// `calculator.wit`
wit_bindgen::generate!({
	// the name of the world in the `*.wit` input file
	world: "calculator",
});

しかし、cargo build --target wasm32-unknown-unknown コマンドを使用して Rust ファイルを WebAssembly にコンパイルすると、エクスポートされた calc 関数の実装が欠落しているため、コンパイルエラーが発生します。以下に calc 関数の簡単な実装を示します。

// Use a procedural macro to generate bindings for the world we specified in
// `calculator.wit`
wit_bindgen::generate!({
	// the name of the world in the `*.wit` input file
	world: "calculator",
});

struct Calculator;

impl Guest for Calculator {

    fn calc(op: Operation) -> u32 {
		match op {
			Operation::Add(operands) => operands.left + operands.right,
			Operation::Sub(operands) => operands.left - operands.right,
			Operation::Mul(operands) => operands.left * operands.right,
			Operation::Div(operands) => operands.left / operands.right,
		}
	}
}

// Export the Calculator to the extension code.
export!(Calculator);

ファイルの最後にある export!(Calculator); ステートメントは、拡張機能が API を呼び出せるように、WebAssembly コードから Calculator をエクスポートします。

wit2ts ツールは、VS Code 拡張機能内で WebAssembly コードと対話するために必要な TypeScript バインディングを生成するために使用されます。このツールは、主に以下の VS Code 拡張機能アーキテクチャの特定の要件を満たすために VS Code チームによって開発されました。

  • VS Code API は拡張機能ホストワーカー内でのみアクセス可能です。拡張機能ホストワーカーから生成された追加のワーカーは、VS Code API にアクセスできません。これは、各ワーカーが通常、ほぼすべてのランタイム API にアクセスできる NodeJS やブラウザーのような環境とは対照的です。
  • 複数の拡張機能が同じ拡張機能ホストワーカーを共有します。拡張機能は、そのワーカー上で長時間実行される同期的な計算を実行することを避けるべきです。

これらのアーキテクチャ要件は、VS Code 用の WASI Preview 1 を実装したときには既に存在していました。しかし、私たちの初期の実装は手動で書かれていました。Component Model のより広範な採用を予測し、コンポーネントと VS Code 固有のホスト実装との統合を容易にするツールを開発しました。

コマンド wit2ts --outDir ./src ./wit は、src フォルダに calculator.ts ファイルを生成します。これには WebAssembly コード用の TypeScript バインディングが含まれています。これらのバインディングを利用する簡単な拡張機能は次のようになります。

import * as vscode from 'vscode';
import { WasmContext, Memory } from '@vscode/wasm-component-model';

// Import the code generated by wit2ts
import { calculator, Types } from './calculator';

export async function activate(context: vscode.ExtensionContext): Promise<void> {
  // The channel for printing the result.
  const channel = vscode.window.createOutputChannel('Calculator');
  context.subscriptions.push(channel);

  // Load the Wasm module
  const filename = vscode.Uri.joinPath(
    context.extensionUri,
    'target',
    'wasm32-unknown-unknown',
    'debug',
    'calculator.wasm'
  );
  const bits = await vscode.workspace.fs.readFile(filename);
  const module = await WebAssembly.compile(bits);

  // The context for the WASM module
  const wasmContext: WasmContext.Default = new WasmContext.Default();

  // Instantiate the module
  const instance = await WebAssembly.instantiate(module, {});
  // Bind the WASM memory to the context
  wasmContext.initialize(new Memory.Default(instance.exports));

  // Bind the TypeScript Api
  const api = calculator._.exports.bind(
    instance.exports as calculator._.Exports,
    wasmContext
  );

  context.subscriptions.push(
    vscode.commands.registerCommand('vscode-samples.wasm-component-model.run', () => {
      channel.show();
      channel.appendLine('Running calculator example');
      const add = Types.Operation.Add({ left: 1, right: 2 });
      channel.appendLine(`Add ${api.calc(add)}`);
      const sub = Types.Operation.Sub({ left: 10, right: 8 });
      channel.appendLine(`Sub ${api.calc(sub)}`);
      const mul = Types.Operation.Mul({ left: 3, right: 7 });
      channel.appendLine(`Mul ${api.calc(mul)}`);
      const div = Types.Operation.Div({ left: 10, right: 2 });
      channel.appendLine(`Div ${api.calc(div)}`);
    })
  );
}

上記のコードをコンパイルして VS Code for the Web で実行すると、Calculator チャンネルに以下の出力が生成されます。

この例の完全なソースコードは VS Code 拡張機能サンプルリポジトリにあります。

@vscode/wasm-component-model の内部

wit2ts ツールによって生成されたソースコードを調べると、それが @vscode/wasm-component-model npm モジュールに依存していることがわかります。このモジュールは、Component Model の Canonical ABI の VS Code 実装として機能し、対応する Python コードから着想を得ています。このブログ記事を理解するために Component Model の内部を理解する必要はありませんが、特に JavaScript/TypeScript と WebAssembly コード間でデータがどのように渡されるかについて、その仕組みを少し明らかにします。

wit-bindgenjco のように WIT ファイルのバインディングを生成する他のツールとは異なり、wit2ts はメタモデルを作成します。これは、さまざまなユースケースのために実行時にバインディングを生成するために使用できます。この柔軟性により、VS Code 内での拡張機能開発のアーキテクチャ要件を満たすことができます。このアプローチを使用することで、バインディングを「promisify」し、WebAssembly コードをワーカーで実行できるようにします。私たちはこのメカニズムを利用して、VS Code 用の WASI 0.2 preview を実装しています。

バインディングを生成する際に、関数が calculator._.imports.create (アンダースコアに注意) のような名前で参照されていることにお気づきかもしれません。WIT ファイル内のシンボルとの名前の衝突を避けるため (例えば、imports という名前の型定義が存在する可能性があるため)、API 関数は _ 名前空間に配置されます。メタモデル自体は $ 名前空間に存在します。したがって、calculator.$.exports.calc はエクスポートされた calc 関数のメタデータを表します。

上記の例では、calc 関数に渡される add 操作パラメーターは、操作コード、左辺値、右辺値の3つのフィールドで構成されています。Component Model の Canonical ABI によると、引数は値渡しされます。また、データがどのようにシリアライズされ、WebAssembly 関数に渡され、反対側でデシリアライズされるかも規定しています。このプロセスにより、2つの操作オブジェクトが生成されます。1つは JavaScript ヒープ上に、もう1つは線形 WebAssembly メモリ内にあります。次の図はこれを示しています。

Diagram illustrating how parameters are passed.

以下の表は、利用可能な WIT 型、VS Code Component Model 実装における JavaScript オブジェクトへのマッピング、および使用される対応する TypeScript 型を一覧にしたものです。

WIT JavaScript TypeScript
u8 number type u8 = number;
u16 number type u16 = number;
u32 number type u32 = number;
u64 bigint type u64 = bigint;
s8 number type s8 = number;
s16 number type s16 = number;
s32 number type s32 = number;
s64 bigint type s64 = bigint;
float32 number type float32 = number;
float64 number type float64 = number;
bool boolean boolean
string string string
char string[0] string
record オブジェクトリテラル 型宣言
list<T> [] Array<T>
tuple<T1, T2> [] [T1, T2]
enum 文字列値 文字列 enum
flags number bigint
variant オブジェクトリテラル 判別共用体
option<T> 変数 ? and (T | undefined)
result<ok, err> 例外またはオブジェクトリテラル 例外または result 型

Component Model は低レベル (C スタイル) のポインターをサポートしていないことに注意することが重要です。そのため、オブジェクトグラフや再帰的なデータ構造を渡すことはできません。この点で、JSON と同じ制限を共有しています。データのコピーを最小限に抑えるため、Component Model はリソースの概念を導入しており、これについてはこのブログ記事の今後のセクションで詳しく説明します。

jco プロジェクトも、type コマンドを使用して WebAssembly コンポーネント用の JavaScript/TypeScript バインディングの生成をサポートしています。前述の通り、私たちは VS Code の特定のニーズを満たすために独自のツールを開発しました。しかし、可能な限りツール間の整合性を確保するために、jco チームと隔週で会議を開いています。基本的な要件は、両方のツールが WIT データ型に対して同じ JavaScript および TypeScript 表現を使用することです。また、2つのツール間でコードを共有する可能性も探っています。

WebAssembly コードから TypeScript を呼び出す

WIT ファイルはホスト (VS Code 拡張機能) と WebAssembly コードとの間の相互作用を記述し、双方向の通信を容易にします。私たちの例では、この機能により WebAssembly コードがその活動のトレースをログに記録できるようになります。これを可能にするために、WIT ファイルを次のように変更します。

world calculator {

	/// ....

	/// A log function implemented on the host side.
	import log: func(msg: string);

	/// ...
}

Rust 側では、log 関数を呼び出すことができるようになります。

fn calc(op: Operation) -> u32 {
	log(&format!("Starting calculation: {:?}", op));
	let result = match op {
		// ...
	};
	log(&format!("Finished calculation: {:?}", op));
	result
}

TypeScript 側では、拡張機能開発者が必要とする唯一のアクションは、log 関数の実装を提供することです。その後、VS Code Component Model は必要なバインディングの生成を容易にし、それらをインポートとして WebAssembly インスタンスに渡します。

export async function activate(context: vscode.ExtensionContext): Promise<void> {
  // ...

  // The channel for printing the log.
  const log = vscode.window.createOutputChannel('Calculator - Log', { log: true });
  context.subscriptions.push(log);

  // The implementation of the log function that is called from WASM
  const service: calculator.Imports = {
    log: (msg: string) => {
      log.info(msg);
    }
  };

  // Create the bindings to import the log function into the WASM module
  const imports = calculator._.imports.create(service, wasmContext);
  // Instantiate the module
  const instance = await WebAssembly.instantiate(module, imports);

  // ...
}

最初の例と比較して、WebAssembly.instantiate 呼び出しには、2番目の引数として calculator._.imports.create(service, wasmContext) の結果が含まれるようになりました。この imports.create 呼び出しは、サービス実装から低レベルの WASM バインディングを生成します。最初の例では、インポートが不要だったため、空のオブジェクトリテラルを渡しました。今回は、VS Code デスクトップ環境のデバッガー下で拡張機能を実行します。Connor Peet の素晴らしい仕事のおかげで、Rust コードにブレークポイントを設定し、VS Code デバッガーを使用してステップ実行することが可能になりました。

Component Model リソースの使用

WebAssembly Component Model は、状態をカプセル化し管理するための標準化されたメカニズムであるリソースの概念を導入しています。この状態は、呼び出し境界の一方 (例えば、TypeScript コード内) で管理され、もう一方 (例えば、WebAssembly コード内) でアクセスおよび操作されます。リソースは WASI preview 0.2 API で広範囲に使用されており、ファイルディスクリプタが典型的な例です。この設定では、状態は拡張機能ホストによって管理され、WebAssembly コードによってアクセスおよび操作されます。

リソースは逆方向にも機能し、その状態は WebAssembly コードによって管理され、拡張機能コードによってアクセスおよび操作されます。このアプローチは、VS Code が WebAssembly でステートフルなサービスを実装し、それらを TypeScript 側からアクセスする場合に特に有益です。ヒューレット・パッカードの携帯電卓で使用されているような逆ポーランド記法をサポートする電卓を実装するリソースを以下の例で定義します。

// wit/calculator.wit
package vscode:example;

interface types {

	enum operation {
		add,
		sub,
		mul,
		div
	}

	resource engine {
		constructor();
		push-operand: func(operand: u32);
		push-operation: func(operation: operation);
		execute: func() -> u32;
	}
}
world calculator {
	export types;
}

以下は、Rust での電卓リソースの簡単な実装です。

impl EngineImpl {
	fn new() -> Self {
		EngineImpl {
			left: None,
			right: None,
		}
	}

	fn push_operand(&mut self, operand: u32) {
		if self.left == None {
			self.left = Some(operand);
		} else {
			self.right = Some(operand);
		}
	}

	fn push_operation(&mut self, operation: Operation) {
        let left = self.left.unwrap();
        let right = self.right.unwrap();
        self.left = Some(match operation {
			Operation::Add => left + right,
			Operation::Sub => left - right,
			Operation::Mul => left * right,
			Operation::Div => left / right,
		});
	}

	fn execute(&mut self) -> u32 {
		self.left.unwrap()
	}
}

TypeScript コードでは、以前と同じ方法でエクスポートをバインドします。唯一の違いは、バインディングプロセスが、WebAssembly コード内で calculator リソースをインスタンス化および管理するために使用されるプロキシクラスを提供するようになったことです。

// Bind the JavaScript Api
const api = calculator._.exports.bind(
  instance.exports as calculator._.Exports,
  wasmContext
);

context.subscriptions.push(
  vscode.commands.registerCommand('vscode-samples.wasm-component-model.run', () => {
    channel.show();
    channel.appendLine('Running calculator example');

    // Create a new calculator engine
    const calculator = new api.types.Engine();

    // Push some operands and operations
    calculator.pushOperand(10);
    calculator.pushOperand(20);
    calculator.pushOperation(Types.Operation.add);
    calculator.pushOperand(2);
    calculator.pushOperation(Types.Operation.mul);

    // Calculate the result
    const result = calculator.execute();
    channel.appendLine(`Result: ${result}`);
  })
);

対応するコマンドを実行すると、出力チャンネルに Result: 60 と表示されます。前述のように、リソースの状態は呼び出し境界の一方の側に存在し、もう一方の側からハンドルを使用してアクセスされます。リソースと相互作用するメソッドに渡される引数を除き、データのコピーは発生しません。

Diagram illustrating how resources are accessed.

この例の完全なソースコードは VS Code 拡張機能サンプルリポジトリで入手できます。

Rust から直接 VS Code API を使用する

Component Model リソースは、WebAssembly コンポーネントとホストの間で状態をカプセル化し管理するために使用できます。この機能により、リソースを利用して VS Code API を正規の方法で WebAssembly コードに公開することができます。このアプローチの利点は、拡張機能全体を WebAssembly にコンパイルされる言語で作成できる点にあります。私たちはこのアプローチの探求を開始しており、以下は Rust で書かれた拡張機能のソースコードです。

use std::rc::Rc;

#[export_name = "activate"]
pub fn activate() -> vscode::Disposables {
	let mut disposables: vscode::Disposables = vscode::Disposables::new();

	// Create an output channel.
	let channel: Rc<vscode::OutputChannel> = Rc::new(vscode::window::create_output_channel("Rust Extension", Some("plaintext")));

	// Register a command handler
	let channel_clone = channel.clone();
	disposables.push(vscode::commands::register_command("testbed-component-model-vscode.run", move || {
		channel_clone.append_line("Open documents");

		// Print the URI of all open documents
		for document in vscode::workspace::text_documents() {
			channel.append_line(&format!("Document: {}", document.uri()));
		}
	}));
	return disposables;
}

#[export_name = "deactivate"]
pub fn deactivate() {
}

このコードが TypeScript で書かれた拡張機能に似ていることに注意してください。

この探求は有望に見えましたが、今のところこれを進めないことにしました。主な理由は、WASM での非同期サポートの欠如です。多くの VS Code API は非同期であり、それらを直接 WebAssembly コードにプロキシするのは困難です。WebAssembly コードを別のワーカーで実行し、WASI Preview 1 サポートで使用されているのと同じ同期メカニズムを WebAssembly ワーカーと拡張機能ホストワーカーの間で採用することもできます。しかし、このアプローチは、同期 API 呼び出し中に予期しない動作を引き起こす可能性があります。なぜなら、これらの呼び出しは実際には非同期に実行されるためです。その結果、2つの同期呼び出しの間で観測可能な状態が変わる可能性があります (例えば、setX(5); getX(); が 5 を返さないかもしれません)。

さらに、0.3 プレビューのタイムフレームで WASI に完全な非同期サポートを導入する取り組みが進行中です。Luke Wagner は WASM I/O 2024 で非同期サポートの現状についてアップデートを提供しました。私たちはこのサポートを待つことにしました。なぜなら、これにより、より完全でクリーンな実装が可能になるからです。

対応する WIT ファイル、Rust コード、TypeScript コードに興味がある場合は、vscode-wasm リポジトリの rust-api フォルダにあります。

今後の展望

現在、WebAssembly コードを拡張機能開発に利用できるさらなる分野をカバーする、続編のブログ記事を準備中です。主なトピックは以下の通りです。

  • WebAssembly で 言語サーバーを作成する。
  • 生成されたメタモデルを使用して、長時間実行される WebAssembly コードを透過的に別のワーカーにオフロードする。

VS Code に特化した Component Model の実装が整ったことで、私たちは VS Code 用の WASI 0.2 preview の実装に向けた取り組みを続けています。

よろしくお願いいたします。

DirkとVS Codeチーム

ハッピーコーディング!