import { spawn } from "node:child_process"; import { mkdtempSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import * as path from "node:path"; import { pathToFileURL } from "node:url"; import type { InlineFormatInspectionBackend, InlineFormatInspectionDiagnostic, InlineFormatInspectionPosition, InlineFormatInspectionRange, InlineFormatInspectionRequest, InlineFormatInspectionResult, InlineFormatSemanticToken, } from "./index.js"; type JsonRpcResponse = { id?: number; result?: unknown; error?: { message?: string; }; }; type JsonRpcNotification = { method?: string; params?: unknown; }; type LspPosition = { line: number; character: number; }; type LspRange = { start: LspPosition; end: LspPosition; }; type LspDiagnostic = { severity?: number; message: string; range: LspRange; source?: string; code?: number | string; }; type LspPublishDiagnosticsParams = { uri?: string; diagnostics?: LspDiagnostic[]; }; type LspLocation = { uri: string; range: LspRange; }; type LspLocationLink = { targetUri: string; targetRange: LspRange; targetSelectionRange: LspRange; }; type LspSemanticTokensLegend = { tokenTypes?: string[]; tokenModifiers?: string[]; }; type LspSemanticTokens = { data?: number[]; }; type LspInitializeResult = { capabilities?: { semanticTokensProvider?: { legend?: LspSemanticTokensLegend; }; }; }; type NotificationWaiter = { method: string; predicate: (params: unknown) => boolean; resolve: (params: unknown) => void; reject: (error: Error) => void; }; const PYTHON_LANGUAGE = "python"; const BASEDPYRIGHT_BACKEND_NAME = "inline-format-basedpyright"; const LSP_TIMEOUT_MS = 8000; function createScaffoldSummary( request: InlineFormatInspectionRequest, extra: string, ): string { return `Prepared virtual ${request.document.language} document ${request.document.filePath} for ${request.kind} inspection. ${extra}`; } function escapeSymbolNameForRegex(symbolName: string): string { return symbolName.replace(/[.*+?^${}()|[\]\\]/gu, "\\$&"); } function findTextualSymbolRanges( source: string, symbolName: string, ): InlineFormatInspectionRange[] { const pattern = new RegExp( `\\b${escapeSymbolNameForRegex(symbolName)}\\b`, "gu", ); return source.split("\n").flatMap((line, lineIndex) => Array.from(line.matchAll(pattern), (match) => { const startColumn = match.index ?? 0; return { start: { lineIndex, columnIndex: startColumn, }, end: { lineIndex, columnIndex: startColumn + symbolName.length, }, }; }), ); } function findFirstIdentifierPosition( source: string, ): InlineFormatInspectionPosition | null { const lines = source.split("\n"); const identifierPattern = /[A-Za-z_][A-Za-z0-9_]*/u; for (const [lineIndex, line] of lines.entries()) { const match = identifierPattern.exec(line); if (match?.index !== undefined) { return { lineIndex, columnIndex: match.index, }; } } return null; } function findIdentifierRangeAtPosition( source: string, position: InlineFormatInspectionPosition, ): InlineFormatInspectionRange | undefined { const line = source.split("\n")[position.lineIndex]; if (line === undefined) { return undefined; } const identifierPattern = /[A-Za-z_][A-Za-z0-9_]*/gu; for (const match of line.matchAll(identifierPattern)) { const startColumn = match.index; if (startColumn === undefined) { continue; } const endColumn = startColumn + match[0].length; if ( position.columnIndex < startColumn || position.columnIndex >= endColumn ) { continue; } return { start: { lineIndex: position.lineIndex, columnIndex: startColumn, }, end: { lineIndex: position.lineIndex, columnIndex: endColumn, }, }; } return undefined; } function resolvePythonInspectionPosition( request: InlineFormatInspectionRequest, ): InlineFormatInspectionPosition | null { if (request.position !== undefined) { return request.position; } const symbolName = request.symbolName?.trim(); if (symbolName !== undefined && symbolName.length > 0) { const symbolRanges = findTextualSymbolRanges( request.document.content, symbolName, ); return symbolRanges[0]?.start ?? null; } return findFirstIdentifierPosition(request.document.content); } function toLspPosition(position: InlineFormatInspectionPosition): LspPosition { return { line: position.lineIndex, character: position.columnIndex, }; } function fromLspRange(range: LspRange): InlineFormatInspectionRange { return { start: { lineIndex: range.start.line, columnIndex: range.start.character, }, end: { lineIndex: range.end.line, columnIndex: range.end.character, }, }; } function diagnosticSeverityFromLsp( severity: number | undefined, ): InlineFormatInspectionDiagnostic["severity"] { switch (severity) { case 1: return "error"; case 2: return "warning"; default: return "info"; } } function mapDiagnostics( diagnostics: readonly LspDiagnostic[], ): InlineFormatInspectionDiagnostic[] { return diagnostics.map((diagnostic) => ({ severity: diagnosticSeverityFromLsp(diagnostic.severity), message: diagnostic.message, range: fromLspRange(diagnostic.range), ...(diagnostic.source !== undefined ? { source: diagnostic.source } : {}), ...(diagnostic.code !== undefined ? { code: String(diagnostic.code) } : {}), })); } function normalizeHoverText(contents: unknown): string | null { if (typeof contents === "string") { return contents.trim() || null; } if ( typeof contents === "object" && contents !== null && "value" in contents && typeof contents.value === "string" ) { return contents.value.trim() || null; } if (Array.isArray(contents)) { const parts = contents .map((entry) => normalizeHoverText(entry)) .filter((entry): entry is string => entry !== null); return parts.length > 0 ? parts.join("\n") : null; } return null; } function normalizeLocations(result: unknown): LspLocation[] { if (!Array.isArray(result)) { return []; } return result.flatMap((entry) => { if ( typeof entry === "object" && entry !== null && "uri" in entry && "range" in entry ) { const uri = typeof entry.uri === "string" ? entry.uri : undefined; const range = entry.range as LspRange | undefined; return uri !== undefined && range !== undefined ? [{ uri, range }] : []; } if ( typeof entry === "object" && entry !== null && "targetUri" in entry && ("targetSelectionRange" in entry || "targetRange" in entry) ) { const link = entry as LspLocationLink; return [ { uri: link.targetUri, range: link.targetSelectionRange ?? link.targetRange, }, ]; } return []; }); } function normalizeSemanticTokensLegend( result: unknown, ): LspSemanticTokensLegend | null { if ( typeof result !== "object" || result === null || !("capabilities" in result) ) { return null; } const capabilities = result.capabilities; if ( typeof capabilities !== "object" || capabilities === null || !("semanticTokensProvider" in capabilities) ) { return null; } const semanticTokensProvider = capabilities.semanticTokensProvider; if ( typeof semanticTokensProvider !== "object" || semanticTokensProvider === null || !("legend" in semanticTokensProvider) ) { return null; } const legend = semanticTokensProvider.legend; const legendRecord = legend as { tokenTypes?: unknown; tokenModifiers?: unknown; }; const tokenTypes = Array.isArray(legendRecord.tokenTypes) && legendRecord.tokenTypes.every( (entry): entry is string => typeof entry === "string", ) ? legendRecord.tokenTypes : undefined; const tokenModifiers = Array.isArray(legendRecord.tokenModifiers) && legendRecord.tokenModifiers.every( (entry): entry is string => typeof entry === "string", ) ? legendRecord.tokenModifiers : undefined; if (tokenTypes === undefined && tokenModifiers === undefined) { return null; } return { ...(tokenTypes !== undefined ? { tokenTypes } : {}), ...(tokenModifiers !== undefined ? { tokenModifiers } : {}), }; } function getDocumentLine( source: string, lineIndex: number, ): string | undefined { return source.split("\n")[lineIndex]; } function decodeSemanticTokenModifiers( encodedModifiers: number, legend: LspSemanticTokensLegend, ): string[] { const modifiers = legend.tokenModifiers ?? []; return modifiers.filter( (_, index) => (encodedModifiers & (1 << index)) !== 0, ); } function mapSemanticTokens( source: string, encoded: LspSemanticTokens | null, legend: LspSemanticTokensLegend | null, ): { legend: LspSemanticTokensLegend | null; tokens: InlineFormatSemanticToken[]; } { const tokenTypes = legend?.tokenTypes ?? []; const data = encoded?.data ?? []; const tokens = [] as { range: InlineFormatInspectionRange; tokenType: string; modifiers: string[]; text?: string; }[]; let line = 0; let column = 0; for (let index = 0; index < data.length; index += 5) { const deltaLine = data[index] ?? 0; const deltaStart = data[index + 1] ?? 0; const length = data[index + 2] ?? 0; const tokenTypeIndex = data[index + 3] ?? 0; const modifierBits = data[index + 4] ?? 0; line += deltaLine; column = deltaLine === 0 ? column + deltaStart : deltaStart; if (length <= 0) { continue; } const lineText = getDocumentLine(source, line); if (lineText === undefined) { continue; } const endColumn = Math.min(column + length, lineText.length); if (endColumn <= column) { continue; } const tokenType = tokenTypes[tokenTypeIndex] ?? `token-${String(tokenTypeIndex)}`; const text = lineText.slice(column, endColumn).trim(); tokens.push({ range: { start: { lineIndex: line, columnIndex: column, }, end: { lineIndex: line, columnIndex: endColumn, }, }, tokenType, modifiers: decodeSemanticTokenModifiers(modifierBits, legend ?? {}), ...(text.length > 0 ? { text } : {}), }); } return { legend, tokens }; } class JsonRpcStdioClient { private readonly process = spawn("basedpyright-langserver", ["--stdio"], { stdio: ["pipe", "pipe", "pipe"], }); private readonly pendingRequests = new Map< number, { resolve: (value: unknown) => void; reject: (error: Error) => void; } >(); private readonly notificationWaiters: NotificationWaiter[] = []; private readonly stderrChunks: string[] = []; private stdoutBuffer = Buffer.alloc(0); private nextRequestId = 0; constructor() { this.process.stdout.on("data", (chunk: Buffer) => { this.stdoutBuffer = Buffer.concat([this.stdoutBuffer, chunk]); this.processStdoutBuffer(); }); this.process.stderr.on("data", (chunk: Buffer) => { this.stderrChunks.push(chunk.toString("utf8")); }); this.process.on("exit", (code, signal) => { const error = new Error( `basedpyright-langserver exited before the request completed (code=${String(code)} signal=${String(signal)} stderr=${this.stderrChunks.join("")}).`, ); for (const pending of this.pendingRequests.values()) { pending.reject(error); } this.pendingRequests.clear(); while (this.notificationWaiters.length > 0) { this.notificationWaiters.shift()?.reject(error); } }); } dispose(): void { this.process.kill(); } async request(method: string, params: unknown): Promise { const id = this.nextRequestId++; const promise = new Promise((resolve, reject) => { this.pendingRequests.set(id, { resolve, reject }); }); this.writeMessage({ jsonrpc: "2.0", id, method, params, }); return await withTimeout( promise, LSP_TIMEOUT_MS, `Timed out waiting for ${method} response from basedpyright-langserver.`, ); } notify(method: string, params: unknown): void { this.writeMessage({ jsonrpc: "2.0", method, params, }); } async waitForNotification( method: string, predicate: (params: unknown) => boolean, ): Promise { const promise = new Promise((resolve, reject) => { this.notificationWaiters.push({ method, predicate, resolve, reject, }); }); return await withTimeout( promise, LSP_TIMEOUT_MS, `Timed out waiting for ${method} notification from basedpyright-langserver.`, ); } private processStdoutBuffer(): void { const separator = Buffer.from("\r\n\r\n", "utf8"); while (true) { const headerEnd = this.stdoutBuffer.indexOf(separator); if (headerEnd === -1) { return; } const headerText = this.stdoutBuffer.slice(0, headerEnd).toString("utf8"); const contentLengthMatch = /Content-Length:\s*(\d+)/iu.exec(headerText); if (contentLengthMatch?.[1] === undefined) { throw new Error( `Invalid LSP header from basedpyright-langserver: ${headerText}`, ); } const contentLength = Number(contentLengthMatch[1]); const messageStart = headerEnd + separator.length; const messageEnd = messageStart + contentLength; if (this.stdoutBuffer.length < messageEnd) { return; } const messageText = this.stdoutBuffer .slice(messageStart, messageEnd) .toString("utf8"); this.stdoutBuffer = this.stdoutBuffer.slice(messageEnd); this.dispatchMessage( JSON.parse(messageText) as JsonRpcResponse & JsonRpcNotification, ); } } private dispatchMessage( message: JsonRpcResponse & JsonRpcNotification, ): void { if (typeof message.id === "number") { const pending = this.pendingRequests.get(message.id); if (pending === undefined) { return; } this.pendingRequests.delete(message.id); if (message.error !== undefined) { pending.reject( new Error(message.error.message ?? "Unknown LSP error."), ); return; } pending.resolve(message.result); return; } if (message.method === undefined) { return; } const waiterIndex = this.notificationWaiters.findIndex( (waiter) => waiter.method === message.method && waiter.predicate(message.params), ); if (waiterIndex === -1) { return; } const [waiter] = this.notificationWaiters.splice(waiterIndex, 1); waiter?.resolve(message.params); } private writeMessage(message: Record): void { const body = Buffer.from(JSON.stringify(message), "utf8"); const header = Buffer.from( `Content-Length: ${String(body.length)}\r\n\r\n`, "utf8", ); this.process.stdin.write(Buffer.concat([header, body])); } } async function withTimeout( promise: Promise, timeoutMs: number, message: string, ): Promise { let timeoutHandle: NodeJS.Timeout | undefined; try { return await Promise.race([ promise, new Promise((_, reject) => { timeoutHandle = setTimeout(() => reject(new Error(message)), timeoutMs); }), ]); } finally { if (timeoutHandle !== undefined) { clearTimeout(timeoutHandle); } } } async function withBasedPyrightSession( request: InlineFormatInspectionRequest, callback: (context: { client: JsonRpcStdioClient; filePath: string; fileUri: string; initialDiagnostics: readonly LspDiagnostic[]; semanticTokensLegend: LspSemanticTokensLegend | null; }) => Promise, ): Promise { const tempDir = mkdtempSync( path.join(tmpdir(), "pi-inline-format-basedpyright-"), ); const filePath = path.join( tempDir, path.basename(request.document.filePath) || "inline-snippet.py", ); const fileUri = pathToFileURL(filePath).toString(); writeFileSync(filePath, request.document.content, "utf8"); const client = new JsonRpcStdioClient(); try { const initializeResult = (await client.request("initialize", { processId: process.pid, clientInfo: { name: "pi-inline-format-intel", }, rootUri: pathToFileURL( request.document.region.projectRoot ?? tempDir, ).toString(), capabilities: {}, workspaceFolders: [ { uri: pathToFileURL( request.document.region.projectRoot ?? tempDir, ).toString(), name: path.basename(request.document.region.projectRoot ?? tempDir), }, ], })) as LspInitializeResult; client.notify("initialized", {}); client.notify("textDocument/didOpen", { textDocument: { uri: fileUri, languageId: PYTHON_LANGUAGE, version: 1, text: request.document.content, }, }); const diagnosticsParams = (await client.waitForNotification( "textDocument/publishDiagnostics", (params) => typeof params === "object" && params !== null && "uri" in params && params.uri === fileUri, )) as LspPublishDiagnosticsParams; return await callback({ client, filePath, fileUri, initialDiagnostics: diagnosticsParams.diagnostics ?? [], semanticTokensLegend: normalizeSemanticTokensLegend(initializeResult), }); } finally { client.dispose(); rmSync(tempDir, { recursive: true, force: true }); } } function createDefinitionSummary( symbolName: string | undefined, definitionCount: number, sameFileDefinitionCount: number, ): string { const label = symbolName ?? "the selected symbol"; if (definitionCount === 0) { return `Basedpyright could not resolve a definition for ${label}.`; } if (sameFileDefinitionCount === definitionCount) { return `Basedpyright resolved ${String(definitionCount)} definition(s) for ${label}.`; } return `Basedpyright resolved ${String(definitionCount)} definition(s) for ${label}; ${String(sameFileDefinitionCount)} map directly into the current virtual document.`; } function createHoverSummary(hoverText: string | null): string { if (hoverText !== null) { return `Resolved hover information via basedpyright. ${hoverText}`; } return "Resolved hover request via basedpyright, but hover text was empty."; } function createDiagnosticsSummary(count: number): string { return count === 0 ? "Basedpyright reported no diagnostics for the current virtual document." : `Basedpyright reported ${String(count)} diagnostic(s) for the current virtual document.`; } function createSemanticTokensSummary(count: number): string { return count === 0 ? "Basedpyright reported no semantic tokens for the current virtual document." : `Basedpyright reported ${String(count)} semantic token(s) for the current virtual document.`; } export const basedPyrightInspectionBackend: InlineFormatInspectionBackend = { name: BASEDPYRIGHT_BACKEND_NAME, languages: [PYTHON_LANGUAGE], async inspect(request): Promise { if (request.document.language !== PYTHON_LANGUAGE) { return null; } const resolvedPosition = resolvePythonInspectionPosition(request); const symbolName = request.symbolName?.trim(); if (request.kind === "document-highlights") { const textualRanges = symbolName === undefined || symbolName.length === 0 ? [] : findTextualSymbolRanges(request.document.content, symbolName); return { backendName: this.name, language: request.document.language, kind: request.kind, summary: createScaffoldSummary( request, symbolName === undefined || symbolName.length === 0 ? "Document highlights need a symbol name or position. basedpyright prototype support is not wired for this surface yet." : `Document highlights for ${symbolName} are not wired yet in the basedpyright prototype. Textual occurrences: ${String(textualRanges.length)}.`, ), ...(textualRanges.length > 0 ? { ranges: textualRanges } : {}), ...(symbolName !== undefined && symbolName.length > 0 ? { payload: { symbolName, textualOccurrenceCount: textualRanges.length, }, } : {}), }; } return await withBasedPyrightSession(request, async (context) => { if (request.kind === "diagnostics") { const diagnostics = mapDiagnostics(context.initialDiagnostics); return { backendName: this.name, language: request.document.language, kind: request.kind, summary: createDiagnosticsSummary(diagnostics.length), diagnostics, payload: { diagnosticCount: diagnostics.length, }, }; } if (request.kind === "semantic-tokens") { const { legend, tokens } = mapSemanticTokens( request.document.content, (await context.client.request("textDocument/semanticTokens/full", { textDocument: { uri: context.fileUri, }, })) as LspSemanticTokens | null, context.semanticTokensLegend, ); return { backendName: this.name, language: request.document.language, kind: request.kind, summary: createSemanticTokensSummary(tokens.length), ...(tokens.length > 0 ? { ranges: tokens.map((token) => token.range) } : {}), payload: { tokenCount: tokens.length, tokens, ...(legend !== null ? { legend } : {}), }, }; } if (resolvedPosition === null) { return { backendName: this.name, language: request.document.language, kind: request.kind, summary: symbolName !== undefined && symbolName.length > 0 ? `Basedpyright could not locate symbol ${symbolName} in the current virtual document.` : "Basedpyright could not locate a suitable symbol in the current virtual document.", }; } if (request.kind === "definition") { const locations = normalizeLocations( await context.client.request("textDocument/definition", { textDocument: { uri: context.fileUri, }, position: toLspPosition(resolvedPosition), }), ); const sameFileRanges = locations .filter((location) => location.uri === context.fileUri) .map((location) => fromLspRange(location.range)); return { backendName: this.name, language: request.document.language, kind: request.kind, summary: createDefinitionSummary( symbolName, locations.length, sameFileRanges.length, ), ...(sameFileRanges.length > 0 ? { ranges: sameFileRanges } : {}), payload: { ...(symbolName !== undefined && symbolName.length > 0 ? { symbolName } : {}), definitionCount: locations.length, sameFileDefinitionCount: sameFileRanges.length, definitionFiles: [ ...new Set(locations.map((location) => location.uri)), ], }, }; } const hoverResult = (await context.client.request("textDocument/hover", { textDocument: { uri: context.fileUri, }, position: toLspPosition(resolvedPosition), })) as { contents?: unknown; range?: LspRange; } | null; const hoverText = normalizeHoverText(hoverResult?.contents); const hoverRange = hoverResult?.range !== undefined ? fromLspRange(hoverResult.range) : undefined; const fallbackHoverRange = hoverRange ?? findIdentifierRangeAtPosition( request.document.content, resolvedPosition, ); if (request.kind === "explain-symbol") { const textualRanges = symbolName === undefined || symbolName.length === 0 ? [] : findTextualSymbolRanges(request.document.content, symbolName); return { backendName: this.name, language: request.document.language, kind: request.kind, summary: symbolName === undefined || symbolName.length === 0 ? createHoverSummary(hoverText) : `Explained symbol ${symbolName} via basedpyright. ${hoverText ?? "Hover text was empty."}`, ...(textualRanges.length > 0 ? { ranges: textualRanges } : fallbackHoverRange !== undefined ? { ranges: [fallbackHoverRange] } : {}), payload: { ...(symbolName !== undefined && symbolName.length > 0 ? { symbolName } : {}), hover: hoverText, }, }; } return { backendName: this.name, language: request.document.language, kind: request.kind, summary: symbolName === undefined || symbolName.length === 0 ? createHoverSummary(hoverText) : `Inspected symbol ${symbolName} via basedpyright. ${hoverText ?? "Hover text was empty."}`, ...(fallbackHoverRange !== undefined ? { ranges: [fallbackHoverRange] } : {}), payload: { ...(symbolName !== undefined && symbolName.length > 0 ? { symbolName } : {}), hover: hoverText, }, }; }); }, };