flake-update-20260505
  1/**
  2 * LSP Hook Extension for pi-coding-agent
  3 *
  4 * Provides automatic diagnostics feedback (default: agent end).
  5 * Can run after each write/edit or once per agent response.
  6 *
  7 * Usage:
  8 *   pi --extension ./lsp.ts
  9 *
 10 * Or load the directory to get both hook and tool:
 11 *   pi --extension ./lsp/
 12 */
 13
 14import * as path from "node:path";
 15import * as fs from "node:fs";
 16import * as os from "node:os";
 17import { type ExtensionAPI, type ExtensionContext } from "@mariozechner/pi-coding-agent";
 18import { Text } from "@mariozechner/pi-tui";
 19import { type Diagnostic } from "vscode-languageserver-protocol";
 20import { LSP_SERVERS, formatDiagnostic, getOrCreateManager, shutdownManager } from "./lsp-core.js";
 21
 22type HookScope = "session" | "global";
 23type HookMode = "edit_write" | "agent_end" | "disabled";
 24
 25const DIAGNOSTICS_WAIT_MS_DEFAULT = 3000;
 26
 27function diagnosticsWaitMsForFile(filePath: string): number {
 28  const ext = path.extname(filePath).toLowerCase();
 29  if (ext === ".kt" || ext === ".kts") return 30000;
 30  if (ext === ".swift") return 20000;
 31  if (ext === ".rs") return 20000;
 32  return DIAGNOSTICS_WAIT_MS_DEFAULT;
 33}
 34const DIAGNOSTICS_PREVIEW_LINES = 10;
 35const LSP_IDLE_SHUTDOWN_MS = 2 * 60 * 1000;
 36const DIM = "\x1b[2m", GREEN = "\x1b[32m", YELLOW = "\x1b[33m", RESET = "\x1b[0m";
 37const DEFAULT_HOOK_MODE: HookMode = "agent_end";
 38const SETTINGS_NAMESPACE = "lsp";
 39const LSP_CONFIG_ENTRY = "lsp-hook-config";
 40
 41const WARMUP_MAP: Record<string, string> = {
 42  "pubspec.yaml": ".dart",
 43  "package.json": ".ts",
 44  "pyproject.toml": ".py",
 45  "go.mod": ".go",
 46  "Cargo.toml": ".rs",
 47  "settings.gradle": ".kt",
 48  "settings.gradle.kts": ".kt",
 49  "build.gradle": ".kt",
 50  "build.gradle.kts": ".kt",
 51  "pom.xml": ".kt",
 52  "gradlew": ".kt",
 53  "gradle.properties": ".kt",
 54  "Package.swift": ".swift",
 55};
 56
 57const MODE_LABELS: Record<HookMode, string> = {
 58  edit_write: "After each edit/write",
 59  agent_end: "At agent end",
 60  disabled: "Disabled",
 61};
 62
 63function normalizeHookMode(value: unknown): HookMode | undefined {
 64  if (value === "edit_write" || value === "agent_end" || value === "disabled") return value;
 65  if (value === "turn_end") return "agent_end";
 66  return undefined;
 67}
 68
 69interface HookConfigEntry {
 70  scope: HookScope;
 71  hookMode?: HookMode;
 72}
 73
 74export default function (pi: ExtensionAPI) {
 75  type LspActivity = "idle" | "loading" | "working";
 76
 77  let activeClients: Set<string> = new Set();
 78  let statusUpdateFn: ((key: string, text: string | undefined) => void) | null = null;
 79  let hookMode: HookMode = DEFAULT_HOOK_MODE;
 80  let hookScope: HookScope = "global";
 81  let activity: LspActivity = "idle";
 82  let diagnosticsAbort: AbortController | null = null;
 83  let shuttingDown = false;
 84  let idleShutdownTimer: NodeJS.Timeout | null = null;
 85
 86  const touchedFiles: Map<string, boolean> = new Map();
 87  const globalSettingsPath = path.join(os.homedir(), ".pi", "agent", "settings.json");
 88
 89  function readSettingsFile(filePath: string): Record<string, unknown> {
 90    try {
 91      if (!fs.existsSync(filePath)) return {};
 92      const raw = fs.readFileSync(filePath, "utf-8");
 93      const parsed = JSON.parse(raw);
 94      return parsed && typeof parsed === "object" ? parsed as Record<string, unknown> : {};
 95    } catch {
 96      return {};
 97    }
 98  }
 99
100  function getGlobalHookMode(): HookMode | undefined {
101    const settings = readSettingsFile(globalSettingsPath);
102    const lspSettings = settings[SETTINGS_NAMESPACE];
103    const hookValue = (lspSettings as { hookMode?: unknown; hookEnabled?: unknown } | undefined)?.hookMode;
104    const normalized = normalizeHookMode(hookValue);
105    if (normalized) return normalized;
106
107    const legacyEnabled = (lspSettings as { hookEnabled?: unknown } | undefined)?.hookEnabled;
108    if (typeof legacyEnabled === "boolean") return legacyEnabled ? "edit_write" : "disabled";
109    return undefined;
110  }
111
112  function setGlobalHookMode(mode: HookMode): boolean {
113    try {
114      const settings = readSettingsFile(globalSettingsPath);
115      const existing = settings[SETTINGS_NAMESPACE];
116      const nextNamespace = (existing && typeof existing === "object")
117        ? { ...(existing as Record<string, unknown>), hookMode: mode }
118        : { hookMode: mode };
119
120      settings[SETTINGS_NAMESPACE] = nextNamespace;
121      fs.mkdirSync(path.dirname(globalSettingsPath), { recursive: true });
122      fs.writeFileSync(globalSettingsPath, JSON.stringify(settings, null, 2), "utf-8");
123      return true;
124    } catch {
125      return false;
126    }
127  }
128
129  function getLastHookEntry(ctx: ExtensionContext): HookConfigEntry | undefined {
130    const branchEntries = ctx.sessionManager.getBranch();
131    let latest: HookConfigEntry | undefined;
132
133    for (const entry of branchEntries) {
134      if (entry.type === "custom" && entry.customType === LSP_CONFIG_ENTRY) {
135        latest = entry.data as HookConfigEntry | undefined;
136      }
137    }
138
139    return latest;
140  }
141
142  function restoreHookState(ctx: ExtensionContext): void {
143    const entry = getLastHookEntry(ctx);
144    if (entry?.scope === "session") {
145      const normalized = normalizeHookMode(entry.hookMode);
146      if (normalized) {
147        hookMode = normalized;
148        hookScope = "session";
149        return;
150      }
151
152      const legacyEnabled = (entry as { hookEnabled?: unknown }).hookEnabled;
153      if (typeof legacyEnabled === "boolean") {
154        hookMode = legacyEnabled ? "edit_write" : "disabled";
155        hookScope = "session";
156        return;
157      }
158    }
159
160    const globalSetting = getGlobalHookMode();
161    hookMode = globalSetting ?? DEFAULT_HOOK_MODE;
162    hookScope = "global";
163  }
164
165  function persistHookEntry(entry: HookConfigEntry): void {
166    pi.appendEntry<HookConfigEntry>(LSP_CONFIG_ENTRY, entry);
167  }
168
169  function labelForMode(mode: HookMode): string {
170    return MODE_LABELS[mode];
171  }
172
173  function messageContentToText(content: unknown): string {
174    if (typeof content === "string") return content;
175    if (Array.isArray(content)) {
176      return content
177        .map((item) => (item && typeof item === "object" && "type" in item && (item as any).type === "text")
178          ? String((item as any).text ?? "")
179          : "")
180        .filter(Boolean)
181        .join("\n");
182    }
183    return "";
184  }
185
186  function formatDiagnosticsForDisplay(text: string): string {
187    return text
188      .replace(/\n?This file has errors, please fix\n/gi, "\n")
189      .replace(/<\/?file_diagnostics>\n?/gi, "")
190      .replace(/\n{3,}/g, "\n\n")
191      .trim();
192  }
193
194  function setActivity(next: LspActivity): void {
195    activity = next;
196    updateLspStatus();
197  }
198
199  function clearIdleShutdownTimer(): void {
200    if (!idleShutdownTimer) return;
201    clearTimeout(idleShutdownTimer);
202    idleShutdownTimer = null;
203  }
204
205  async function shutdownLspServersForIdle(): Promise<void> {
206    diagnosticsAbort?.abort();
207    diagnosticsAbort = null;
208    setActivity("idle");
209
210    await shutdownManager();
211    activeClients.clear();
212    updateLspStatus();
213  }
214
215  function scheduleIdleShutdown(): void {
216    clearIdleShutdownTimer();
217
218    idleShutdownTimer = setTimeout(() => {
219      idleShutdownTimer = null;
220      if (shuttingDown) return;
221      void shutdownLspServersForIdle();
222    }, LSP_IDLE_SHUTDOWN_MS);
223
224    (idleShutdownTimer as any).unref?.();
225  }
226
227  function updateLspStatus(): void {
228    if (!statusUpdateFn) return;
229
230    const clients = activeClients.size > 0 ? [...activeClients].join(", ") : "";
231    const clientsText = clients ? `${DIM}${clients}${RESET}` : "";
232    const activityHint = activity === "idle" ? "" : `${DIM}•${RESET}`;
233
234    if (hookMode === "disabled") {
235      const text = clientsText
236        ? `${YELLOW}LSP${RESET} ${DIM}(tool)${RESET}: ${clientsText}`
237        : `${YELLOW}LSP${RESET} ${DIM}(tool)${RESET}`;
238      statusUpdateFn("lsp", text);
239      return;
240    }
241
242    let text = `${GREEN}LSP${RESET}`;
243    if (activityHint) text += ` ${activityHint}`;
244    if (clientsText) text += ` ${clientsText}`;
245    statusUpdateFn("lsp", text);
246  }
247
248  function normalizeFilePath(filePath: string, cwd: string): string {
249    return path.isAbsolute(filePath) ? filePath : path.resolve(cwd, filePath);
250  }
251
252  pi.registerMessageRenderer("lsp-diagnostics", (message, options, theme) => {
253    const content = formatDiagnosticsForDisplay(messageContentToText(message.content));
254    if (!content) return new Text("", 0, 0);
255
256    const expanded = options.expanded === true;
257    const lines = content.split("\n");
258    const maxLines = expanded ? lines.length : DIAGNOSTICS_PREVIEW_LINES;
259    const display = lines.slice(0, maxLines);
260    const remaining = lines.length - display.length;
261
262    const styledLines = display.map((line) => {
263      if (line.startsWith("File: ")) return theme.fg("muted", line);
264      return theme.fg("toolOutput", line);
265    });
266
267    if (!expanded && remaining > 0) {
268      styledLines.push(theme.fg("dim", `... (${remaining} more lines)`));
269    }
270
271    return new Text(styledLines.join("\n"), 0, 0);
272  });
273
274  function getServerConfig(filePath: string) {
275    const ext = path.extname(filePath);
276    return LSP_SERVERS.find((s) => s.extensions.includes(ext));
277  }
278
279  function ensureActiveClientForFile(filePath: string, cwd: string): string | undefined {
280    const absPath = normalizeFilePath(filePath, cwd);
281    const cfg = getServerConfig(absPath);
282    if (!cfg) return undefined;
283
284    if (!activeClients.has(cfg.id)) {
285      activeClients.add(cfg.id);
286      updateLspStatus();
287    }
288
289    return absPath;
290  }
291
292  function extractLspFiles(input: Record<string, unknown>): string[] {
293    const files: string[] = [];
294
295    if (typeof input.file === "string") files.push(input.file);
296    if (Array.isArray(input.files)) {
297      for (const item of input.files) {
298        if (typeof item === "string") files.push(item);
299      }
300    }
301
302    return files;
303  }
304
305  function buildDiagnosticsOutput(
306    filePath: string,
307    diagnostics: Diagnostic[],
308    cwd: string,
309    includeFileHeader: boolean,
310  ): { notification: string; errorCount: number; output: string } {
311    const absPath = path.isAbsolute(filePath) ? filePath : path.resolve(cwd, filePath);
312    const relativePath = path.relative(cwd, absPath);
313    const errorCount = diagnostics.filter((e) => e.severity === 1).length;
314
315    const MAX = 5;
316    const lines = diagnostics.slice(0, MAX).map((e) => {
317      const sev = e.severity === 1 ? "ERROR" : "WARN";
318      return `${sev}[${e.range.start.line + 1}] ${e.message.split("\n")[0]}`;
319    });
320
321    let notification = `📋 ${relativePath}\n${lines.join("\n")}`;
322    if (diagnostics.length > MAX) notification += `\n... +${diagnostics.length - MAX} more`;
323
324    const header = includeFileHeader ? `File: ${relativePath}\n` : "";
325    const output = `\n${header}This file has errors, please fix\n<file_diagnostics>\n${diagnostics.map(formatDiagnostic).join("\n")}\n</file_diagnostics>\n`;
326
327    return { notification, errorCount, output };
328  }
329
330  async function collectDiagnostics(
331    filePath: string,
332    ctx: ExtensionContext,
333    includeWarnings: boolean,
334    includeFileHeader: boolean,
335    notify = true,
336  ): Promise<string | undefined> {
337    const manager = getOrCreateManager(ctx.cwd);
338    const absPath = ensureActiveClientForFile(filePath, ctx.cwd);
339    if (!absPath) return undefined;
340
341    try {
342      const result = await manager.touchFileAndWait(absPath, diagnosticsWaitMsForFile(absPath));
343      if (!result.receivedResponse) return undefined;
344
345      const diagnostics = includeWarnings
346        ? result.diagnostics
347        : result.diagnostics.filter((d) => d.severity === 1);
348      if (!diagnostics.length) return undefined;
349
350      const report = buildDiagnosticsOutput(filePath, diagnostics, ctx.cwd, includeFileHeader);
351
352      if (notify) {
353        if (ctx.hasUI) ctx.ui.notify(report.notification, report.errorCount > 0 ? "error" : "warning");
354        else console.error(report.notification);
355      }
356
357      return report.output;
358    } catch {
359      return undefined;
360    }
361  }
362
363  pi.registerCommand("lsp", {
364    description: "LSP settings (auto diagnostics hook)",
365    handler: async (_args, ctx) => {
366      if (!ctx.hasUI) {
367        ctx.ui.notify("LSP settings require UI", "warning");
368        return;
369      }
370
371      const currentMark = " ✓";
372      const modeOptions = ([
373        "edit_write",
374        "agent_end",
375        "disabled",
376      ] as HookMode[]).map((mode) => ({
377        mode,
378        label: mode === hookMode ? `${labelForMode(mode)}${currentMark}` : labelForMode(mode),
379      }));
380
381      const modeChoice = await ctx.ui.select(
382        "LSP auto diagnostics hook mode:",
383        modeOptions.map((option) => option.label),
384      );
385      if (!modeChoice) return;
386
387      const nextMode = modeOptions.find((option) => option.label === modeChoice)?.mode;
388      if (!nextMode) return;
389
390      const scopeOptions = [
391        {
392          scope: "session" as HookScope,
393          label: "Session only",
394        },
395        {
396          scope: "global" as HookScope,
397          label: "Global (all sessions)",
398        },
399      ];
400
401      const scopeChoice = await ctx.ui.select(
402        "Apply LSP auto diagnostics hook setting to:",
403        scopeOptions.map((option) => option.label),
404      );
405      if (!scopeChoice) return;
406
407      const scope = scopeOptions.find((option) => option.label === scopeChoice)?.scope;
408      if (!scope) return;
409      if (scope === "global") {
410        const ok = setGlobalHookMode(nextMode);
411        if (!ok) {
412          ctx.ui.notify("Failed to update global settings", "error");
413          return;
414        }
415      }
416
417      hookMode = nextMode;
418      hookScope = scope;
419      touchedFiles.clear();
420      persistHookEntry({ scope, hookMode: nextMode });
421      updateLspStatus();
422      ctx.ui.notify(`LSP hook: ${labelForMode(hookMode)} (${hookScope})`, "info");
423    },
424  });
425
426  pi.on("session_start", async (_event, ctx) => {
427    restoreHookState(ctx);
428    statusUpdateFn = ctx.hasUI && ctx.ui.setStatus ? ctx.ui.setStatus.bind(ctx.ui) : null;
429    updateLspStatus();
430
431    if (hookMode === "disabled") return;
432
433    const manager = getOrCreateManager(ctx.cwd);
434
435    for (const [marker, ext] of Object.entries(WARMUP_MAP)) {
436      if (fs.existsSync(path.join(ctx.cwd, marker))) {
437        setActivity("loading");
438        manager.getClientsForFile(path.join(ctx.cwd, `dummy${ext}`))
439          .then((clients) => {
440            if (clients.length > 0) {
441              const cfg = LSP_SERVERS.find((s) => s.extensions.includes(ext));
442              if (cfg) activeClients.add(cfg.id);
443            }
444          })
445          .catch(() => {})
446          .finally(() => setActivity("idle"));
447        break;
448      }
449    }
450  });
451
452  pi.on("session_switch", async (_event, ctx) => {
453    restoreHookState(ctx);
454    updateLspStatus();
455  });
456
457  pi.on("session_tree", async (_event, ctx) => {
458    restoreHookState(ctx);
459    updateLspStatus();
460  });
461
462  pi.on("session_fork", async (_event, ctx) => {
463    restoreHookState(ctx);
464    updateLspStatus();
465  });
466
467  pi.on("session_shutdown", async () => {
468    shuttingDown = true;
469    clearIdleShutdownTimer();
470    diagnosticsAbort?.abort();
471    diagnosticsAbort = null;
472    setActivity("idle");
473
474    await shutdownManager();
475    activeClients.clear();
476    statusUpdateFn?.("lsp", undefined);
477  });
478
479  pi.on("tool_call", async (event, ctx) => {
480    const input = (event.input && typeof event.input === "object")
481      ? event.input as Record<string, unknown>
482      : {};
483
484    if (event.toolName === "lsp") {
485      clearIdleShutdownTimer();
486      const files = extractLspFiles(input);
487      for (const file of files) {
488        ensureActiveClientForFile(file, ctx.cwd);
489      }
490      return;
491    }
492
493    if (event.toolName !== "read" && event.toolName !== "write" && event.toolName !== "edit") return;
494
495    clearIdleShutdownTimer();
496    const filePath = typeof input.path === "string" ? input.path : undefined;
497    if (!filePath) return;
498
499    const absPath = ensureActiveClientForFile(filePath, ctx.cwd);
500    if (!absPath) return;
501
502    void getOrCreateManager(ctx.cwd).getClientsForFile(absPath).catch(() => {});
503  });
504
505  pi.on("agent_start", async () => {
506    clearIdleShutdownTimer();
507    diagnosticsAbort?.abort();
508    diagnosticsAbort = null;
509    setActivity("idle");
510    touchedFiles.clear();
511  });
512
513  function agentWasAborted(event: any): boolean {
514    const messages = Array.isArray(event?.messages) ? event.messages : [];
515    return messages.some((m: any) =>
516      m &&
517      typeof m === "object" &&
518      (m as any).role === "assistant" &&
519      (((m as any).stopReason === "aborted") || ((m as any).stopReason === "error"))
520    );
521  }
522
523  pi.on("agent_end", async (event, ctx) => {
524    try {
525      if (hookMode !== "agent_end") return;
526
527      if (agentWasAborted(event)) {
528        // Don't run diagnostics on aborted/error runs.
529        touchedFiles.clear();
530        return;
531      }
532
533      if (touchedFiles.size === 0) return;
534      if (!ctx.isIdle() || ctx.hasPendingMessages()) return;
535
536      const abort = new AbortController();
537      diagnosticsAbort?.abort();
538      diagnosticsAbort = abort;
539
540      setActivity("working");
541
542      const files = Array.from(touchedFiles.entries());
543      touchedFiles.clear();
544
545      try {
546        const outputs: string[] = [];
547        for (const [filePath, includeWarnings] of files) {
548          if (shuttingDown || abort.signal.aborted) return;
549          if (!ctx.isIdle() || ctx.hasPendingMessages()) {
550            abort.abort();
551            return;
552          }
553
554          const output = await collectDiagnostics(filePath, ctx, includeWarnings, true, false);
555          if (abort.signal.aborted) return;
556          if (output) outputs.push(output);
557        }
558
559        if (shuttingDown || abort.signal.aborted) return;
560
561        if (outputs.length) {
562          pi.sendMessage({
563            customType: "lsp-diagnostics",
564            content: outputs.join("\n"),
565            display: true,
566          }, {
567            triggerTurn: true,
568            deliverAs: "followUp",
569          });
570        }
571      } finally {
572        if (diagnosticsAbort === abort) diagnosticsAbort = null;
573        if (!shuttingDown) setActivity("idle");
574      }
575    } finally {
576      if (!shuttingDown) scheduleIdleShutdown();
577    }
578  });
579
580  pi.on("tool_result", async (event, ctx) => {
581    if (event.toolName !== "write" && event.toolName !== "edit") return;
582
583    const filePath = event.input.path as string;
584    if (!filePath) return;
585
586    const absPath = ensureActiveClientForFile(filePath, ctx.cwd);
587    if (!absPath) return;
588
589    if (hookMode === "disabled") return;
590
591    if (hookMode === "agent_end") {
592      const includeWarnings = event.toolName === "write";
593      const existing = touchedFiles.get(absPath) ?? false;
594      touchedFiles.set(absPath, existing || includeWarnings);
595      return;
596    }
597
598    const includeWarnings = event.toolName === "write";
599    const output = await collectDiagnostics(absPath, ctx, includeWarnings, false);
600    if (!output) return;
601
602    return { content: [...event.content, { type: "text" as const, text: output }] as Array<{ type: "text"; text: string }> };
603  });
604}