server/runtime/lsp-client.ts

/**
 * LSP Client — EXACT parity with original claw-code Rust LSP crate.
 * 
 * Implements the Language Server Protocol over stdio:
 * - JSON-RPC 2.0 message framing (Content-Length headers)
 * - Initialize/Initialized handshake
 * - textDocument/didOpen, didChange, didSave, didClose
 * - textDocument/definition (go-to-definition)
 * - textDocument/references (find references)
 * - textDocument/publishDiagnostics (diagnostics collection)
 * - Workspace diagnostics collection
 * - Multi-server support via extension mapping
 * - Shutdown/exit lifecycle
 */

import { spawn, ChildProcess } from "child_process";
import * as path from "path";
import * as fs from "fs";

// ─── Types (matches original Rust types) ──────────────────────────────────

export interface LspPosition {
  line: number;
  character: number;
}

export interface LspRange {
  start: LspPosition;
  end: LspPosition;
}

export interface LspDiagnostic {
  range: LspRange;
  severity?: number; // 1=Error, 2=Warning, 3=Info, 4=Hint
  source?: string;
  message: string;
  code?: string | number;
}

export interface SymbolLocation {
  path: string;
  range: LspRange;
}

export interface FileDiagnostics {
  path: string;
  uri: string;
  diagnostics: LspDiagnostic[];
}

export interface WorkspaceDiagnostics {
  files: FileDiagnostics[];
  totalDiagnostics(): number;
  errorCount(): number;
  warningCount(): number;
}

export interface LspContextEnrichment {
  filePath: string;
  diagnostics: WorkspaceDiagnostics;
  definitions: SymbolLocation[];
  references: SymbolLocation[];
}

export interface LspServerConfig {
  name: string;
  command: string;
  args: string[];
  env: Record<string, string>;
  workspaceRoot: string;
  languageIds: Record<string, string>; // extension -> languageId
  initializationOptions?: any;
}

// ─── LSP Client (matches original LspClient) ─────────────────────────────

class LspClient {
  private config: LspServerConfig;
  private process: ChildProcess | null = null;
  private nextRequestId = 1;
  private pendingRequests = new Map<number, {
    resolve: (value: any) => void;
    reject: (error: Error) => void;
  }>();
  private diagnosticsMap = new Map<string, LspDiagnostic[]>();
  private openDocuments = new Map<string, number>(); // path -> version
  private buffer = "";
  private initialized = false;

  constructor(config: LspServerConfig) {
    this.config = config;
  }

  async connect(): Promise<void> {
    this.process = spawn(this.config.command, this.config.args, {
      cwd: this.config.workspaceRoot,
      env: { ...process.env, ...this.config.env },
      stdio: ["pipe", "pipe", "pipe"],
    });

    // Read stdout for JSON-RPC messages
    this.process.stdout?.on("data", (data: Buffer) => {
      this.buffer += data.toString();
      this.processBuffer();
    });

    // Drain stderr to prevent buffer hang (matches original Rust stderr drain)
    // Without this, LSP server can hang when stderr buffer fills up
    this.process.stderr?.on("data", (data: Buffer) => {
      // Actively consume stderr to prevent pipe buffer from filling up.
      // Log at debug level for troubleshooting.
      if (process.env.DEBUG_LSP) {
        process.stderr.write(`[LSP:${this.config.name}:stderr] ${data.toString()}`);
      }
    });
    this.process.stderr?.resume(); // ensure flowing mode even if no listener

    this.process.on("error", (err) => {
      console.error(`LSP server ${this.config.name} error: ${err.message}`);
      this.rejectAllPending(err);
    });

    this.process.on("exit", (code) => {
      if (code !== 0 && code !== null) {
        console.error(`LSP server ${this.config.name} exited with code ${code}`);
      }
      this.rejectAllPending(new Error(`LSP server exited with code ${code}`));
    });

    // Initialize handshake
    await this.initialize();
  }

  private async initialize(): Promise<void> {
    const workspaceUri = `file://${this.config.workspaceRoot}`;

    await this.request("initialize", {
      processId: process.pid,
      rootUri: workspaceUri,
      rootPath: this.config.workspaceRoot,
      workspaceFolders: [{
        uri: workspaceUri,
        name: this.config.name,
      }],
      initializationOptions: this.config.initializationOptions || null,
      capabilities: {
        textDocument: {
          publishDiagnostics: {
            relatedInformation: true,
          },
          definition: {
            linkSupport: true,
          },
          references: {},
        },
        workspace: {
          configuration: false,
          workspaceFolders: true,
        },
        general: {
          positionEncodings: ["utf-16"],
        },
      },
    });

    await this.notify("initialized", {});
    this.initialized = true;
  }

