@oh-my-pi/pi-coding-agent 13.3.6 → 13.3.8

package/src/lsp/index.ts

@@ -9,6 +9,7 @@ import lspDescription from "../prompts/tools/lsp.md" with { type: "text" };
 import type { ToolSession } from "../tools";
 import { resolveToCwd } from "../tools/path-utils";
 import { ToolAbortError, throwIfAborted } from "../tools/tool-errors";
+import { clampTimeout } from "../tools/tool-timeouts";
 import {
 	ensureFileOpen,
 	getActiveClients,
@@ -27,6 +28,9 @@ import { applyTextEditsToString, applyWorkspaceEdit } from "./edits";
 import { detectLspmux } from "./lspmux";
 import { renderCall, renderResult } from "./render";
 import {
+	type CodeAction,
+	type CodeActionContext,
+	type Command,
 	type Diagnostic,
 	type DocumentSymbol,
 	type Hover,
@@ -42,16 +46,25 @@ import {
 	type WorkspaceEdit,
 } from "./types";
 import {
+	applyCodeAction,
+	collectGlobMatches,
+	dedupeWorkspaceSymbols,
 	extractHoverText,
 	fileToUri,
+	filterWorkspaceSymbols,
+	formatCodeAction,
 	formatDiagnostic,
 	formatDiagnosticsSummary,
 	formatDocumentSymbol,
 	formatLocation,
 	formatSymbolInformation,
 	formatWorkspaceEdit,
+	hasGlobPattern,
+	readLocationContext,
+	resolveSymbolColumn,
 	sortDiagnostics,
 	symbolKindToIcon,
+	uriToFile,
 } from "./utils";
 
 export type { LspServerStatus } from "./client";
@@ -238,6 +251,10 @@ function getLspServerForFile(config: LspConfig, filePath: string): [string, Serv
 }
 
 const DIAGNOSTIC_MESSAGE_LIMIT = 50;
+const SINGLE_DIAGNOSTICS_WAIT_TIMEOUT_MS = 3000;
+const BATCH_DIAGNOSTICS_WAIT_TIMEOUT_MS = 400;
+const MAX_GLOB_DIAGNOSTIC_TARGETS = 20;
+const WORKSPACE_SYMBOL_LIMIT = 200;
 
 function limitDiagnosticMessages(messages: string[]): string[] {
 	if (messages.length <= DIAGNOSTIC_MESSAGE_LIMIT) {
@@ -246,15 +263,52 @@ function limitDiagnosticMessages(messages: string[]): string[] {
 	return messages.slice(0, DIAGNOSTIC_MESSAGE_LIMIT);
 }
 
-function getServerForWorkspaceAction(config: LspConfig, action: string): [string, ServerConfig] | null {
-	const entries = getLspServers(config);
-	if (entries.length === 0) return null;
+const LOCATION_CONTEXT_LINES = 1;
+const REFERENCE_CONTEXT_LIMIT = 50;
 
-	if (action === "symbols" || action === "reload") {
-		return entries[0];
-	}
+function normalizeLocationResult(result: Location | Location[] | LocationLink | LocationLink[] | null): Location[] {
+	if (!result) return [];
+	const raw = Array.isArray(result) ? result : [result];
+	return raw.flatMap(loc => {
+		if ("uri" in loc) {
+			return [loc as Location];
+		}
+		if ("targetUri" in loc) {
+			const link = loc as LocationLink;
+			return [{ uri: link.targetUri, range: link.targetSelectionRange ?? link.targetRange }];
+		}
+		return [];
+	});
+}
 
-	return null;
+async function formatLocationWithContext(location: Location, cwd: string): Promise<string> {
+	const header = `  ${formatLocation(location, cwd)}`;
+	const context = await readLocationContext(
+		uriToFile(location.uri),
+		location.range.start.line + 1,
+		LOCATION_CONTEXT_LINES,
+	);
+	if (context.length === 0) {
+		return header;
+	}
+	return `${header}\n${context.map(lineText => `    ${lineText}`).join("\n")}`;
+}
+async function reloadServer(client: LspClient, serverName: string, signal?: AbortSignal): Promise<string> {
+	let output = `Restarted ${serverName}`;
+	const reloadMethods = ["rust-analyzer/reloadWorkspace", "workspace/didChangeConfiguration"];
+	for (const method of reloadMethods) {
+		try {
+			await sendRequest(client, method, method.includes("Configuration") ? { settings: {} } : null, signal);
+			output = `Reloaded ${serverName}`;
+			break;
+		} catch {
+			// Method not supported, try next
+		}
+	}
+	if (output.startsWith("Restarted")) {
+		client.proc.kill();
+	}
+	return output;
 }
 
 async function waitForDiagnostics(
@@ -887,7 +941,10 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 		_onUpdate?: AgentToolUpdateCallback<LspToolDetails>,
 		_context?: AgentToolContext,
 	): Promise<AgentToolResult<LspToolDetails>> {
-		const { action, file, files, line, column, query, new_name, apply, include_declaration } = params;
+		const { action, file, line, symbol, occurrence, query, new_name, apply, timeout } = params;
+		const timeoutSec = clampTimeout("lsp", timeout);
+		const timeoutSignal = AbortSignal.timeout(timeoutSec * 1000);
+		signal = signal ? AbortSignal.any([signal, timeoutSignal]) : timeoutSignal;
 		throwIfAborted(signal);
 
 		const config = getConfig(this.session.cwd);
@@ -916,8 +973,7 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 
 		// Diagnostics can be batch or single-file - queries all applicable servers
 		if (action === "diagnostics") {
-			const targets = files?.length ? files : file ? [file] : null;
-			if (!targets) {
+			if (!file) {
 				// No file specified - run workspace diagnostics
 				const result = await runWorkspaceDiagnostics(this.session.cwd, signal);
 				return {
@@ -931,9 +987,34 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 				};
 			}
 
-			const detailed = Boolean(files?.length);
+			let targets: string[];
+			let truncatedGlobTargets = false;
+			if (hasGlobPattern(file)) {
+				const globMatches = await collectGlobMatches(file, this.session.cwd, MAX_GLOB_DIAGNOSTIC_TARGETS);
+				targets = globMatches.matches;
+				truncatedGlobTargets = globMatches.truncated;
+			} else {
+				targets = [file];
+			}
+
+			if (targets.length === 0) {
+				return {
+					content: [{ type: "text", text: `No files matched pattern: ${file}` }],
+					details: { action, success: true, request: params },
+				};
+			}
+
+			const detailed = targets.length > 1 || truncatedGlobTargets;
+			const diagnosticsWaitTimeoutMs = detailed
+				? Math.min(BATCH_DIAGNOSTICS_WAIT_TIMEOUT_MS, timeoutSec * 1000)
+				: Math.min(SINGLE_DIAGNOSTICS_WAIT_TIMEOUT_MS, timeoutSec * 1000);
 			const results: string[] = [];
 			const allServerNames = new Set<string>();
+			if (truncatedGlobTargets) {
+				results.push(
+					`${theme.status.warning} Pattern matched more than ${MAX_GLOB_DIAGNOSTIC_TARGETS} files; showing first ${MAX_GLOB_DIAGNOSTIC_TARGETS}. Narrow the glob or use workspace diagnostics.`,
+				);
+			}
 
 			for (const target of targets) {
 				throwIfAborted(signal);
@@ -962,7 +1043,13 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 						const client = await getOrCreateClient(serverConfig, this.session.cwd);
 						const minVersion = client.diagnosticsVersion;
 						await refreshFile(client, resolved, signal);
-						const diagnostics = await waitForDiagnostics(client, uri, 3000, signal, minVersion);
+						const diagnostics = await waitForDiagnostics(
+							client,
+							uri,
+							diagnosticsWaitTimeoutMs,
+							signal,
+							minVersion,
+						);
 						allDiagnostics.push(...diagnostics);
 					} catch (err) {
 						if (err instanceof ToolAbortError || signal?.aborted) {
@@ -1029,10 +1116,107 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 		}
 
 		const resolvedFile = file ? resolveToCwd(file, this.session.cwd) : null;
-		const serverInfo = resolvedFile
-			? getLspServerForFile(config, resolvedFile)
-			: getServerForWorkspaceAction(config, action);
+		if (action === "symbols" && !resolvedFile) {
+			const normalizedQuery = query?.trim();
+			if (!normalizedQuery) {
+				return {
+					content: [{ type: "text", text: "Error: query parameter required for workspace symbol search" }],
+					details: { action, success: false, request: params },
+				};
+			}
+			const servers = getLspServers(config);
+			if (servers.length === 0) {
+				return {
+					content: [{ type: "text", text: "No language server found for this action" }],
+					details: { action, success: false, request: params },
+				};
+			}
+			const aggregatedSymbols: SymbolInformation[] = [];
+			const respondingServers = new Set<string>();
+			for (const [workspaceServerName, workspaceServerConfig] of servers) {
+				throwIfAborted(signal);
+				try {
+					const workspaceClient = await getOrCreateClient(workspaceServerConfig, this.session.cwd);
+					const workspaceResult = (await sendRequest(
+						workspaceClient,
+						"workspace/symbol",
+						{ query: normalizedQuery },
+						signal,
+					)) as SymbolInformation[] | null;
+					if (!workspaceResult || workspaceResult.length === 0) {
+						continue;
+					}
+					respondingServers.add(workspaceServerName);
+					aggregatedSymbols.push(...filterWorkspaceSymbols(workspaceResult, normalizedQuery));
+				} catch (err) {
+					if (err instanceof ToolAbortError || signal?.aborted) {
+						throw err;
+					}
+				}
+			}
+			const dedupedSymbols = dedupeWorkspaceSymbols(aggregatedSymbols);
+			if (dedupedSymbols.length === 0) {
+				return {
+					content: [{ type: "text", text: `No symbols matching "${normalizedQuery}"` }],
+					details: {
+						action,
+						serverName: Array.from(respondingServers).join(", "),
+						success: true,
+						request: params,
+					},
+				};
+			}
+			const limitedSymbols = dedupedSymbols.slice(0, WORKSPACE_SYMBOL_LIMIT);
+			const lines = limitedSymbols.map(s => formatSymbolInformation(s, this.session.cwd));
+			const truncationLine =
+				dedupedSymbols.length > WORKSPACE_SYMBOL_LIMIT
+					? `\n... ${dedupedSymbols.length - WORKSPACE_SYMBOL_LIMIT} additional symbol(s) omitted`
+					: "";
+			return {
+				content: [
+					{
+						type: "text",
+						text: `Found ${dedupedSymbols.length} symbol(s) matching "${normalizedQuery}":\n${lines.map(l => `  ${l}`).join("\n")}${truncationLine}`,
+					},
+				],
+				details: {
+					action,
+					serverName: Array.from(respondingServers).join(", "),
+					success: true,
+					request: params,
+				},
+			};
+		}
 
+		if (action === "reload" && !resolvedFile) {
+			const servers = getLspServers(config);
+			if (servers.length === 0) {
+				return {
+					content: [{ type: "text", text: "No language server found for this action" }],
+					details: { action, success: false, request: params },
+				};
+			}
+			const outputs: string[] = [];
+			for (const [workspaceServerName, workspaceServerConfig] of servers) {
+				throwIfAborted(signal);
+				try {
+					const workspaceClient = await getOrCreateClient(workspaceServerConfig, this.session.cwd);
+					outputs.push(await reloadServer(workspaceClient, workspaceServerName, signal));
+				} catch (err) {
+					if (err instanceof ToolAbortError || signal?.aborted) {
+						throw err;
+					}
+					const errorMessage = err instanceof Error ? err.message : String(err);
+					outputs.push(`Failed to reload ${workspaceServerName}: ${errorMessage}`);
+				}
+			}
+			return {
+				content: [{ type: "text", text: outputs.join("\n") }],
+				details: { action, serverName: servers.map(([name]) => name).join(", "), success: true, request: params },
+			};
+		}
+
+		const serverInfo = resolvedFile ? getLspServerForFile(config, resolvedFile) : null;
 		if (!serverInfo) {
 			return {
 				content: [{ type: "text", text: "No language server found for this action" }],
@@ -1051,7 +1235,11 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 			}
 
 			const uri = targetFile ? fileToUri(targetFile) : "";
-			const position = { line: (line || 1) - 1, character: (column || 1) - 1 };
+			const resolvedLine = line ?? 1;
+			const resolvedCharacter = targetFile
+				? await resolveSymbolColumn(targetFile, resolvedLine, symbol, occurrence)
+				: 0;
+			const position = { line: resolvedLine - 1, character: resolvedCharacter };
 
 			let output: string;
 
@@ -1071,33 +1259,66 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 						signal,
 					)) as Location | Location[] | LocationLink | LocationLink[] | null;
 
-					if (!result) {
+					const locations = normalizeLocationResult(result);
+
+					if (locations.length === 0) {
 						output = "No definition found";
 					} else {
-						const raw = Array.isArray(result) ? result : [result];
-						const locations = raw.flatMap(loc => {
-							if ("uri" in loc) {
-								return [loc as Location];
-							}
-							if ("targetUri" in loc) {
-								// Use targetSelectionRange (the precise identifier range) with fallback to targetRange
-								const link = loc as LocationLink;
-								return [{ uri: link.targetUri, range: link.targetSelectionRange ?? link.targetRange }];
-							}
-							return [];
-						});
+						const lines = await Promise.all(
+							locations.map(location => formatLocationWithContext(location, this.session.cwd)),
+						);
+						output = `Found ${locations.length} definition(s):\n${lines.join("\n")}`;
+					}
+					break;
+				}
 
-						if (locations.length === 0) {
-							output = "No definition found";
-						} else {
-							output = `Found ${locations.length} definition(s):\n${locations
-								.map(loc => `  ${formatLocation(loc, this.session.cwd)}`)
-								.join("\n")}`;
-						}
+				case "type_definition": {
+					const result = (await sendRequest(
+						client,
+						"textDocument/typeDefinition",
+						{
+							textDocument: { uri },
+							position,
+						},
+						signal,
+					)) as Location | Location[] | LocationLink | LocationLink[] | null;
+
+					const locations = normalizeLocationResult(result);
+
+					if (locations.length === 0) {
+						output = "No type definition found";
+					} else {
+						const lines = await Promise.all(
+							locations.map(location => formatLocationWithContext(location, this.session.cwd)),
+						);
+						output = `Found ${locations.length} type definition(s):\n${lines.join("\n")}`;
 					}
 					break;
 				}
 
+				case "implementation": {
+					const result = (await sendRequest(
+						client,
+						"textDocument/implementation",
+						{
+							textDocument: { uri },
+							position,
+						},
+						signal,
+					)) as Location | Location[] | LocationLink | LocationLink[] | null;
+
+					const locations = normalizeLocationResult(result);
+
+					if (locations.length === 0) {
+						output = "No implementation found";
+					} else {
+						const lines = await Promise.all(
+							locations.map(location => formatLocationWithContext(location, this.session.cwd)),
+						);
+						output = `Found ${locations.length} implementation(s):\n${lines.join("\n")}`;
+					}
+					break;
+				}
 				case "references": {
 					const result = (await sendRequest(
 						client,
@@ -1105,7 +1326,7 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 						{
 							textDocument: { uri },
 							position,
-							context: { includeDeclaration: include_declaration ?? true },
+							context: { includeDeclaration: true },
 						},
 						signal,
 					)) as Location[] | null;
@@ -1113,7 +1334,19 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 					if (!result || result.length === 0) {
 						output = "No references found";
 					} else {
-						const lines = result.map(loc => `  ${formatLocation(loc, this.session.cwd)}`);
+						const contextualReferences = result.slice(0, REFERENCE_CONTEXT_LIMIT);
+						const plainReferences = result.slice(REFERENCE_CONTEXT_LIMIT);
+						const contextualLines = await Promise.all(
+							contextualReferences.map(location => formatLocationWithContext(location, this.session.cwd)),
+						);
+						const plainLines = plainReferences.map(location => `  ${formatLocation(location, this.session.cwd)}`);
+						const lines = plainLines.length
+							? [
+									...contextualLines,
+									`  ... ${plainLines.length} additional reference(s) shown without context`,
+									...plainLines,
+								]
+							: contextualLines;
 						output = `Found ${result.length} reference(s):\n${lines.join("\n")}`;
 					}
 					break;
@@ -1138,52 +1371,118 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 					break;
 				}
 
-				case "symbols": {
-					// If no file, do workspace symbol search (requires query)
-					if (!targetFile) {
-						if (!query) {
-							return {
-								content: [
-									{ type: "text", text: "Error: query parameter required for workspace symbol search" },
-								],
-								details: { action, serverName, success: false },
-							};
+				case "code_actions": {
+					const diagnostics = client.diagnostics.get(uri) ?? [];
+					const context: CodeActionContext = {
+						diagnostics,
+						only: !apply && query ? [query] : undefined,
+						triggerKind: 1,
+					};
+
+					const result = (await sendRequest(
+						client,
+						"textDocument/codeAction",
+						{
+							textDocument: { uri },
+							range: { start: position, end: position },
+							context,
+						},
+						signal,
+					)) as (CodeAction | Command)[] | null;
+
+					if (!result || result.length === 0) {
+						output = "No code actions available";
+						break;
+					}
+
+					if (apply === true && query) {
+						const normalizedQuery = query.trim();
+						if (normalizedQuery.length === 0) {
+							output = "Error: query parameter required when apply=true for code_actions";
+							break;
 						}
-						const result = (await sendRequest(client, "workspace/symbol", { query }, signal)) as
-							| SymbolInformation[]
-							| null;
-						if (!result || result.length === 0) {
-							output = `No symbols matching "${query}"`;
-						} else {
-							const lines = result.map(s => formatSymbolInformation(s, this.session.cwd));
-							output = `Found ${result.length} symbol(s) matching "${query}":\n${lines.map(l => `  ${l}`).join("\n")}`;
+						const parsedIndex = /^\d+$/.test(normalizedQuery) ? Number.parseInt(normalizedQuery, 10) : null;
+						const selectedAction = result.find(
+							(actionItem, index) =>
+								(parsedIndex !== null && index === parsedIndex) ||
+								actionItem.title.toLowerCase().includes(normalizedQuery.toLowerCase()),
+						);
+
+						if (!selectedAction) {
+							const actionLines = result.map((actionItem, index) => `  ${formatCodeAction(actionItem, index)}`);
+							output = `No code action matches "${normalizedQuery}". Available actions:\n${actionLines.join("\n")}`;
+							break;
 						}
-					} else {
-						// File-based document symbols
-						const result = (await sendRequest(
-							client,
-							"textDocument/documentSymbol",
-							{
-								textDocument: { uri },
+
+						const appliedAction = await applyCodeAction(selectedAction, {
+							resolveCodeAction: async actionItem =>
+								(await sendRequest(client, "codeAction/resolve", actionItem, signal)) as CodeAction,
+							applyWorkspaceEdit: async edit => applyWorkspaceEdit(edit, this.session.cwd),
+							executeCommand: async commandItem => {
+								await sendRequest(
+									client,
+									"workspace/executeCommand",
+									{
+										command: commandItem.command,
+										arguments: commandItem.arguments ?? [],
+									},
+									signal,
+								);
 							},
-							signal,
-						)) as (DocumentSymbol | SymbolInformation)[] | null;
+						});
 
-						if (!result || result.length === 0) {
-							output = "No symbols found";
+						if (!appliedAction) {
+							output = `Action "${selectedAction.title}" has no workspace edit or command to apply`;
+							break;
+						}
+
+						const summaryLines: string[] = [];
+						if (appliedAction.edits.length > 0) {
+							summaryLines.push("  Workspace edit:");
+							summaryLines.push(...appliedAction.edits.map(item => `    ${item}`));
+						}
+						if (appliedAction.executedCommands.length > 0) {
+							summaryLines.push("  Executed command(s):");
+							summaryLines.push(...appliedAction.executedCommands.map(commandName => `    ${commandName}`));
+						}
+
+						output = `Applied "${appliedAction.title}":\n${summaryLines.join("\n")}`;
+						break;
+					}
+
+					const actionLines = result.map((actionItem, index) => `  ${formatCodeAction(actionItem, index)}`);
+					output = `${result.length} code action(s):\n${actionLines.join("\n")}`;
+					break;
+				}
+				case "symbols": {
+					if (!targetFile) {
+						output = "Error: file parameter required for document symbols";
+						break;
+					}
+					// File-based document symbols
+					const result = (await sendRequest(
+						client,
+						"textDocument/documentSymbol",
+						{
+							textDocument: { uri },
+						},
+						signal,
+					)) as (DocumentSymbol | SymbolInformation)[] | null;
+
+					if (!result || result.length === 0) {
+						output = "No symbols found";
+					} else {
+						const relPath = path.relative(this.session.cwd, targetFile);
+						if ("selectionRange" in result[0]) {
+							const lines = (result as DocumentSymbol[]).flatMap(s => formatDocumentSymbol(s));
+							output = `Symbols in ${relPath}:\n${lines.join("\n")}`;
 						} else {
-							const relPath = path.relative(this.session.cwd, targetFile);
-							if ("selectionRange" in result[0]) {
-								const lines = (result as DocumentSymbol[]).flatMap(s => formatDocumentSymbol(s));
-								output = `Symbols in ${relPath}:\n${lines.join("\n")}`;
-							} else {
-								const lines = (result as SymbolInformation[]).map(s => {
-									const line = s.location.range.start.line + 1;
-									const icon = symbolKindToIcon(s.kind);
-									return `${icon} ${s.name} @ line ${line}`;
-								});
-								output = `Symbols in ${relPath}:\n${lines.join("\n")}`;
-							}
+							const lines = (result as SymbolInformation[]).map(s => {
+								const line = s.location.range.start.line + 1;
+								const icon = symbolKindToIcon(s.kind);
+								return `${icon} ${s.name} @ line ${line}`;
+							});
+							output = `Symbols in ${relPath}:\n${lines.join("\n")}`;
 						}
 					}
 					break;
@@ -1224,26 +1523,7 @@ export class LspTool implements AgentTool<typeof lspSchema, LspToolDetails, Them
 				}
 
 				case "reload": {
-					// Try graceful reload first, fall back to kill
-					output = `Restarted ${serverName}`;
-					const reloadMethods = ["rust-analyzer/reloadWorkspace", "workspace/didChangeConfiguration"];
-					for (const method of reloadMethods) {
-						try {
-							await sendRequest(
-								client,
-								method,
-								method.includes("Configuration") ? { settings: {} } : null,
-								signal,
-							);
-							output = `Reloaded ${serverName}`;
-							break;
-						} catch {
-							// Method not supported, try next
-						}
-					}
-					if (output.startsWith("Restarted")) {
-						client.proc.kill();
-					}
+					output = await reloadServer(client, serverName, signal);
 					break;
 				}