extension.ts

import * as vscode from 'vscode';
import { LanguageClient, LanguageClientOptions, ServerOptions, TransportKind } from 'vscode-languageclient/node';

import { type ResolvedReference, Resolver, parse } from '@authzed/spicedb-parser-js';

import { getInstallCommand, languageServerBinaryPath } from './binary';
import { CheckWatchProvider } from './checkwatchprovider';

export function activate(context: vscode.ExtensionContext) {
  console.log('spicedb-vscode is now active');

  const checkWatchProvider = new CheckWatchProvider(context.extensionUri);

  context.subscriptions.push(vscode.window.registerWebviewViewProvider(CheckWatchProvider.viewType, checkWatchProvider));

  context.subscriptions.push(
    vscode.commands.registerCommand('spicedb.addCheckWatch', () => {
      checkWatchProvider.addWatch();
    }),
  );

  context.subscriptions.push(
    vscode.window.onDidChangeActiveTextEditor((editor) => {
      if (editor) {
        checkWatchProvider.performUpdate(editor.document.uri.fsPath);
      }
    }),
  );

  context.subscriptions.push(
    vscode.workspace.onDidChangeTextDocument((e: vscode.TextDocumentChangeEvent) => {
      if (e.document.uri.fsPath.endsWith('.zed')) {
        if (vscode.window.activeTextEditor?.document.uri.fsPath === e.document.uri.fsPath) {
          checkWatchProvider.setActiveSchema(e.document.uri.fsPath, e.document.getText());
        }
      }

      if (e.document.uri.fsPath.endsWith('.zed.yaml')) {
        if (vscode.window.activeTextEditor?.document.uri.fsPath === e.document.uri.fsPath) {
          checkWatchProvider.setActiveYaml(e.document.uri.fsPath, e.document.getText());
        }
      }
    }),
  );

  // TODO: Move this into the language server.
  vscode.languages.registerDocumentSemanticTokensProvider(
    'spicedb',
    {
      provideDocumentSemanticTokens: function (
        document: vscode.TextDocument,
        _token: vscode.CancellationToken,
      ): vscode.ProviderResult<vscode.SemanticTokens> {
        const text = document.getText();
        const parserResult = parse(text);
        const data: number[] = [];
        if (parserResult.error) {
          return {
            data: new Uint32Array(data),
            resultId: undefined,
          };
        }

        // Data format:
        // - Line number (0-indexed, and offset from the *previous line*)
        // - Column position (0-indexed)
        // - Token length
        // - Token type index
        // - Modifier index
        let prevLine = 0;
        let prevChar = 0;

        const appendData = (lineNumber: number, colPosition: number, length: number, tokenType: number, modifierIndex: number) => {
          data.push(
            lineNumber - prevLine,
            prevLine === lineNumber ? colPosition - prevChar : colPosition,
            length,
            tokenType,
            modifierIndex,
          );

          prevLine = lineNumber;
          prevChar = colPosition;
        };

        // Resolve all type references and relation/permission references in expressions and color based on their kind and resolution
        // status.
        const resolution = new Resolver(parserResult.schema!);
        resolution.resolvedReferences().forEach((resolved: ResolvedReference) => {
          const lineNumber = resolved.reference.range.startIndex.line - 1; // parser ranges are 1-indexed
          const colPosition = resolved.reference.range.startIndex.column - 1;

          switch (resolved.kind) {
            case 'type':
              if (resolved.referencedTypeAndRelation === undefined) {
                appendData(lineNumber, colPosition, resolved.reference.path.length, /* type.unknown */ 3, 0);
                return;
              }

              appendData(lineNumber, colPosition, resolved.reference.path.length, /* type */ 0, 0);

              if (resolved.reference.relationName) {
                if (resolved.referencedTypeAndRelation.relation !== undefined) {
                  appendData(
                    lineNumber,
                    colPosition + 1 + resolved.reference.path.length,
                    resolved.reference.relationName.length,
                    /* member */ 2,
                    0,
                  );
                } else if (resolved.referencedTypeAndRelation.permission !== undefined) {
                  appendData(
                    lineNumber,
                    colPosition + 1 + resolved.reference.path.length,
                    resolved.reference.relationName.length,
                    /* property */ 1,
                    0,
                  );
                } else {
                  appendData(
                    lineNumber,
                    colPosition + 1 + resolved.reference.path.length,
                    resolved.reference.relationName.length,
                    /* member.unknown */ 3,
                    0,
                  );
                }
              }
              break;

            case 'expression':
              if (resolved.resolvedRelationOrPermission === undefined) {
                appendData(lineNumber, colPosition, resolved.reference.relationName.length, /* property.unknown */ 5, 0);
              } else {
                switch (resolved.resolvedRelationOrPermission.kind) {
                  case 'permission':
                    appendData(lineNumber, colPosition, resolved.reference.relationName.length, /* member */ 2, 0);
                    break;

                  case 'relation':
                    appendData(lineNumber, colPosition, resolved.reference.relationName.length, /* property */ 1, 0);
                    break;
                }
              }
              break;
          }
        });

        return {
          data: new Uint32Array(data),
          resultId: undefined,
        };
      },
    },
    {
      tokenTypes: ['type', 'property', 'member', 'type.unknown', 'member.unknown', 'property.unknown'],
      tokenModifiers: ['declaration'],
    },
  );

  startLanguageServer(context);
}

async function startLanguageServer(context: vscode.ExtensionContext) {
  // Start the LSP hooks using the language server found in the SpiceDB binary.
  const serverBinary = await languageServerBinaryPath(context);
  if (!serverBinary) {
    const installCommand = getInstallCommand();
    if (installCommand.startsWith('https://')) {
      const OpenInstall = 'Open Installation Instructions';
      vscode.window
        .showInformationMessage('`spicedb` binary is not installed. Click below to open the installation instructions.', OpenInstall)
        .then((selection) => {
          if (selection === OpenInstall) {
            vscode.env.openExternal(vscode.Uri.parse(installCommand));
          }
        });
    } else {
      vscode.window.showErrorMessage('`spicedb` binary is not installed. Please install it using `' + installCommand + '`');
    }

    console.error('Failed to find the SpiceDB language server binary');
    return;
  }

  const serverOptions: ServerOptions = {
    run: {
      command: serverBinary,
      args: ['lsp'],
      transport: TransportKind.stdio,
    },
    debug: {
      command: serverBinary,
      args: ['lsp'],
      transport: TransportKind.stdio,
    },
  };

  const clientOptions: LanguageClientOptions = {
    documentSelector: [{ scheme: 'file', language: 'spicedb' }],
    synchronize: {
      fileEvents: vscode.workspace.createFileSystemWatcher('**/.zed'),
    },
  };

  // Create the language client and start the client.
  const client = new LanguageClient('spicedbLanguageServer', 'SpiceDB Language Server', serverOptions, clientOptions);

  // Start the client. This will also launch the server.
  await client.start();
  console.log('spicedb-vscode is now active with its LSP running');
}

export function deactivate() {}