@oh-my-pi/pi-coding-agent 5.8.0 → 6.1.0

package/src/core/tools/lsp/lspmux.ts

@@ -0,0 +1,249 @@
+import { homedir, platform } from "node:os";
+import { join } from "node:path";
+import { TOML } from "bun";
+import { logger } from "../../logger";
+
+/**
+ * lspmux integration for LSP server multiplexing.
+ *
+ * When lspmux is available and running, this module wraps supported LSP server
+ * commands to use lspmux client mode, enabling server instance sharing across
+ * multiple editor windows.
+ *
+ * Integration is transparent: if lspmux is unavailable, falls back to direct spawning.
+ */
+
+// =============================================================================
+// Types
+// =============================================================================
+
+interface LspmuxConfig {
+	instance_timeout?: number;
+	gc_interval?: number;
+	listen?: [string, number] | string;
+	connect?: [string, number] | string;
+	log_filters?: string;
+	pass_environment?: string[];
+}
+
+interface LspmuxState {
+	available: boolean;
+	running: boolean;
+	binaryPath: string | null;
+	config: LspmuxConfig | null;
+}
+
+// =============================================================================
+// Constants
+// =============================================================================
+
+/**
+ * Servers that benefit from lspmux multiplexing.
+ *
+ * lspmux can multiplex any LSP server, but it's most beneficial for servers
+ * with high startup cost or significant memory usage.
+ */
+const DEFAULT_SUPPORTED_SERVERS = new Set([
+	"rust-analyzer",
+	// Other servers can be added after testing with lspmux
+]);
+
+/** Timeout for liveness check (ms) */
+const LIVENESS_TIMEOUT_MS = 1000;
+
+/** Cache duration for lspmux state (5 minutes) */
+const STATE_CACHE_TTL_MS = 5 * 60 * 1000;
+
+// =============================================================================
+// Config Path
+// =============================================================================
+
+/**
+ * Get the lspmux config path based on platform.
+ * Matches Rust's `dirs::config_dir()` behavior.
+ */
+function getConfigPath(): string {
+	const home = homedir();
+	switch (platform()) {
+		case "win32":
+			return join(process.env.APPDATA ?? join(home, "AppData", "Roaming"), "lspmux", "config.toml");
+		case "darwin":
+			return join(home, "Library", "Application Support", "lspmux", "config.toml");
+		default:
+			return join(process.env.XDG_CONFIG_HOME ?? join(home, ".config"), "lspmux", "config.toml");
+	}
+}
+
+// =============================================================================
+// State Management
+// =============================================================================
+
+let cachedState: LspmuxState | null = null;
+let cacheTimestamp = 0;
+
+/**
+ * Parse lspmux config.toml file.
+ */
+async function parseConfig(): Promise<LspmuxConfig | null> {
+	try {
+		const file = Bun.file(getConfigPath());
+		if (!(await file.exists())) {
+			return null;
+		}
+		return TOML.parse(await file.text()) as LspmuxConfig;
+	} catch {
+		return null;
+	}
+}
+
+/**
+ * Check if lspmux server is running via `lspmux status`.
+ */
+async function checkServerRunning(binaryPath: string): Promise<boolean> {
+	try {
+		const proc = Bun.spawn([binaryPath, "status"], {
+			stdout: "pipe",
+			stderr: "pipe",
+		});
+
+		const exited = await Promise.race([
+			proc.exited,
+			new Promise<null>((resolve) => setTimeout(() => resolve(null), LIVENESS_TIMEOUT_MS)),
+		]);
+
+		if (exited === null) {
+			proc.kill();
+			return false;
+		}
+
+		return exited === 0;
+	} catch {
+		return false;
+	}
+}
+
+/**
+ * Detect lspmux availability and state.
+ * Results are cached for STATE_CACHE_TTL_MS.
+ *
+ * Set OMP_DISABLE_LSPMUX=1 to disable.
+ */
+export async function detectLspmux(): Promise<LspmuxState> {
+	const now = Date.now();
+	if (cachedState && now - cacheTimestamp < STATE_CACHE_TTL_MS) {
+		return cachedState;
+	}
+
+	if (process.env.OMP_DISABLE_LSPMUX === "1") {
+		cachedState = { available: false, running: false, binaryPath: null, config: null };
+		cacheTimestamp = now;
+		return cachedState;
+	}
+
+	const binaryPath = Bun.which("lspmux");
+	if (!binaryPath) {
+		cachedState = { available: false, running: false, binaryPath: null, config: null };
+		cacheTimestamp = now;
+		return cachedState;
+	}
+
+	const [config, running] = await Promise.all([parseConfig(), checkServerRunning(binaryPath)]);
+
+	cachedState = { available: true, running, binaryPath, config };
+	cacheTimestamp = now;
+
+	if (running) {
+		logger.debug("lspmux detected and running", { binaryPath });
+	}
+
+	return cachedState;
+}
+
+/**
+ * Invalidate the cached lspmux state.
+ * Call this if you know the server state has changed.
+ */
+export function invalidateLspmuxCache(): void {
+	cachedState = null;
+	cacheTimestamp = 0;
+}
+
+// =============================================================================
+// Command Wrapping
+// =============================================================================
+
+/**
+ * Check if a server command is supported by lspmux.
+ */
+export function isLspmuxSupported(command: string): boolean {
+	// Extract base command name (handle full paths)
+	const baseName = command.split("/").pop() ?? command;
+	return DEFAULT_SUPPORTED_SERVERS.has(baseName);
+}
+
+export interface LspmuxWrappedCommand {
+	command: string;
+	args: string[];
+	env?: Record<string, string>;
+}
+
+/**
+ * Wrap a server command to use lspmux client mode.
+ *
+ * @param originalCommand - The original LSP server command (e.g., "rust-analyzer")
+ * @param originalArgs - Original command arguments
+ * @param state - lspmux state from detectLspmux()
+ * @returns Wrapped command, args, and env vars; or original if lspmux unavailable
+ */
+export function wrapWithLspmux(
+	originalCommand: string,
+	originalArgs: string[] | undefined,
+	state: LspmuxState,
+): LspmuxWrappedCommand {
+	if (!state.available || !state.running || !state.binaryPath) {
+		return { command: originalCommand, args: originalArgs ?? [] };
+	}
+
+	if (!isLspmuxSupported(originalCommand)) {
+		return { command: originalCommand, args: originalArgs ?? [] };
+	}
+
+	const baseName = originalCommand.split("/").pop() ?? originalCommand;
+	const isDefaultRustAnalyzer = baseName === "rust-analyzer" && originalCommand === "rust-analyzer";
+	const hasArgs = originalArgs && originalArgs.length > 0;
+
+	// rust-analyzer from $PATH with no args - lspmux's default, simplest case
+	if (isDefaultRustAnalyzer && !hasArgs) {
+		return { command: state.binaryPath, args: [] };
+	}
+
+	// Use explicit `client` subcommand with LSPMUX_SERVER env var
+	// Use `--` to separate lspmux options from server args
+	const args = hasArgs ? ["client", "--", ...originalArgs] : ["client"];
+	return {
+		command: state.binaryPath,
+		args,
+		env: { LSPMUX_SERVER: originalCommand },
+	};
+}
+
+/**
+ * Get lspmux-wrapped command if available, otherwise return original.
+ * This is the main entry point for config.ts integration.
+ *
+ * @param command - Original LSP server command
+ * @param args - Original command arguments
+ * @returns Command and args to use (possibly wrapped with lspmux)
+ */
+export async function getLspmuxCommand(command: string, args?: string[]): Promise<LspmuxWrappedCommand> {
+	const state = await detectLspmux();
+	return wrapWithLspmux(command, args, state);
+}
+
+/**
+ * Check if lspmux is currently active and usable.
+ */
+export async function isLspmuxActive(): Promise<boolean> {
+	const state = await detectLspmux();
+	return state.available && state.running;
+}

