locode/src/codeintel/lspManager.ts

import { spawn, type ChildProcess } from "node:child_process";
import path from "node:path";
import { readFile as fsReadFile } from "node:fs/promises";
import {
  createProtocolConnection,
  DidChangeTextDocumentNotification,
  DidOpenTextDocumentNotification,
  DefinitionRequest,
  ReferencesRequest,
  type ProtocolConnection,
  type TextDocumentIdentifier,
  type Position,
  type Location,
  type Diagnostic,
} from "vscode-languageserver-protocol";
import { StreamMessageReader, StreamMessageWriter } from 'vscode-languageserver-protocol/node';
import { URI } from "vscode-uri";
import type { MarkupContent } from "vscode-languageserver-protocol";

/** LSP diagnostic messages can be either a plain string or a { kind, value } MarkupContent object.
 *  locode's tool surface deals in plain strings, so flatten either form to text. */
function messageToString(message: string | MarkupContent): string {
  if (typeof message === "string") return message;
  return message?.value ?? "";
}

/** A connected LSP server for one language, plus its child process so we can clean it up. */
interface LspHandle {
  connection: ProtocolConnection;
  child: ChildProcess;
  languageId: string;
  /** Open documents we've already sent didOpen for, so we send didChange (not didOpen) on edits. */
  openDocs: Set<string>;
}

/** Maps a file extension to a language id (the LSP "languageId" string) and the server command to
 *  spawn for it. Only one server per language is ever spawned (lazy, on first use). A missing entry
 *  means locode has no built-in mapping — the user can still point a server at it via config in a
 *  future extension. The command is resolved on the PATH; if it isn't installed the spawn fails and
 *  the tool returns a clear "install X" error rather than a silent no-op. */
interface LanguageSpec {
  languageId: string;
  extensions: string[];
  /** The server command (no args). Must be on PATH. */
  command: string;
  /** Args passed to the server command. */
  args?: string[];
}

// The built-in language→server mappings. Mutable so `configureLanguageSpecs` can merge in user
// overrides/additions from config (see config.ts `lspServers`). One clangd spec covers both C
// and C++ — clangd handles both, and merging avoids spawning a second clangd for a mixed C/C++
// project (two servers keyed by separate languageIds would each index the same headers twice).
let LANGUAGE_SPECS: LanguageSpec[] = [
  // TypeScript / JavaScript — `typescript-language-server` wraps tsserver and speaks LSP. The most
  // common local-model codebase shape, so it's the first one locode wires up.
  {
    languageId: "typescript",
    extensions: [".ts", ".tsx", ".mts", ".cts", ".js", ".jsx", ".mjs", ".cjs"],
    command: "typescript-language-server",
    args: ["--stdio"],
  },
  {
    languageId: "python",
    extensions: [".py", ".pyi"],
    command: "pyright-langserver",
    args: ["--stdio"],
  },
  {
    languageId: "go",
    extensions: [".go"],
    command: "gopls",
    args: ["serve"],
  },
  {
    languageId: "rust",
    extensions: [".rs"],
    command: "rust-analyzer",
  },
  // C and C++ share clangd. The languageId is "c" (clangd treats .cpp/.hpp the same way);
  // all C/C++ extensions route to the single clangd process.
  {
    languageId: "c",
    extensions: [".c", ".h", ".cpp", ".cc", ".cxx", ".hpp", ".hh", ".hxx"],
    command: "clangd",
  },
];

/** Merge user-configured LSP server entries (from `locode config set lspServers`) into the
  *  built-in specs. An entry keyed by a built-in languageId overrides that spec's command/args
  *  and, if `extensions` is provided, which file extensions route to it. An entry keyed by a new
  *  languageId (e.g. "java", "ruby") adds a brand-new mapping — it MUST supply `extensions` so
  *  files can be routed to it. Call once at startup; idempotent against the built-in list.
  *
  *  Entries missing a `command` are ignored (a server we can't spawn is useless), and entries for
  *  new ids without `extensions` are ignored too (no way to route files to them). */
export function configureLanguageSpecs(overrides: Record<string, { command: string; args?: string[]; extensions?: string[] }>): void {
  const merged: LanguageSpec[] = LANGUAGE_SPECS.map((spec) => {
    const ov = overrides[spec.languageId];
    if (!ov) return spec;
    return {
      languageId: spec.languageId,
      extensions: ov.extensions ?? spec.extensions,
      command: ov.command,
      args: ov.args,
    };
  });
  for (const [languageId, ov] of Object.entries(overrides)) {
    if (merged.some((s) => s.languageId === languageId)) continue; // already a built-in we overrode
    if (!ov.command || !ov.extensions || ov.extensions.length === 0) continue;
    merged.push({ languageId, extensions: ov.extensions, command: ov.command, args: ov.args });
  }
  LANGUAGE_SPECS = merged;
}

