main
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}