package/src/core/tools/task/executor.ts

@@ -12,6 +12,8 @@ import type { MCPManager } from "../../mcp/manager";
 import type { ModelRegistry } from "../../model-registry";
 import { checkPythonKernelAvailability } from "../../python-kernel";
 import type { ToolSession } from "..";
+import { createLspTool } from "../lsp/index";
+import type { LspParams } from "../lsp/types";
 import { createPythonTool } from "../python";
 import { ensureArtifactsDir, getArtifactPaths } from "./artifacts";
 import { resolveModelPattern } from "./model-resolver";
@@ -27,6 +29,7 @@ import {
 	TASK_SUBAGENT_PROGRESS_CHANNEL,
 } from "./types";
 import type {
+	LspToolCallRequest,
 	MCPToolCallRequest,
 	MCPToolMetadata,
 	PythonToolCallCancel,
@@ -307,6 +310,9 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 		pythonProxyEnabled = availability.ok;
 	}
 
+	const lspEnabled = enableLsp ?? true;
+	const lspToolRequested = lspEnabled && (toolNames === undefined || toolNames.includes("lsp"));
+
 	let worker: Worker;
 	try {
 		worker = new Worker(new URL("./worker.ts", import.meta.url), { type: "module" });
@@ -385,6 +391,17 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 	const pythonTool = pythonProxyEnabled ? createPythonTool(pythonToolSession) : null;
 	const pythonCallControllers = new Map<string, AbortController>();
 
+	const lspToolSession: ToolSession = {
+		cwd,
+		hasUI: false,
+		enableLsp: lspEnabled,
+		getSessionFile: () => pythonSessionFile,
+		getSessionSpawns: () => spawnsEnv,
+		settings: options.settingsManager as ToolSession["settings"],
+		settingsManager: options.settingsManager,
+	};
+	const lspTool = lspToolRequested ? createLspTool(lspToolSession) : null;
+
 	// Accumulate usage incrementally from message_end events (no memory for streaming events)
 	const accumulatedUsage = {
 		input: 0,
@@ -678,12 +695,13 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 			outputSchema,
 			sessionFile,
 			spawnsEnv,
-			enableLsp,
+			enableLsp: lspEnabled,
 			serializedAuth: options.authStorage?.serialize(),
 			serializedModels: options.modelRegistry?.serialize(),
 			serializedSettings,
 			mcpTools: options.mcpManager ? extractMCPToolMetadata(options.mcpManager) : undefined,
 			pythonToolProxy: pythonProxyEnabled,
+			lspToolProxy: Boolean(lspTool),
 		},
 	};
 