/** Picks the LanguageSpec for a file path, or null if no extension matches. */
function specForFile(filePath: string): LanguageSpec | null {
  const ext = path.extname(filePath).toLowerCase();
  if (!ext) return null;
  return LANGUAGE_SPECS.find((s) => s.extensions.includes(ext)) ?? null;
}

/** A per-workspace (cwd) registry of live LSP servers, keyed by language id. One server per
 *  language per cwd — a second project gets its own manager (locode is single-session-per-process
 *  today, but keying on cwd keeps it correct if that ever changes). */
const handles = new Map<string, LspHandle>();

/** Convert an absolute filesystem path to an LSP file:// URI string. */
function toUri(absPath: string): string {
  return URI.file(absPath).toString();
}

interface LocResult {
  path: string;
  line: number;
  column: number;
}

function toLocation(loc: Location): LocResult {
  return {
    path: URI.parse(loc.uri).fsPath,
    line: loc.range.start.line + 1,
    column: loc.range.start.character + 1,
  };
}

/** Spawns the LSP server for `spec`, initializes it, and returns a live handle. Throws a clear,
 *  actionable error if the server binary isn't on the PATH (the most common failure) so the tool
 *  can surface "install typescript-language-server" instead of an opaque spawn ENOENT. */
async function startServer(spec: LanguageSpec, cwd: string): Promise<LspHandle> {
  let child: ChildProcess;
  try {
    child = spawn(spec.command, spec.args ?? [], { cwd, stdio: ["pipe", "pipe", "pipe"] });
  } catch (err) {
    throw new Error(
      `Could not start the LSP server "${spec.command}" for ${spec.languageId}. Is it installed and on your PATH? (${(err as Error).message})`,
    );
  }
  if (!child.stdin || !child.stdout) {
    child.kill();
    throw new Error(`LSP server "${spec.command}" did not open stdio streams.`);
  }

  const reader = new StreamMessageReader(child.stdout);
  const writer = new StreamMessageWriter(child.stdin);
  const connection = createProtocolConnection(reader, writer);

  // Surface stderr so a crashing server isn't a silent void (matches locode's MCP stdio policy).
  child.stderr?.on("data", () => {
    // Discard by default; a future debug mode could surface this. Don't let it back up.
  });

  await connection.sendRequest("initialize", {
    processId: process.pid,
    rootUri: URI.file(cwd).toString(),
    capabilities: {
      // locode consumes definition/references/diagnostics; declare only those so a server doesn't
      // waste effort enabling features we'll never query. Full text sync (change=1) is simplest and
      // correct — we always resend the whole file, never a range edit.
      textDocumentSync: { openClose: true, change: 1 },
      definitionProvider: true,
      referencesProvider: true,
    },
    workspaceFolders: [{ uri: URI.file(cwd).toString(), name: path.basename(cwd) || cwd }],
  });
  // Per LSP spec, the client must send `initialized` after the initialize response.
  await connection.sendNotification("initialized", {});

  // A server crash should reject any in-flight request rather than hanging forever — listen for
  // exit and dispose the connection so the next call throws instead of awaiting a dead process.
  child.on("exit", () => {
    connection.dispose();
    handles.delete(`${cwd}::${spec.languageId}`);
  });

  return { connection, child, languageId: spec.languageId, openDocs: new Set() };
}

/** Returns the (lazily-started) LSP handle for the language owning `filePath`, or throws if no
 *  server is configured/can't start. The first call for a language pays the initialize round-trip;
 *  every later call reuses the live server. */
async function handleForFile(filePath: string, cwd: string): Promise<LspHandle> {
  const spec = specForFile(filePath);
  if (!spec) {
    throw new Error(`No LSP server configured for "${path.extname(filePath)}" (code intelligence supports: ${LANGUAGE_SPECS.map((s) => s.extensions[0]).join(", ")}).`);
  }
  const key = `${cwd}::${spec.languageId}`;
  let handle = handles.get(key);
  if (!handle) {
    handle = await startServer(spec, cwd);
    handles.set(key, handle);
  }
  return handle;
}

/** Ensures the LSP server knows the current on-disk contents of `filePath`. Sends didOpen the
 *  first time a file is touched, didChange on subsequent syncs (the file was edited on disk since).
 *  Reads the file fresh each time — locode's tools write to disk before this runs, so the disk is
 *  the source of truth, not any in-memory buffer. */