  async ensureDocumentOpen(filePath: string): Promise<void> {
    if (this.openDocuments.has(filePath)) return;
    const contents = fs.readFileSync(filePath, "utf-8");
    await this.openDocument(filePath, contents);
  }

  async openDocument(filePath: string, text: string): Promise<void> {
    const uri = fileUrl(filePath);
    const languageId = this.getLanguageId(filePath);
    if (!languageId) {
      throw new Error(`Unsupported document: ${filePath}`);
    }

    await this.notify("textDocument/didOpen", {
      textDocument: {
        uri,
        languageId,
        version: 1,
        text,
      },
    });

    this.openDocuments.set(filePath, 1);
  }

  async changeDocument(filePath: string, text: string): Promise<void> {
    const uri = fileUrl(filePath);
    const version = (this.openDocuments.get(filePath) || 0) + 1;
    this.openDocuments.set(filePath, version);

    await this.notify("textDocument/didChange", {
      textDocument: { uri, version },
      contentChanges: [{ text }],
    });
  }

  async saveDocument(filePath: string): Promise<void> {
    const uri = fileUrl(filePath);
    await this.notify("textDocument/didSave", {
      textDocument: { uri },
    });
  }

  async closeDocument(filePath: string): Promise<void> {
    const uri = fileUrl(filePath);
    await this.notify("textDocument/didClose", {
      textDocument: { uri },
    });
    this.openDocuments.delete(filePath);
  }

  async goToDefinition(filePath: string, position: LspPosition): Promise<SymbolLocation[]> {
    await this.ensureDocumentOpen(filePath);

    const response = await this.request("textDocument/definition", {
      textDocument: { uri: fileUrl(filePath) },
      position,
    });

    return parseLocationResponse(response);
  }

  async findReferences(
    filePath: string,
    position: LspPosition,
    includeDeclaration = true
  ): Promise<SymbolLocation[]> {
    await this.ensureDocumentOpen(filePath);

    const response = await this.request("textDocument/references", {
      textDocument: { uri: fileUrl(filePath) },
      position,
      context: { includeDeclaration },
    });

    return parseLocationResponse(response);
  }

  getDiagnosticsSnapshot(): Map<string, LspDiagnostic[]> {
    return new Map(this.diagnosticsMap);
  }

  async shutdown(): Promise<void> {
    if (!this.process) return;

    try {
      await this.request("shutdown", {});
      await this.notify("exit", null);
      // Wait for graceful exit before force-killing (matches original Rust behavior)
      await new Promise<void>((resolve) => {
        const timeout = setTimeout(() => resolve(), 3000);
        this.process?.on("exit", () => {
          clearTimeout(timeout);
          resolve();
        });
      });
    } catch {
      // Ignore errors during shutdown
    }

    try {
      if (this.process && !this.process.killed) {
        this.process.kill();
      }
    } catch {}

    this.process = null;
    this.initialized = false;
  }

  // ─── JSON-RPC Protocol ──────────────────────────────────────────────

  private async request(method: string, params: any): Promise<any> {
    return new Promise((resolve, reject) => {
      const id = this.nextRequestId++;
      this.pendingRequests.set(id, { resolve, reject });

      const message = JSON.stringify({
        jsonrpc: "2.0",
        id,
        method,
        params,
      });

      this.sendMessage(message);

      // Timeout after 10s
      setTimeout(() => {
        if (this.pendingRequests.has(id)) {
          this.pendingRequests.delete(id);
          reject(new Error(`LSP request ${method} timed out after 10s`));
        }
      }, 10000);
    });
  }

  private async notify(method: string, params: any): Promise<void> {
    const message = JSON.stringify({
      jsonrpc: "2.0",
      method,
      params,
    });
    this.sendMessage(message);
  }

  private sendMessage(body: string): void {
    if (!this.process?.stdin?.writable) {
      throw new Error("LSP server stdin not available");
    }

    const header = `Content-Length: ${Buffer.byteLength(body)}\r\n\r\n`;
    this.process.stdin.write(header);
    this.process.stdin.write(body);
  }