@@ -792,6 +810,40 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 			}
 		};
 
+		const handleLspCall = async (request: LspToolCallRequest) => {
+			if (!lspTool) {
+				worker.postMessage({
+					type: "lsp_tool_result",
+					callId: request.callId,
+					error: "LSP proxy not available",
+				});
+				return;
+			}
+			try {
+				const result = await withTimeout(
+					lspTool.execute(request.callId, request.params as LspParams, signal),
+					request.timeoutMs,
+				);
+				worker.postMessage({
+					type: "lsp_tool_result",
+					callId: request.callId,
+					result: { content: result.content ?? [], details: result.details },
+				});
+			} catch (error) {
+				const message =
+					request.timeoutMs !== undefined && error instanceof Error && error.message.includes("timed out")
+						? `LSP tool call timed out after ${request.timeoutMs}ms`
+						: error instanceof Error
+							? error.message
+							: String(error);
+				worker.postMessage({
+					type: "lsp_tool_result",
+					callId: request.callId,
+					error: message,
+				});
+			}
+		};
+
 		const onMessage = (event: WorkerMessageEvent<SubagentWorkerResponse>) => {
 			const message = event.data;
 			if (!message || resolved) return;
@@ -807,6 +859,10 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
 				handlePythonCancel(message as PythonToolCallCancel);
 				return;
 			}