async function syncDocument(handle: LspHandle, absPath: string, cwd: string): Promise<void> {
  const uri = toUri(absPath);
  const content = await fsReadFile(absPath, "utf-8");
  if (!handle.openDocs.has(uri)) {
    await handle.connection.sendNotification(DidOpenTextDocumentNotification.type, {
      textDocument: { uri, languageId: handle.languageId, version: 1, text: content },
    });
    handle.openDocs.add(uri);
  } else {
    await handle.connection.sendNotification(DidChangeTextDocumentNotification.type, {
      textDocument: { uri, version: Date.now() },
      contentChanges: [{ text: content }],
    });
  }
}

export interface DefinitionResult {
  /** The file/line/column of the symbol's definition. Multiple entries if the symbol has more than
   *  one definition (interface implementations, overloads, partial classes). Empty if the server
   *  found none (undefined symbol, or the server couldn't resolve it). */
  definitions: LocResult[];
}

/** Resolves where the symbol at `line`/`column` (1-indexed) in `filePath` is defined. Syncs the
 *  document first so the server's view matches disk. Returns an empty list (not an error) when the
 *  server has no definition to offer — that's a legitimate "not found", not a failure. */
export async function getDefinition(
  filePath: string,
  line: number,
  column: number,
  cwd: string,
): Promise<DefinitionResult> {
  const absPath = path.resolve(cwd, filePath);
  const handle = await handleForFile(filePath, cwd);
  await syncDocument(handle, absPath, cwd);
  const pos: Position = { line: line - 1, character: column - 1 };
  const result = (await handle.connection.sendRequest(DefinitionRequest.type, {
    textDocument: { uri: toUri(absPath) } as TextDocumentIdentifier,
    position: pos,
  })) as Location | Location[] | null;
  const locs = Array.isArray(result) ? result : result ? [result] : [];
  return { definitions: locs.map(toLocation) };
}

export interface ReferencesResult {
  /** Every place the symbol at `line`/`column` is referenced (including its definition). */
  references: LocResult[];
}

/** Finds every reference to the symbol at `line`/`column` in `filePath`. `includeDeclaration`
 *  defaults to true (matches most IDE "find all references" behavior). */
export async function getReferences(
  filePath: string,
  line: number,
  column: number,
  cwd: string,
  includeDeclaration = true,
): Promise<ReferencesResult> {
  const absPath = path.resolve(cwd, filePath);
  const handle = await handleForFile(filePath, cwd);
  await syncDocument(handle, absPath, cwd);
  const pos: Position = { line: line - 1, character: column - 1 };
  const result = (await handle.connection.sendRequest(ReferencesRequest.type, {
    textDocument: { uri: toUri(absPath) } as TextDocumentIdentifier,
    position: pos,
    context: { includeDeclaration },
  })) as Location[] | null;
  return { references: (result ?? []).map(toLocation) };
}

/** Notifies the LSP server that `filePath` changed on disk, so a subsequent `diagnostics` call
 *  reflects the new content. Called from the FileChanged hook path after edit_file/write_file. If no
 *  server is running for this language (or the file isn't one we manage), this is a no-op — it must
 *  never throw from a hook context, since hooks fire on every mutating tool. */
export async function notifyFileChanged(filePath: string, cwd: string): Promise<void> {
  try {
    const spec = specForFile(filePath);
    if (!spec) return;
    const key = `${cwd}::${spec.languageId}`;
    const handle = handles.get(key);
    if (!handle) return; // No server started yet — diagnostics will sync on first query.
    await syncDocument(handle, path.resolve(cwd, filePath), cwd);
  } catch {
    // Best-effort: a hook context can't propagate errors into the turn.
  }
}

type Severity = "error" | "warning" | "information" | "hint";

export interface DiagnosticsResult {
  diagnostics: { path: string; line: number; column: number; severity: Severity; message: string; source?: string }[];
}

/** The most recent diagnostics the server has published for `filePath`. LSP pushes diagnostics via
 *  `textDocument/publishDiagnostics` notifications; locode collects them per-URI as they arrive and
 *  returns the latest snapshot here. Forces a document sync first so the snapshot is current. */
const diagnosticsByUri = new Map<string, Diagnostic[]>();

// Per-URI resolvers waiting on the next publishDiagnostics notification. getDiagnostics arms one
// for the file it just synced, then races it against a timeout — so a slow server (tsserver on a
// large file) still gets a chance to publish the fresh snapshot rather than the caller reading a
// stale one after a single event-loop turn. Resolved and cleared by the publishDiagnostics handler.
const diagWaiters = new Map<string, () => void>();

/** Wait for the next publishDiagnostics for `uri`, or give up after `timeoutMs`. Resolves true if
 *  a publish arrived, false on timeout. The waiter is removed either way. */
