Simple LSP client, and the way to build twoslash

client.ts

import type { ChildProcessWithoutNullStreams } from 'node:child_process';

const LengthPrefix = 'Content-Length: ';

/// Simple language server client over JSON-RPC over stdio.
export class LanguageClient {
  readonly process: ChildProcessWithoutNullStreams;
  private readonly _messageListeners: Set<(msg: unknown) => void> = new Set();
  constructor(proc: ChildProcessWithoutNullStreams) {
    this.process = proc;
    const decoder = new TextDecoder();
    this.process.stdout.on('data', chunk => this._onStdout(typeof chunk == 'string' ? chunk : decoder.decode(chunk)));
  }
  public onMessage(listener: (msg: unknown) => void): void {
    this._messageListeners.add(listener);
  }
  private _nextId = 1;
  private readonly _requests: Map<number, (result: unknown) => void> = new Map();
  public request<T>(method: string, params: unknown = {}): Promise<T> {
    const id = this._nextId++;
    return new Promise<any>((resolve) => {
      this._requests.set(id, resolve);
      this._write({ jsonrpc: '2.0', id, method, params });
      // console.log('-->', method, params);
    });
  }
  private _write(message: unknown): void {
    const content = JSON.stringify(message);
    const body = `Content-Length: ${Buffer.byteLength(content, 'utf8')}\r\n\r\n${content}`;
    this.process.stdin.write(body);
  }
  private last = '';
  private _onStdout(chunk: string): void {
    this.last += chunk;
    while (this.last.startsWith(LengthPrefix)) {
      const endOfLength = this.last.indexOf('\r\n');
      if (endOfLength === -1) break;
      const lengthStr = this.last.slice(LengthPrefix.length, endOfLength);
      const length = parseInt(lengthStr, 10);
      const totalMessageLength = endOfLength + 4 + length;
      if (this.last.length >= totalMessageLength) {
        const message = this.last.slice(endOfLength + 4, totalMessageLength);
        this._onLine(message);
        this.last = this.last.slice(totalMessageLength);
      } else {
        break;
      }
    }
  }
  private _onLine(line: string): void {
    const message = JSON.parse(line);
    this._onMessage(message);
  }
  private _onMessage(message: any): void {
    // console.log('<--', message);
    if (message.id !== undefined && this._requests.has(message.id)) {
      const resolve = this._requests.get(message.id)!;
      this._requests.delete(message.id);
      resolve(message.result);
    } else {
      for (const listener of this._messageListeners) {
        listener(message);
      }
    }
  }
}

index.ts

export * from './types.ts'
export * from './twoslasher.ts'

twoslasher.ts

import type { StaticQuickInfo, TwoSlashLSPOptions, TwoSlashReturn } from './types.ts';

import { spawn } from 'node:child_process';
import { mkdtemp, rm, writeFile } from 'node:fs/promises';
import { homedir, tmpdir } from 'node:os';
import { join } from 'node:path';
import { pathToFileURL } from 'node:url';
import { LanguageClient } from './client.ts';

/// Runs the LSP server against the code snippet and returns type information on tokens.
///
/// @param code The twoslash markup'd code.
/// @param filename The filename to use for the temporary file.
/// @param options Options to setup the LSP server.
export async function twoslasher(code: string, filename: string, options: TwoSlashLSPOptions): Promise<TwoSlashReturn> {
  const workingDir = await mkdtemp(join(tmpdir(), 'twoslash-'));

  try {
    const entryFile = join(workingDir, filename);

    const codeLines = code.split(/\r\n?|\n/g);
    const bannerLines = options.banner?.split(/\r\n?|\n/g) ?? [];
    const footerLines = options.footer?.split(/\r\n?|\n/g) ?? [];

    const fullCode = [...bannerLines, ...codeLines, ...footerLines].join('\n');

    await writeFile(entryFile, fullCode);

    const server = spawn(options.command[0], options.command.slice(1), {
      cwd: workingDir,
      env: { ...process.env, ...options.env },
    });

    const client = new LanguageClient(server);

    await client.request('initialize', {
      processId: null,
      rootUri: pathToFileURL(workingDir).toString(),
      capabilities: {},
    });

    // Open the file
    await client.request('textDocument/didOpen', {
      textDocument: {
        uri: pathToFileURL(entryFile).toString(),
        languageId: 'untitled',
        version: 1,
        text: fullCode,
      },
    });

    const seen = new Set<string>();
    const staticQuickInfos: StaticQuickInfo[] = [];
    for (let line = 0; line < codeLines.length; line++) {
      const lineText = codeLines[line];
      for (let character = 0; character < lineText.length; character++) {
        interface HoverInfo {
          contents: string | { value: string } | Array<string | { value: string }>;
          range?: {
            start: { line: number; character: number };
            end: { line: number; character: number };
          };
        }
        const info = await client.request<HoverInfo | null>('textDocument/hover', {
          textDocument: { uri: pathToFileURL(entryFile).toString() },
          position: { line: line + bannerLines.length, character },
        });
        if (info?.contents) {
          // FIXME: cross-line ranges
          const start = info.range?.start.character ?? character;
          const length = (info.range?.end.character ?? character + 1) - (info.range?.start.character ?? character);
          const targetString = lineText.slice(start, start + length);
          const text = Array.isArray(info.contents)
            ? info.contents.map(c => (typeof c === 'string' ? c : c.value)).join('\n')
            : typeof info.contents === 'string'
              ? info.contents
              : info.contents.value;

          const key = `${start}|${length}|${text}`;
          if (seen.has(key)) {
            continue;
          }
          seen.add(key);

          staticQuickInfos.push({
            targetString,
            text,
            start,
            length,
            line,
            character,
          });
        }
      }
    }

    await client.request('shutdown');
    server.kill();

    return {
      code: codeLines.join('\n'),
      staticQuickInfos,
    };
  } catch (e) {
    console.error('Error during twoslash LSP processing:', e);
    return null!;
  } finally {
    await rm(workingDir, { recursive: true });
  }
}

const data = await twoslasher(`println("\\{@float.not_a_number}")`, 'a.mbt', {
  command: [join(homedir(), '.moon/bin/moonbit-lsp')],
  banner: 'fn main {',
  footer: '}',
})
console.log(data)

types.ts

export interface TwoSlashLSPOptions {
  /// The command to start the LSP server in stdio mode.
  readonly command: readonly string[];
  /// Additional environment variables.
  readonly env?: Readonly<Record<string, string | undefined>>;
  /// Helper utility so that you don't have to --cut-- everything out.
  readonly banner?: string;
  /// Helper utility so that you don't have to --cut-- everything out.
  readonly footer?: string;
}

export interface TwoSlashReturn {
  /// The input code, newlines will be normalized to \n.
  code: string;
  /// An array of static quick infos returned from the LSP server.
  staticQuickInfos: readonly StaticQuickInfo[];
}

export interface StaticQuickInfo {
  /// Content of the quick info node (mainly for debug).
  targetString: string;
  /// The base LSP response (the type).
  text: string;
  /// The index of the text in the file.
  start: number;
  /// How long the identifier is.
  length: number;
  /// Line number.
  line: number;
  /// Character on the line.
  character: number;
}