+			if (message.type === "lsp_tool_call") {
+				handleLspCall(message as LspToolCallRequest);
+				return;
+			}
 			if (message.type === "event") {
 				try {
 					processEvent(message.event);

package/src/core/tools/task/worker-protocol.ts

@@ -64,6 +64,24 @@ export interface PythonToolCallCancel {
 	reason?: string;
 }
 
+export interface LspToolCallRequest {
+	type: "lsp_tool_call";
+	callId: string;
+	params: Record<string, unknown>;
+	timeoutMs?: number;
+}
+
+export interface LspToolCallResponse {
+	type: "lsp_tool_result";
+	callId: string;
+	result?: {
+		content: Array<{ type: string; text?: string; [key: string]: unknown }>;
+		details?: unknown;
+		isError?: boolean;
+	};
+	error?: string;
+}
+
 export interface SubagentWorkerStartPayload {
 	cwd: string;
 	task: string;
@@ -80,6 +98,7 @@ export interface SubagentWorkerStartPayload {
 	serializedSettings?: Settings;
 	mcpTools?: MCPToolMetadata[];
 	pythonToolProxy?: boolean;
+	lspToolProxy?: boolean;
 }
 
 export type SubagentWorkerRequest =
@@ -87,11 +106,13 @@ export type SubagentWorkerRequest =
 	| { type: "abort" }
 	| MCPToolCallResponse
 	| PythonToolCallResponse
-	| PythonToolCallCancel;
+	| PythonToolCallCancel
+	| LspToolCallResponse;
 
 export type SubagentWorkerResponse =
 	| { type: "event"; event: AgentEvent }
 	| { type: "done"; exitCode: number; durationMs: number; error?: string; aborted?: boolean }
 	| MCPToolCallRequest
 	| PythonToolCallRequest
-	| PythonToolCallCancel;
+	| PythonToolCallCancel
+	| LspToolCallRequest;

package/src/core/tools/task/worker.ts

@@ -16,18 +16,22 @@
 import type { AgentEvent, ThinkingLevel } from "@oh-my-pi/pi-agent-core";
 import type { Api, Model } from "@oh-my-pi/pi-ai";
 import type { TSchema } from "@sinclair/typebox";
+import lspDescription from "../../../prompts/tools/lsp.md" with { type: "text" };
 import type { AgentSessionEvent } from "../../agent-session";
 import { AuthStorage } from "../../auth-storage";
 import type { CustomTool } from "../../custom-tools/types";
 import { logger } from "../../logger";
 import { ModelRegistry } from "../../model-registry";
 import { parseModelPattern, parseModelString } from "../../model-resolver";
+import { renderPromptTemplate } from "../../prompt-templates";
 import { createAgentSession, discoverAuthStorage, discoverModels } from "../../sdk";
 import { SessionManager } from "../../session-manager";
 import { SettingsManager } from "../../settings-manager";
 import { untilAborted } from "../../utils";
+import { type LspToolDetails, lspSchema } from "../lsp/types";
 import { getPythonToolDescription, type PythonToolDetails, type PythonToolParams, pythonSchema } from "../python";
 import type {
+	LspToolCallResponse,
 	MCPToolCallResponse,
 	MCPToolMetadata,
 	PythonToolCallResponse,
@@ -58,11 +62,19 @@ interface PendingPythonCall {
 	timeoutId?: ReturnType<typeof setTimeout>;
 }
 
+interface PendingLspCall {
+	resolve: (result: LspToolCallResponse["result"]) => void;
+	reject: (error: Error) => void;
+	timeoutId?: ReturnType<typeof setTimeout>;
+}
+
 const pendingMCPCalls = new Map<string, PendingMCPCall>();
 const pendingPythonCalls = new Map<string, PendingPythonCall>();
+const pendingLspCalls = new Map<string, PendingLspCall>();
 const MCP_CALL_TIMEOUT_MS = 60_000;
 let mcpCallIdCounter = 0;
 let pythonCallIdCounter = 0;
+let lspCallIdCounter = 0;
 
 function generateMCPCallId(): string {
 	return `mcp_${Date.now()}_${++mcpCallIdCounter}`;
@@ -72,6 +84,10 @@ function generatePythonCallId(): string {
 	return `python_${Date.now()}_${++pythonCallIdCounter}`;
 }
 
+function generateLspCallId(): string {
+	return `lsp_${Date.now()}_${++lspCallIdCounter}`;
+}
+
 function callMCPToolViaParent(
 	toolName: string,
 	params: Record<string, unknown>,
@@ -193,6 +209,65 @@ function callPythonToolViaParent(
 	});
 }
 
+function callLspToolViaParent(
+	params: Record<string, unknown>,
+	signal?: AbortSignal,
+	timeoutMs?: number,
+): Promise<LspToolCallResponse["result"]> {
+	return new Promise((resolve, reject) => {
+		const callId = generateLspCallId();
+		if (signal?.aborted) {
+			reject(new Error("Aborted"));
+			return;
+		}
+
+		const timeoutId =
+			typeof timeoutMs === "number" && Number.isFinite(timeoutMs)
+				? setTimeout(() => {
+						pendingLspCalls.delete(callId);
+						reject(new Error(`LSP call timed out after ${timeoutMs}ms`));
+					}, timeoutMs)
+				: undefined;
+
+		const cleanup = () => {
+			if (timeoutId) {
+				clearTimeout(timeoutId);
+			}
+			pendingLspCalls.delete(callId);
+		};
+
+		if (typeof signal?.addEventListener === "function") {
+			signal.addEventListener(
+				"abort",
+				() => {
+					cleanup();
+					reject(new Error("Aborted"));
+				},
+				{ once: true },
+			);
+		}
+
+		pendingLspCalls.set(callId, {
+			resolve: (result) => {
+				cleanup();
+				resolve(result ?? { content: [] });
+			},
+			reject: (error) => {
+				cleanup();
+				reject(error);
+			},
+			timeoutId,
+		});
+
+		postMessageSafe({
+			type: "lsp_tool_call",
+			callId,
+			params,
+			timeoutMs,
+		} as SubagentWorkerResponse);
+	});
+}
+
 function handleMCPToolResult(response: MCPToolCallResponse): void {
 	const pending = pendingMCPCalls.get(response.callId);
 	if (!pending) return;
@@ -213,12 +288,24 @@ function handlePythonToolResult(response: PythonToolCallResponse): void {
 	}
 }
 
+function handleLspToolResult(response: LspToolCallResponse): void {
+	const pending = pendingLspCalls.get(response.callId);
+	if (!pending) return;
+	if (response.error) {
+		pending.reject(new Error(response.error));
+	} else {
+		pending.resolve(response.result);
+	}
+}
+
 function rejectPendingCalls(reason: string): void {
 	const error = new Error(reason);
 	const mcpCalls = Array.from(pendingMCPCalls.values());
 	const pythonCalls = Array.from(pendingPythonCalls.values());
+	const lspCalls = Array.from(pendingLspCalls.values());
 	pendingMCPCalls.clear();
 	pendingPythonCalls.clear();
+	pendingLspCalls.clear();
 	for (const pending of mcpCalls) {
 		clearTimeout(pending.timeoutId);
 		pending.reject(error);
@@ -227,6 +314,10 @@ function rejectPendingCalls(reason: string): void {
 		clearTimeout(pending.timeoutId);
 		pending.reject(error);
 	}
+	for (const pending of lspCalls) {
+		clearTimeout(pending.timeoutId);
+		pending.reject(error);
+	}
 }
 
 function createMCPProxyTool(metadata: MCPToolMetadata): CustomTool<TSchema> {
@@ -296,6 +387,40 @@ function createPythonProxyTool(): CustomTool<typeof pythonSchema> {
 	};
 }
 
+function createLspProxyTool(): CustomTool<typeof lspSchema> {
+	return {
+		name: "lsp",
+		label: "LSP",
+		description: renderPromptTemplate(lspDescription),
+		parameters: lspSchema,
+		execute: async (_toolCallId, params, _onUpdate, _ctx, signal) => {
+			try {
+				const result = await callLspToolViaParent(params as Record<string, unknown>, signal);
+				return {
+					content:
+						result?.content?.map((c) =>
+							c.type === "text"
+								? { type: "text" as const, text: c.text ?? "" }
+								: { type: "text" as const, text: JSON.stringify(c) },
+						) ?? [],
+					details: result?.details as LspToolDetails | undefined,
+				};
+			} catch (error) {
+				const { action } = params;
+				return {
+					content: [
+						{
+							type: "text" as const,
+							text: `LSP error: ${error instanceof Error ? error.message : String(error)}`,
+						},
+					],
+					details: { action, success: false } as LspToolDetails,
+				};
+			}
+		},
+	};
+}
+
 interface WorkerMessageEvent<T> {
 	data: T;
 }
@@ -423,12 +548,17 @@ async function runTask(runState: RunState, payload: SubagentWorkerStartPayload):
 			checkAbort();
 		}
 
-		// Create MCP/python proxy tools if provided
+		// Create MCP/python/LSP proxy tools if provided
 		const mcpProxyTools: CustomTool<TSchema>[] = payload.mcpTools?.map(createMCPProxyTool) ?? [];
 		const pythonProxyTools: CustomTool<TSchema>[] = payload.pythonToolProxy
 			? [createPythonProxyTool() as unknown as CustomTool<TSchema>]
 			: [];
-		const proxyTools = [...mcpProxyTools, ...pythonProxyTools];
+		const lspProxyTools: CustomTool<TSchema>[] = payload.lspToolProxy
+			? [createLspProxyTool() as unknown as CustomTool<TSchema>]
+			: [];
+		const proxyTools = [...mcpProxyTools, ...pythonProxyTools, ...lspProxyTools];
+		const enableLsp = payload.enableLsp ?? true;
+		const lspProxyEnabled = payload.lspToolProxy ?? false;
 
 		// Resolve model override (equivalent to CLI's parseModelPattern with --model)
 		const { model, thinkingLevel: modelThinkingLevel } = resolveModelOverride(payload.model, modelRegistry);
@@ -465,7 +595,7 @@ async function runTask(runState: RunState, payload: SubagentWorkerStartPayload):
 			hasUI: false,
 			// Pass spawn restrictions to nested tasks
 			spawns: payload.spawnsEnv,
-			enableLsp: payload.enableLsp ?? true,
+			enableLsp: enableLsp && !lspProxyEnabled,
 			// Disable local MCP discovery if using proxy tools
 			enableMCP: !payload.mcpTools,
 			// Add proxy tools
@@ -703,7 +833,7 @@ self.addEventListener("messageerror", () => {
 	reportFatal("Failed to deserialize parent message");
 });
 
-// Message handler - receives start/abort/mcp_tool_result commands from parent
+// Message handler - receives start/abort/tool_result commands from parent
 globalThis.addEventListener("message", (event: WorkerMessageEvent<SubagentWorkerRequest>) => {
 	const message = event.data;
 	if (!message) return;
@@ -723,6 +853,11 @@ globalThis.addEventListener("message", (event: WorkerMessageEvent<SubagentWorker
 		return;
 	}
 
+	if (message.type === "lsp_tool_result") {
+		handleLspToolResult(message);
+		return;
+	}
+
 	if (message.type === "start") {
 		// Only allow one task per worker
 		if (activeRun) return;