  private processBuffer(): void {
    while (true) {
      // Find Content-Length header
      const headerEnd = this.buffer.indexOf("\r\n\r\n");
      if (headerEnd === -1) break;

      const headerSection = this.buffer.substring(0, headerEnd);
      const contentLengthMatch = headerSection.match(/Content-Length:\s*(\d+)/i);
      if (!contentLengthMatch) {
        // Skip malformed header
        this.buffer = this.buffer.substring(headerEnd + 4);
        continue;
      }

      const contentLength = parseInt(contentLengthMatch[1], 10);
      const bodyStart = headerEnd + 4;
      const bodyEnd = bodyStart + contentLength;

      if (this.buffer.length < bodyEnd) break; // Not enough data yet

      const body = this.buffer.substring(bodyStart, bodyEnd);
      this.buffer = this.buffer.substring(bodyEnd);

      try {
        const message = JSON.parse(body);
        this.handleMessage(message);
      } catch (err) {
        console.error(`LSP: Failed to parse message: ${err}`);
      }
    }
  }

  private handleMessage(message: any): void {
    // Response to a request
    if (message.id !== undefined && message.id !== null) {
      const pending = this.pendingRequests.get(message.id);
      if (pending) {
        this.pendingRequests.delete(message.id);
        if (message.error) {
          pending.reject(new Error(`LSP error: ${JSON.stringify(message.error)}`));
        } else {
          pending.resolve(message.result);
        }
      }
      return;
    }

    // Notification
    if (message.method === "textDocument/publishDiagnostics") {
      const params = message.params;
      if (params.diagnostics && params.diagnostics.length > 0) {
        this.diagnosticsMap.set(params.uri, params.diagnostics);
      } else {
        this.diagnosticsMap.delete(params.uri);
      }
    }
  }

  private rejectAllPending(error: Error): void {
    for (const [id, pending] of Array.from(this.pendingRequests.entries())) {
      pending.reject(error);
      this.pendingRequests.delete(id);
    }
  }

  private getLanguageId(filePath: string): string | null {
    const ext = path.extname(filePath).replace(".", "").toLowerCase();
    return this.config.languageIds[ext] || null;
  }
}

// ─── LSP Manager (matches original LspManager) ───────────────────────────

export class LspManager {
  private serverConfigs: Map<string, LspServerConfig> = new Map();
  private extensionMap: Map<string, string> = new Map(); // ext -> serverName
  private clients: Map<string, LspClient> = new Map();

  constructor(configs: LspServerConfig[]) {
    for (const config of configs) {
      this.serverConfigs.set(config.name, config);
      for (const ext of Object.keys(config.languageIds)) {
        this.extensionMap.set(ext.toLowerCase(), config.name);
      }
    }
  }

  isSupported(filePath: string): boolean {
    const ext = path.extname(filePath).replace(".", "").toLowerCase();
    return this.extensionMap.has(ext);
  }

  async openDocument(filePath: string, text: string): Promise<void> {
    const client = await this.clientForPath(filePath);
    await client.openDocument(filePath, text);
  }

  async syncDocumentFromDisk(filePath: string): Promise<void> {
    const contents = fs.readFileSync(filePath, "utf-8");
    const client = await this.clientForPath(filePath);
    await client.changeDocument(filePath, contents);
    await client.saveDocument(filePath);
  }

  async changeDocument(filePath: string, text: string): Promise<void> {
    const client = await this.clientForPath(filePath);
    await client.changeDocument(filePath, text);
  }

  async saveDocument(filePath: string): Promise<void> {
    const client = await this.clientForPath(filePath);
    await client.saveDocument(filePath);
  }

  async closeDocument(filePath: string): Promise<void> {
    const client = await this.clientForPath(filePath);
    await client.closeDocument(filePath);
  }

  async goToDefinition(filePath: string, position: LspPosition): Promise<SymbolLocation[]> {
    const client = await this.clientForPath(filePath);
    const locations = await client.goToDefinition(filePath, position);
    return dedupeLocations(locations);
  }

  async findReferences(
    filePath: string,
    position: LspPosition,
    includeDeclaration = true
  ): Promise<SymbolLocation[]> {
    const client = await this.clientForPath(filePath);
    const locations = await client.findReferences(filePath, position, includeDeclaration);
    return dedupeLocations(locations);
  }

