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

WebAssemblyを拡張機能開発に利用する - パート2

2024年6月7日 Dirk Bäumer

前回のWebAssemblyを拡張機能開発に利用するに関するブログ記事では、コンポーネントモデルを使用してWebAssemblyコードをVisual Studio Code拡張機能に統合する方法を説明しました。今回のブログ記事では、さらに2つの独立したユースケースに焦点を当てます。(a) 拡張ホストのメインスレッドをブロックしないように、WebAssemblyコードをワーカーで実行する。(b) WebAssemblyにコンパイルされる言語を使用して言語サーバーを作成する。

このブログ記事の例を実行するには、以下のツールが必要です: VS Code、Node.js、Rustコンパイラツールチェーン、wasm-tools、そしてwit-bindgen。

ワーカーでWebAssemblyコードを実行する

前回のブログ記事の例では、WebAssemblyコードはVS Code拡張ホストのメインスレッドで実行されていました。実行時間が短い限りは問題ありません。しかし、長時間の操作は、拡張ホストのメインスレッドが他の拡張機能で利用できるように、ワーカーで実行すべきです。

VS Codeコンポーネントモデルは、ワーカー側と拡張機能のメイン側の両方で必要なグルーコードを自動生成できるようにすることで、これを容易にするメタモデルを提供します。

以下のコードスニペットは、ワーカーに必要なコードを示しています。この例では、コードがworker.tsという名前のファイルに保存されていることを前提としています。

import { Connection, RAL } from '@vscode/wasm-component-model';
import { calculator } from './calculator';

async function main(): Promise<void> {
  const connection = await Connection.createWorker(calculator._);
  connection.listen();
}

main().catch(RAL().console.error);

このコードは、拡張ホストのメインワーカーと通信するための接続を作成し、wit2tsツールによって生成されたcalculatorワールドで接続を初期化します。

拡張機能側では、WebAssemblyモジュールをロードし、それをcalculatorワールドにもバインドします。ワーカーでの実行は非同期で行われるため（例: await api.calc(...)）、計算を実行するための対応する呼び出しは待機する必要があります。

// The channel for printing the result.
const channel = vscode.window.createOutputChannel('Calculator');
context.subscriptions.push(channel);

// 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.Promisified = {
  log: async (msg: string): Promise<void> => {
    // Wait 100ms to slow things down :-)
    await new Promise(resolve => setTimeout(resolve, 100));
    log.info(msg);
  }
};

// Load the WASM model
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);

// Create the worker
const worker = new Worker(
  vscode.Uri.joinPath(context.extensionUri, './out/worker.js').fsPath
);
// Bind the world to the worker
const api = await calculator._.bind(service, module, worker);

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

いくつか重要な点があります。

  • この例で使用されているWITファイルは、前回のブログ記事の計算機例で使用されたものと違いはありません。
  • WebAssemblyコードの実行はワーカーで行われるため、インポートされたサービス（例: 上記のlog関数）の実装はPromiseを返すことができますが、必須ではありません。
  • WebAssemblyは現在、同期実行モデルのみをサポートしています。そのため、WebAssemblyコードを実行するワーカーから拡張ホストのメインスレッドへの、インポートされたサービスを呼び出すためのすべての呼び出しには、以下のステップが必要です。
    1. 呼び出されるサービス（例: log関数を呼び出す）を記述したメッセージを拡張ホストのメインスレッドに投稿します。
    2. Atomics.waitを使用してワーカーの実行を中断します。
    3. 拡張ホストのメインスレッドでメッセージを処理します。
    4. ワーカーを再開し、Atomics.notifyを使用して結果を通知します。

この同期は、測定可能な時間オーバーヘッドを追加します。これらのステップはすべてコンポーネントモデルによって透過的に処理されますが、開発者はこれらを認識し、インポートされたAPIサーフェスを設計する際に考慮すべきです。

この例の完全なソースコードは、VS Code拡張機能サンプルリポジトリで見つけることができます。

WebAssemblyベースの言語サーバー

Web向けのVS CodeにおけるWebAssemblyサポートに取り組み始めたとき、私たちが思い描いたユースケースの1つは、WebAssemblyを使用して言語サーバーを実行することでした。VS CodeのLSPライブラリへの最新の変更と、WebAssemblyとLSPを橋渡しする新しいモジュールの導入により、WebAssembly言語サーバーの実装は、オペレーティングシステムプロセスとして実装するのと同じくらい簡単になりました。

さらに、WebAssembly言語サーバーは、WASI Preview 1を完全にサポートするWebAssembly Core拡張機能上で動作します。これは、たとえファイルがGitHubリポジトリなどのリモートに保存されている場合でも、言語サーバーが自身のプログラミング言語の通常のファイルシステムAPIを使用してワークスペース内のファイルにアクセスできることを意味します。

以下のコードスニペットは、lsp_serverクレートの例のサーバーに基づくRust言語サーバーを示しています。この言語サーバーは言語解析を行わず、GotoDefinitionリクエストに対して定義済みの結果を返すだけです。

match cast::<GotoDefinition>(req) {
    Ok((id, params)) => {
        let uri = params.text_document_position_params.text_document.uri;
        eprintln!("Received gotoDefinition request #{} {}", id, uri.to_string());
        let loc = Location::new(
            uri,
            lsp_types::Range::new(lsp_types::Position::new(0, 0), lsp_types::Position::new(0, 0))
        );
        let mut vec = Vec::new();
        vec.push(loc);
        let result = Some(GotoDefinitionResponse::Array(vec));
        let result = serde_json::to_value(&result).unwrap();
        let resp = Response { id, result: Some(result), error: None };
        connection.sender.send(Message::Response(resp))?;
        continue;
    }
    Err(err @ ExtractError::JsonError { .. }) => panic!("{err:?}"),
    Err(ExtractError::MethodMismatch(req)) => req,
};

