How we built a VS Code extension with Rust, WebAssembly, and TypeScript

We build Oso, a batteries-included framework for building authorization in your application. At the core of Oso is Polar, a declarative language for writing authorization policies. While someday we’ll release a feature that lets a fully-sentient Oso write your authorization policy for you, in the interim we thought it would be a good idea to improve the developer experience of writing Polar policies.

Along with Oso's 0.24 release, we shipped a VS Code extension that surfaces policy errors & warnings in VS Code instead of requiring you to run your application to see them. This shortens the development cycle for editing policies, giving you more time to spend on whatever sparks joy for you (and us more time to spend on making sure the kill orders issued by our sentient beta were a one-time glitch).

In this post, we'll talk through our design decisions in building the extension, and along the way we'll touch on the Language Server Protocol and, more generally, how an IDE extension works. If you’ve ever wanted to write an IDE extension for your language of choice, this post might be for you.

Why We Chose VS Code

We expect our users to write Polar in a variety of environments, and there are a wide range of IDEs out there. We restricted our initial focus to VS Code because we've had the most requests from our community for it. We also knew we were going to want to fan out the functionality beyond VS Code in the future, and building on top of the Language Server Protocol (LSP) would minimize the work required to port the extension's functionality to Vim, IntelliJ, Emacs, and other editors.

There's a huge range of features that an IDE integration can provide, like auto-completion and jump-to-definition. For our initial foray we scoped it down to surfacing policy errors & warnings. When writing Polar policies, these can be parsing errors (like "you're missing a semicolon") or validation errors (like "your policy references an unknown term"). There's a lot of value in surfacing those issues to the developer ASAP. Limiting ourselves to those diagnostics improved the developer experience without giving us an endless list of features to implement.

Using the Language Server Protocol (LSP)

Microsoft’s open-source Language Server Protocol abstracts a common set of IDE functionality across a JSON-RPC bridge. Every IDE that speaks LSP sends the same set of events across that bridge, like "a document has been opened" or "the user has requested all references for a particular symbol." Every language server that speaks LSP can implement features like documentation-on-hover once and get that functionality across every LSP-fluent editor.

Before the Language Server Protocol, if we wanted to build a Polar language server to surface errors & warnings in VS Code, we'd use VSCode-specific APIs. Then, when we wanted to port the same functionality to IntelliJ, we'd have to figure out the corresponding APIs in IntelliJ, and the same for Vim, Emacs, and every other IDE. It would be a lot of duplicated effort to build N slight variations of the same feature set.

Design

For the initial design, we knew a few things:

• We wanted to write the language server in Rust. The Polar interpreter is written in Rust, so writing the language server in Rust meant it would be easy to build on top of existing APIs in our Rust core. Also, the team knows Rust well, and it's generally our default for greenfield projects.
• The language server would be a new crate, polar-language-server, to avoid concerning polar-core with LSP-specific code.
• The extension would be written in TypeScript, because VS Code expects extensions to have a JavaScript entrypoint file that exposes hooks for VS Code to manage the extension's lifecycle.

This last point brings up the first big question we had to answer: where should the boundary be between TypeScript and Rust?

Why We Chose WebAssembly

If we want to maximize the amount we write in Rust, a minimal TypeScript extension would handle extension activation, kick off the Polar language server (written in Rust) in a separate process, and then get out of the way and let VS Code talk to the server via the Language Server Protocol. However, that explanation glosses over how we distribute the language server. VS Code users will install the extension via the VS Code extension marketplace, but how are they going to get the Polar language server onto their system?

If we write the server in Rust, we have four primary options for distribution. We can quickly rule out two of them:

1. Publish the new crate to crates.io and have folks cargo install the crate if they want to use the VS Code extension. This was easy to rule out because we can't ask folks to install a compatible Rust toolchain with Cargo just to use our VS Code extension.
2. Build platform-specific binaries from the Rust crate and either ask folks to manually download the correct one or do some host identification magic to attempt to download the correct one for their platform. This was also pretty easy to rule out because, as the proud parents of a release process consisting of well over 100 GitHub Actions jobs, we're all-too-familiar with the complexities of multi-platform builds.

We’re left with two viable options:

1. Leverage our existing multi-platform pipeline and bundle the language server into the Oso language libraries as an executable. We already ship an executable REPL with every library, so there's some prior art here.
2. Compile the server to WebAssembly (Wasm) and run the whole extension (client & server) in Node.js, VS Code's native runtime.

Comparing and contrasting bundling the language server into the language libraries (3) and compiling to Wasm (4) in the above list:

• Bundling would mean the version of Oso in the language server would always match the version installed in the project because the language server would come from the installed library. Going the Wasm route would mean we'd always ship the latest version of Oso in the extension, and folks whose projects weren't using the latest version might see different errors and warnings during development than at runtime if, for example, we introduce a new validation check or make a breaking change to the parser.
• Bundling would require writing & maintaining significantly more new code than compiling to Wasm. For bundling, we'd have to figure out how to slot the new language server API into the existing language libraries that we build and distribute for Python, Node.js, Ruby, Rust, Java, and Go, whereas compiling a new crate to Wasm wouldn't impact the language libraries at all.
• Bundling would result in distributing the extension & language server separately, so we would have to think about how to handle the two pieces running different major versions of the Language Server Protocol. It's fair to assume this wouldn't be a common issue, but it's still something we don't need to think about at all with the Wasm option, where we distribute the language server with the extension as an additional .wasm file.
  • It's also worth noting that the Language Server Protocol startup handshake involves the server telling the client which pieces of the protocol it supports, so, even if we distributed the extension and language server separately (the bundling option), we'd be able to enable new client functionality without breaking backwards compatibility with older servers running the same major LSP version.

Both options involve tradeoffs, and we really could've gone either way. However, we ultimately decided to go with option 4—compiling the server to WebAssembly and packaging the client and server together as a JavaScript package—because it seemed like it would be quicker to implement, would make packaging and distribution relatively straightforward, and the language server scaffolding packages in JavaScript's ecosystem are more mature than Rust's.

Deciding Against a Rust LSP Library

There are three main Rust crates that provide scaffolding for a language server: tower-lsp, lspower (a fork of tower-lsp), and lsp-server. Both tower-lsp and lspower require implementing the language server as an asynchronous service powered by Tower, a framework for async Rust. We didn't think the benefits of asynchrony would be worth the added complexity to the language server. (For what it's worth, the folks behind Rust Analyzer, the most mature LSP implementation in Rust, seem to agree.) For one, the extension runs in a separate process and doesn't impact the editor's responsiveness, and anyway parsing a Polar policy and surfacing errors and warnings is a very fast operation (~1 ms). Additionally, there are open questions about the model at the core of tower-lsp and lspower and whether Tower's request/response model is a natural fit for the Language Server Protocol, which more naturally maps to a pair of streams from client→server and server→client. We also looked into lsp-server, the synchronous language server scaffold used by Rust Analyzer, but the lack of documentation and explicit aim to not be "a good general purpose library at the moment" scared us off.

Ultimately, we decided that none of the three scaffolding crates were worth the tradeoffs. We decided to write the language server ourselves.

Implementation

With our sights set on Wasm, implementing the language server in Rust largely went off without a hitch, and we landed on a clean, three-part design consisting of a TypeScript entrypoint module, a second, tiny TypeScript module that delegates connection handling to Microsoft's vscode-languageserver package, and the polar-language-server crate compiled to Wasm.

The extension's entrypoint is a TypeScript module that exposes activate() and deactivate() functions so that VS Code can manage the extension's lifecycle. On activation, the extension kicks off a separate language server process for each VS Code workspace in order to support multi-root workspaces.

There are a few layers between the process that we kick off per workspace and the polar-language-server Wasm module. At the top is an instance of the LanguageClient class provided by Microsoft's vscode-languageclient package. The LanguageClient instance is in charge of kicking off the language server, handling communication between VS Code and the server, and passing some configuration to VS Code like which types of files we're interested in for a particular workspace.

When the LanguageClient boots up the tiny "server" TypeScript module, we initialize the Wasm-ified polar-language-server and manage the server side of the LSP connection from TypeScript. By keeping connection management in TypeScript, we're able to delegate all of that logic to the vscode-languageserver package instead of having to reimplement it ourselves in polar-language-server. We forward all LSP notification messages (e.g., textDocument/didOpen) straight to polar-language-server, where the PolarLanguageServer::on_notification method de-serializes the message from JSON into a Rust type provided by the excellent lsp_types crate. We also pass in a JavaScript callback to the polar-language-server constructor so that it can send diagnostics back across the LSP divide to VS Code.

Unencumbered by all of the busywork we delegated to Microsoft's JavaScript packages, the polar-language-server crate has one job: whenever the state of a particular workspace's Polar policy changes (i.e., a .polar file was added, changed, or deleted), it reloads the policy, translates any errors or warnings into VS Code diagnostics, and punts those diagnostics back across the fence to VS Code, which then displays them inline in the editor.

Wrapping Up

By splitting the language server into a thin TypeScript wrapper and a Rust core, we were able to build on top of some nice LSP scaffolding that exists in the JavaScript ecosystem while still implementing all of the new diagnostic-wrangling logic in Rust. The extension now surfaces issues as the policy is edited in VS Code, resulting in less time flipping back-and-forth between editor and terminal and more time focused on writing the authorization policy:

We’d love to hear feedback on the new extension and/or any features or editors you’d like to see supported; please join us on Slack or open a GitHub issue. Also, if this project sounds like something that would be fun to get paid to work on, we’re hiring!