  async collectWorkspaceDiagnostics(): Promise<WorkspaceDiagnostics> {
    const files: FileDiagnostics[] = [];

    for (const client of Array.from(this.clients.values())) {
      const snapshot = client.getDiagnosticsSnapshot();
      for (const [uri, diagnostics] of Array.from(snapshot.entries())) {
        if (diagnostics.length === 0) continue;
        const filePath = uriToPath(uri);
        if (!filePath) continue;
        files.push({ path: filePath, uri, diagnostics });
      }
    }

    files.sort((a, b) => a.path.localeCompare(b.path));

    return {
      files,
      totalDiagnostics() {
        return files.reduce((sum, f) => sum + f.diagnostics.length, 0);
      },
      errorCount() {
        return files.reduce(
          (sum, f) => sum + f.diagnostics.filter((d) => d.severity === 1).length,
          0
        );
      },
      warningCount() {
        return files.reduce(
          (sum, f) => sum + f.diagnostics.filter((d) => d.severity === 2).length,
          0
        );
      },
    };
  }

  async contextEnrichment(
    filePath: string,
    position: LspPosition
  ): Promise<LspContextEnrichment> {
    return {
      filePath,
      diagnostics: await this.collectWorkspaceDiagnostics(),
      definitions: await this.goToDefinition(filePath, position),
      references: await this.findReferences(filePath, position, true),
    };
  }

  async shutdown(): Promise<void> {
    for (const client of Array.from(this.clients.values())) {
      await client.shutdown();
    }
    this.clients.clear();
  }

  private async clientForPath(filePath: string): Promise<LspClient> {
    const ext = path.extname(filePath).replace(".", "").toLowerCase();
    const serverName = this.extensionMap.get(ext);
    if (!serverName) {
      throw new Error(`No LSP server configured for .${ext} files`);
    }

    let client = this.clients.get(serverName);
    if (!client) {
      const config = this.serverConfigs.get(serverName);
      if (!config) {
        throw new Error(`LSP server config not found: ${serverName}`);
      }
      client = new LspClient(config);
      await client.connect();
      this.clients.set(serverName, client);
    }

    return client;
  }
}

// ─── Helpers ──────────────────────────────────────────────────────────────

function fileUrl(filePath: string): string {
  return `file://${path.resolve(filePath)}`;
}

function uriToPath(uri: string): string | null {
  try {
    if (uri.startsWith("file://")) {
      return decodeURIComponent(uri.substring(7));
    }
    return null;
  } catch {
    return null;
  }
}

function parseLocationResponse(response: any): SymbolLocation[] {
  if (!response) return [];

  // Single location
  if (response.uri) {
    const p = uriToPath(response.uri);
    if (p) return [{ path: p, range: response.range }];
    return [];
  }

  // Array of locations
  if (Array.isArray(response)) {
    return response
      .map((loc: any) => {
        // Location
        if (loc.uri) {
          const p = uriToPath(loc.uri);
          if (p) return { path: p, range: loc.range };
        }
        // LocationLink
        if (loc.targetUri) {
          const p = uriToPath(loc.targetUri);
          if (p) return { path: p, range: loc.targetSelectionRange || loc.targetRange };
        }
        return null;
      })
      .filter((loc: SymbolLocation | null): loc is SymbolLocation => loc !== null);
  }

  return [];
}

function dedupeLocations(locations: SymbolLocation[]): SymbolLocation[] {
  const seen = new Set<string>();
  return locations.filter((loc) => {
    const key = `${loc.path}:${loc.range.start.line}:${loc.range.start.character}:${loc.range.end.line}:${loc.range.end.character}`;
    if (seen.has(key)) return false;
    seen.add(key);
    return true;
  });
}

// ─── Default TypeScript LSP Config ────────────────────────────────────────

export function createTypeScriptLspConfig(workspaceRoot: string): LspServerConfig {
  return {
    name: "typescript-language-server",
    command: "npx",
    args: ["typescript-language-server", "--stdio"],
    env: {},
    workspaceRoot,
    languageIds: {
      ts: "typescript",
      tsx: "typescriptreact",
      js: "javascript",
      jsx: "javascriptreact",
    },
  };
}

// ─── Singleton ────────────────────────────────────────────────────────────

let lspManagerInstance: LspManager | null = null;

export function getLspManager(workspaceRoot?: string): LspManager {
  if (!lspManagerInstance && workspaceRoot) {
    lspManagerInstance = new LspManager([createTypeScriptLspConfig(workspaceRoot)]);
  }
  if (!lspManagerInstance) {
    lspManagerInstance = new LspManager([
      createTypeScriptLspConfig(process.cwd()),
    ]);
  }
  return lspManagerInstance;
}

export async function initializeLsp(workspaceRoot: string): Promise<LspManager> {
  const manager = getLspManager(workspaceRoot);
  return manager;
}