perf: run shellcheck directly from the lsp launcher (#11707)

docker/DockerfileExtra

@@ -71,7 +71,7 @@ ENV XDG_CACHE_HOME=/pyls/.cache
 RUN uv pip install --system --break-system-packages pipenv tornado python-lsp-jsonrpc ruff==0.16.0 Cython
 
 # Install Node-based language servers
-RUN npm install -g diagnostic-languageserver pyright
+RUN npm install -g pyright
 
 # Setup LSP working directory
 WORKDIR /pyls

docker/test_windmill_extra.ts

@@ -121,6 +121,60 @@ async function testLspWebSocket(): Promise<void> {
 	})
 }
 
+async function testLspShellDiagnostics(): Promise<void> {
+	return new Promise((resolve, reject) => {
+		const timeout = setTimeout(() => {
+			ws.close()
+			reject(new Error('no shellcheck diagnostics within 10s'))
+		}, 10000)
+
+		const ws = new WebSocket(`ws://${HOST}:${LSP_PORT}/ws/diagnostic`)
+		const send = (msg: object) => ws.send(JSON.stringify({ jsonrpc: '2.0', ...msg }))
+
+		ws.onopen = () => {
+			// Editors still send a linter definition; the server's own configuration
+			// applies, so the diagnostics below have to come from shellcheck.
+			send({
+				id: 1,
+				method: 'initialize',
+				params: {
+					capabilities: {},
+					initializationOptions: {
+						linters: { shellcheck: { command: 'true', sourceName: 'shellcheck' } },
+						filetypes: { shell: 'shellcheck' }
+					}
+				}
+			})
+			send({
+				method: 'textDocument/didOpen',
+				params: {
+					textDocument: {
+						uri: 'file:///tmp/monaco/test.sh',
+						languageId: 'shell',
+						version: 1,
+						text: '#!/bin/bash\necho "$unassigned"\n'
+					}
+				}
+			})
+		}
+
+		ws.onmessage = (event) => {
+			const msg = JSON.parse(event.data.toString())
+			if (msg.method !== 'textDocument/publishDiagnostics') return
+			clearTimeout(timeout)
+			ws.close()
+			const codes = msg.params.diagnostics.map((d: { code: string }) => d.code)
+			if (codes.includes('SC2154')) resolve()
+			else reject(new Error(`expected SC2154, got ${JSON.stringify(msg.params.diagnostics)}`))
+		}
+
+		ws.onerror = () => {
+			clearTimeout(timeout)
+			reject(new Error('LSP WebSocket connection failed'))
+		}
+	})
+}
+
 async function runLspTests(): Promise<boolean> {
 	log('\n=== LSP Service Tests (Port ' + LSP_PORT + ') ===')
 
@@ -132,6 +186,7 @@ async function runLspTests(): Promise<boolean> {
 	let passed = true
 	passed = (await runTest('LSP health check', testLspHealth)) && passed
 	passed = (await runTest('LSP WebSocket connection (pyright)', testLspWebSocket)) && passed
+	passed = (await runTest('LSP shell diagnostics', testLspShellDiagnostics)) && passed
 
 	return passed
 }

lsp/pyls_launcher.py

@@ -1,6 +1,9 @@
+import asyncio
 import logging
+import shutil
 import subprocess
 import threading
+import time
 import os
 import urllib.request
 import urllib.error
@@ -149,8 +152,221 @@ class LanguageServerWebSocketHandler(websocket.WebSocketHandler):
 class PyrightLS(LanguageServerWebSocketHandler):
     procargs = ["pipenv", "run", "pyright-langserver", "--stdio"]
 
-class DiagnosticLS(LanguageServerWebSocketHandler):
-    procargs = ["diagnostic-languageserver", "--stdio", "--log-level", "4"]
+SHELLCHECK_BIN = shutil.which("shellcheck")
+# The linter is configured here rather than by the editor's initializationOptions.
+SHELLCHECK_ARGS = ["--norc", "--format=json1", "-"]
+SHELLCHECK_TIMEOUT_SECS = float(os.environ.get("SHELLCHECK_TIMEOUT_SECS", "10"))
+SHELLCHECK_MAX_DOC_CHARS = int(os.environ.get("SHELLCHECK_MAX_DOC_CHARS", "200000"))
+SHELLCHECK_MAX_DOCS_PER_CONN = 8
+# A document whose checks finish faster than this is checked on every change
+# with no pause; see ShellcheckLS._changed.
+SHELLCHECK_INSTANT_SECS = 0.05
+SHELLCHECK_MAX_DEBOUNCE_SECS = 1.0
+# Shared by every connection: shellcheck is CPU-bound and single-threaded, so
+# more concurrent runs than cores only makes all of them slower.
+_shellcheck_slots = asyncio.Semaphore(os.cpu_count() or 2)
+
+_SHELLCHECK_SEVERITIES = {"error": 1, "warning": 2, "info": 3, "style": 3}
+
+
+def _utf16_col(line, col):
+    """shellcheck counts code points, LSP positions count UTF-16 code units."""
+    col = max(col - 1, 0)
+    if line.isascii():
+        return col
+    return len(line[:col].encode("utf-16-le")) // 2
+
+
+def shellcheck_diagnostics(text, output):
+    lines = text.split("\n")
+
+    def position(line, col):
+        line = max(line - 1, 0)
+        return {
+            "line": line,
+            "character": _utf16_col(lines[line], col) if line < len(lines) else 0,
+        }
+
+    diagnostics = []
+    for c in json.loads(output).get("comments", []):
+        code = "SC%d" % c["code"]
+        diagnostics.append(
+            {
+                "range": {
+                    "start": position(c["line"], c["column"]),
+                    "end": position(c["endLine"], c["endColumn"]),
+                },
+                "severity": _SHELLCHECK_SEVERITIES.get(c["level"], 3),
+                "code": code,
+                "codeDescription": {"href": "https://www.shellcheck.net/wiki/" + code},
+                "source": "shellcheck",
+                "message": c["message"],
+            }
+        )
+    return diagnostics
+
+
+class ShellcheckLS(websocket.WebSocketHandler):
+    """Lints shell documents by running shellcheck itself, speaking only the
+    slice of LSP that takes: document sync in, diagnostics out."""
+
+    def open(self):
+        self.docs = {}
+        self.tasks = {}
+        self.cost = {}
+        self.linted = {}
+
+    def check_origin(self, origin):
+        return True
+
+    def on_message(self, message):
+        try:
+            msg = json.loads(message)
+            method = msg.get("method")
+            params = msg.get("params") or {}
+            if method == "initialize":
+                self._send(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": msg.get("id"),
+                        "result": {
+                            "capabilities": {
+                                "textDocumentSync": {"openClose": True, "change": 1}
+                            },
+                            "serverInfo": {"name": "windmill-shellcheck"},
+                        },
+                    }
+                )
+            elif method == "textDocument/didOpen":
+                doc = params["textDocument"]
+                self._changed(doc["uri"], doc.get("version"), doc["text"])
+            elif method == "textDocument/didChange":
+                doc = params["textDocument"]
+                self._changed(
+                    doc["uri"], doc.get("version"), params["contentChanges"][-1]["text"]
+                )
+            elif method == "textDocument/didClose":
+                self._forget(params["textDocument"]["uri"])
+            elif method == "shutdown":
+                self._send({"jsonrpc": "2.0", "id": msg.get("id"), "result": None})
+            elif method is not None and "id" in msg:
+                self._send(
+                    {
+                        "jsonrpc": "2.0",
+                        "id": msg["id"],
+                        "error": {"code": -32601, "message": "Method not found"},
+                    }
+                )
+        except (ValueError, KeyError, IndexError, TypeError, AttributeError):
+            log.warning("shellcheck: dropping malformed message")
+
+    def on_close(self):
+        for uri in list(self.docs):
+            self._forget(uri)
+
+    def _send(self, msg):
+        try:
+            self.write_message(json.dumps(msg))
+        except websocket.WebSocketClosedError:
+            pass
+
+    def _forget(self, uri):
+        task = self.tasks.pop(uri, None)
+        if task:
+            task.cancel()
+        self.docs.pop(uri, None)
+        self.cost.pop(uri, None)
+        self.linted.pop(uri, None)
+
+    def _changed(self, uri, version, text):
+        if not isinstance(uri, str) or not isinstance(text, str):
+            raise TypeError
+        if uri not in self.docs and len(self.docs) >= SHELLCHECK_MAX_DOCS_PER_CONN:
+            return
+        if len(text) > SHELLCHECK_MAX_DOC_CHARS:
+            # Not checked, so not kept either: reported as having no findings.
+            text = None
+        self.docs[uri] = (version, text)
+        if uri in self.linted and self.linted[uri] == text:
+            return
+        running = self.tasks.pop(uri, None)
+        # The pause before a check is what that document's checks cost: one that
+        # is cheap is checked at once, one that is not waits out a typing burst
+        # rather than burning a core on results the next keystroke discards.
+        delay = self.cost.get(uri, 0)
+        if delay < SHELLCHECK_INSTANT_SECS:
+            delay = 0
+        if running and not running.done():
+            running.cancel()
+            delay = max(delay, SHELLCHECK_INSTANT_SECS)
+        self.tasks[uri] = asyncio.ensure_future(self._lint(uri, version, text, delay))
+
+    async def _lint(self, uri, version, text, delay):
+        if delay:
+            await asyncio.sleep(delay)
+        diagnostics = []
+        if SHELLCHECK_BIN and text is not None:
+            async with _shellcheck_slots:
+                started = time.monotonic()
+                try:
+                    output = await self._run(text)
+                    self.cost[uri] = 0
+                finally:
+                    # A killed check counts too, as a lower bound: it is the only
+                    # measure a document edited faster than it can be checked gets.
+                    if uri in self.docs:
+                        self.cost[uri] = min(
+                            max(time.monotonic() - started, self.cost.get(uri, 0)),
+                            SHELLCHECK_MAX_DEBOUNCE_SECS,
+                        )
+            if output is None:
+                return
+            try:
+                diagnostics = shellcheck_diagnostics(text, output)
+            except (ValueError, KeyError, TypeError) as e:
+                log.warning("shellcheck: unreadable output: %s", e)
+                return
+        if self.docs.get(uri) != (version, text):
+            return
+        self.linted[uri] = text
+        params = {"uri": uri, "diagnostics": diagnostics}
+        if version is not None:
+            params["version"] = version
+        self._send(
+            {
+                "jsonrpc": "2.0",
+                "method": "textDocument/publishDiagnostics",
+                "params": params,
+            }
+        )
+
+    async def _run(self, text):
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                SHELLCHECK_BIN,
+                *SHELLCHECK_ARGS,
+                stdin=asyncio.subprocess.PIPE,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.DEVNULL,
+            )
+        except OSError as e:
+            log.warning("shellcheck: could not start: %s", e)
+            return None
+        try:
+            output, _ = await asyncio.wait_for(
+                proc.communicate(text.encode("utf-8", "replace")),
+                SHELLCHECK_TIMEOUT_SECS,
+            )
+            return output
+        except asyncio.TimeoutError:
+            log.warning("shellcheck: timed out on a %d char document", len(text))
+            return None
+        finally:
+            # Also reached when a newer change cancels this check mid-run.
+            if proc.returncode is None:
+                proc.kill()
+                await proc.wait()
+
 
 class RuffLS(LanguageServerWebSocketHandler):
     procargs = ["ruff", "server"]
@@ -201,11 +417,14 @@ if __name__ == "__main__":
     # spawned `ruff server` picks it up.
     start_ruff_config_poller()
 
+    if not SHELLCHECK_BIN:
+        log.warning("shellcheck is not on PATH: shell scripts will not be linted")
+
     port = int(os.environ.get("PORT", "3001"))
     app = web.Application(
         [
             (r"/ws/pyright", PyrightLS),
-            (r"/ws/diagnostic", DiagnosticLS),
+            (r"/ws/diagnostic", ShellcheckLS),
             (r"/ws/ruff", RuffLS),
             (r"/ws/deno", DenoLS),
             (r"/ws/go", GoLS),