function waitForDiagnostics(uri: string, timeoutMs: number): Promise<boolean> {
  return new Promise((resolve) => {
    const timer = setTimeout(() => {
      diagWaiters.delete(uri);
      resolve(false);
    }, timeoutMs);
    diagWaiters.set(uri, () => {
      clearTimeout(timer);
      diagWaiters.delete(uri);
      resolve(true);
    });
  });
}

const SEVERITY_MAP: Record<number, Severity> = {
  1: "error",
  2: "warning",
  3: "information",
  4: "hint",
};

export async function getDiagnostics(filePath: string, cwd: string): Promise<DiagnosticsResult> {
  const absPath = path.resolve(cwd, filePath);
  const handle = await handleForFile(filePath, cwd);
  const uri = toUri(absPath);
  // Attach a per-connection diagnostic collector the first time we use this handle.
  if (!(handle as unknown as { __diagWired?: boolean }).__diagWired) {
    (handle as unknown as { __diagWired?: boolean }).__diagWired = true;
    handle.connection.onNotification("textDocument/publishDiagnostics", (params: { uri: string; diagnostics: Diagnostic[] }) => {
      diagnosticsByUri.set(params.uri, params.diagnostics);
      // Wake a getDiagnostics call waiting on this URI, if any.
      diagWaiters.get(params.uri)?.();
    });
  }
  // Clear any stale snapshot for this URI before syncing so a timeout fallthrough can't return
  // diagnostics from before the edit. The server publishes asynchronously after didChange; race
  // its next publish against a short timeout so a slow server (tsserver on a large file) still
  // gets a chance to compute fresh diagnostics rather than us reading a stale snapshot after one
  // event-loop turn. Fall through to whatever's cached on timeout (possibly empty).
  diagnosticsByUri.delete(uri);
  await syncDocument(handle, absPath, cwd);
  await waitForDiagnostics(uri, 1500);
  const diags = diagnosticsByUri.get(uri) ?? [];
  return {
    diagnostics: diags.map((d) => ({
      path: URI.parse(uri).fsPath,
      line: (d.range?.start.line ?? 0) + 1,
      column: (d.range?.start.character ?? 0) + 1,
      severity: SEVERITY_MAP[d.severity ?? 1] ?? "information",
      message: messageToString(d.message ?? ""),
      source: d.source,
    })),
  };
}

/** Shuts down every live LSP server. Call on locode exit so spawned servers (tsserver, pyright,
 *  gopls, …) don't outlive the process as orphans. Awaits each shutdown so the signals land before
 *  teardown. Best-effort: a stuck server can't block exit forever (the child kill still fires). */
export async function shutdownAll(): Promise<void> {
  const all = [...handles.values()];
  handles.clear();
  await Promise.allSettled(
    all.map(async (h) => {
      try {
        await h.connection.sendRequest("shutdown", null);
        h.connection.sendNotification("exit", {});
      } catch {
        // Already dead — fall through to kill.
      }
      h.child.kill();
    }),
  );
}

/** For tests only: clear the live-handle registry and diagnostic cache without spawning/killing. */
export function _resetForTests(): void {
  handles.clear();
  diagnosticsByUri.clear();
  diagWaiters.clear();
}

/** For tests only: a snapshot of the currently configured language specs (after any
  *  configureLanguageSpecs merge), so tests can assert the merge without spawning a server. */
export function _specsForTests(): readonly LanguageSpec[] {
  return LANGUAGE_SPECS;
}

// The immutable built-in spec list, kept so tests can restore LANGUAGE_SPECS to defaults after a
// configureLanguageSpecs call (the merge is forward-only by design — production applies it once).
const BUILTIN_LANGUAGE_SPECS: readonly LanguageSpec[] = [
  { languageId: "typescript", extensions: [".ts", ".tsx", ".mts", ".cts", ".js", ".jsx", ".mjs", ".cjs"], command: "typescript-language-server", args: ["--stdio"] },
  { languageId: "python", extensions: [".py", ".pyi"], command: "pyright-langserver", args: ["--stdio"] },
  { languageId: "go", extensions: [".go"], command: "gopls", args: ["serve"] },
  { languageId: "rust", extensions: [".rs"], command: "rust-analyzer" },
  { languageId: "c", extensions: [".c", ".h", ".cpp", ".cc", ".cxx", ".hpp", ".hh", ".hxx"], command: "clangd" },
];

/** For tests only: restore the built-in language specs (undo any configureLanguageSpecs merge). */
export function _resetSpecsForTests(): void {
  LANGUAGE_SPECS = BUILTIN_LANGUAGE_SPECS.map((s) => ({ ...s }));
}