言語サーバーの完全なソースコードは、VS Codeの例のリポジトリで見つけることができます。

新しい@vscode/wasm-wasi-lsp npmモジュールを使用して、拡張機能のTypeScriptコード内にWebAssembly言語サーバーを作成できます。WebAssemblyコードをWASIサポート付きのワーカーとしてインスタンス化するには、Web向けのVS CodeでWebAssemblyを実行するというブログ記事で詳しく説明されているWebAssembly Core拡張機能を使用します。

拡張機能のTypeScriptコードも単純です。プレーンテキストファイル用にサーバーを登録します。

import {
  createStdioOptions,
  createUriConverters,
  startServer
} from '@vscode/wasm-wasi-lsp';

export async function activate(context: ExtensionContext) {
  const wasm: Wasm = await Wasm.load();

  const channel = window.createOutputChannel('LSP WASM Server');
  // The server options to run the WebAssembly language server.
  const serverOptions: ServerOptions = async () => {
    const options: ProcessOptions = {
      stdio: createStdioOptions(),
      mountPoints: [{ kind: 'workspaceFolder' }]
    };

    // Load the WebAssembly code
    const filename = Uri.joinPath(
      context.extensionUri,
      'server',
      'target',
      'wasm32-wasip1-threads',
      'release',
      'server.wasm'
    );
    const bits = await workspace.fs.readFile(filename);
    const module = await WebAssembly.compile(bits);

    // Create the wasm worker that runs the LSP server
    const process = await wasm.createProcess(
      'lsp-server',
      module,
      { initial: 160, maximum: 160, shared: true },
      options
    );

    // Hook stderr to the output channel
    const decoder = new TextDecoder('utf-8');
    process.stderr!.onData(data => {
      channel.append(decoder.decode(data));
    });

    return startServer(process);
  };

  const clientOptions: LanguageClientOptions = {
    documentSelector: [{ language: 'plaintext' }],
    outputChannel: channel,
    uriConverters: createUriConverters()
  };

  let client = new LanguageClient('lspClient', 'LSP Client', serverOptions, clientOptions);
  await client.start();
}

このコードを実行すると、プレーンテキストファイルのコンテキストメニューにGoto Definitionエントリが追加されます。このアクションを実行すると、対応するリクエストがLSPサーバーに送信されます。

Running the goto definition action

重要な点として、@vscode/wasm-wasi-lsp npmモジュールは、ドキュメントURIをワークスペース値からWASI Preview 1ホストで認識される値に自動的に変換します。上記の例では、VS Code内のテキストドキュメントのURIは通常vscode-vfs://github/dbaeumer/plaintext-sample/lorem.txtのようなものですが、この値はWASIホスト内で認識されるfile:///workspace/lorem.txtに変換されます。この変換は、言語サーバーがURIをVS Codeに送り返す際にも自動的に行われます。

ほとんどの言語サーバーライブラリはカスタムメッセージをサポートしており、言語サーバープロトコル仕様にまだ存在しない機能を言語サーバーに追加するのが容易になります。以下のコードスニペットは、以前使用したRust言語サーバーに、指定されたワークスペースフォルダー内のファイルを数えるためのカスタムメッセージハンドラーを追加する方法を示しています。

#[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct CountFilesParams {
    pub folder: Url,
}

pub enum CountFilesRequest {}
impl Request for CountFilesRequest {
    type Params = CountFilesParams;
    type Result = u32;
    const METHOD: &'static str = "wasm-language-server/countFilesInDirectory";
}

//...

for msg in &connection.receiver {
    match msg {
		//....
		match cast::<CountFilesRequest>(req) {
    		Ok((id, params)) => {
				eprintln!("Received countFiles request #{} {}", id, params.folder);
        		let result = count_files_in_directory(&params.folder.path());
        		let json = serde_json::to_value(&result).unwrap();
        		let resp = Response { id, result: Some(json), error: None };
        		connection.sender.send(Message::Response(resp))?;
        		continue;
    		}
    		Err(err @ ExtractError::JsonError { .. }) => panic!("{err:?}"),
    		Err(ExtractError::MethodMismatch(req)) => req,
		}
	}
	//...
}

fn count_files_in_directory(path: &str) -> usize {
    WalkDir::new(path)
        .into_iter()
        .filter_map(Result::ok)
        .filter(|entry| entry.file_type().is_file())
        .count()
}

このカスタムリクエストをLSPサーバーに送信するTypeScriptコードは次のようになります。

const folder = workspace.workspaceFolders![0].uri;
const result = await client.sendRequest(CountFilesRequest, {
  folder: client.code2ProtocolConverter.asUri(folder)
});
window.showInformationMessage(`The workspace contains ${result} files.`);

これをvscode-languageserverリポジトリで実行すると、次の通知が表示されます。

Running count all files

言語サーバーは、言語サーバープロトコル仕様で指定されている機能のいずれかを実装する必要があるわけではないことに注意してください。拡張機能がWASI Preview 1ターゲットにのみコンパイルできるライブラリコードを統合したい場合、VS Codeがコンポーネントモデルの実装でWASI 0.2プレビューをサポートするまでは、カスタムメッセージを持つ言語サーバーを実装するのが良い選択肢かもしれません。

次に来るもの

前回のブログ記事で述べたように、VS Code用のWASI 0.2プレビューの実装に向けた取り組みを継続しています。また、コード例をRust以外のWASMにコンパイルされる言語を含むように広げる計画です。

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

DirkとVS Codeチーム

ハッピーコーディング！