@oh-my-pi/pi-coding-agent 17.2.4 → 17.2.5

package/src/lsp/client.ts

@@ -4,12 +4,15 @@ import { MessageFramer } from "../jsonrpc/message-framing";
 import { ToolAbortError, throwIfAborted } from "../tools/tool-errors";
 import { applyWorkspaceEdit } from "./edits";
 import { getLspmuxCommand, isLspmuxSupported } from "./lspmux";
+import { connectSharedLspTransport } from "./mux/daemon";
 import type {
 	LspClient,
 	LspJsonRpcId,
 	LspJsonRpcNotification,
 	LspJsonRpcRequest,
 	LspJsonRpcResponse,
+	LspTransport,
+	LspWriteSink,
 	PublishDiagnosticsParams,
 	ServerConfig,
 	WorkspaceEdit,
@@ -34,6 +37,17 @@ let idleTimeoutMs: number | null = null;
 let idleCheckInterval: NodeJS.Timeout | null = null;
 const IDLE_CHECK_INTERVAL_MS = 60 * 1000;
 
+// Broker-shared server mode (one language server per project shared by every
+// omp instance through the LSP mux daemon). Off by default so embedders and
+// tests that drive getOrCreateClient directly never touch the daemon broker;
+// the SDK turns it on from the `lsp.shared` setting at session creation.
+let sharedLspEnabled = false;
+
+/** Enable or disable attaching to broker-shared language servers. */
+export function setSharedLspEnabled(enabled: boolean): void {
+	sharedLspEnabled = enabled;
+}
+
 /**
  * Configure the idle timeout for LSP clients.
  * @param ms - Timeout in milliseconds, or null/undefined to disable
@@ -209,7 +223,7 @@ class LspDrainAbortError extends Error {
 }
 
 async function writeMessage(
-	sink: Bun.FileSink,
+	sink: LspWriteSink,
 	message: LspJsonRpcRequest | LspJsonRpcNotification | LspJsonRpcResponse,
 	signal?: AbortSignal,
 ): Promise<void> {
@@ -742,7 +756,14 @@ export async function getOrCreateClient(
 			? await getLspmuxCommand(baseCommand, baseArgs)
 			: { command: baseCommand, args: baseArgs };
 
-		const proc = ptree.spawn([command, ...args], {
+		// Prefer the broker-shared server unless an external lspmux wrapper is
+		// already multiplexing this command. Any shared-path failure falls back
+		// to a private spawn so LSP never regresses on broker trouble.
+		let proc: LspTransport | null = null;
+		if (sharedLspEnabled && command === baseCommand) {
+			proc = await connectSharedLspTransport({ command, args, cwd, env, signal });
+		}
+		proc ??= ptree.spawn([command, ...args], {
 			cwd,
 			stdin: "pipe",
 			env: env ? { ...Bun.env, ...env } : undefined,

package/src/lsp/index.ts

@@ -47,6 +47,7 @@ import {
 } from "./edits";
 import { resolveFormatOptions } from "./format-options";
 import { detectLspmux } from "./lspmux";
+import { MUX_RESTART_METHOD } from "./mux/protocol";
 import {
 	type CodeAction,
 	type CodeActionContext,
@@ -535,6 +536,13 @@ async function reloadServer(client: LspClient, serverName: string, signal?: Abor
 		// from the registry by identity and awaiting confirmed process exit) so
 		// the next request cold-starts a fresh client. A kill that never confirms
 		// exit is not a restart: surface the teardown failure truthfully.
+		//
+		// On a broker-shared link a per-session teardown only detaches this
+		// process while the wedged server keeps serving everyone else — ask the
+		// mux to kill the shared server first (best-effort; it also severs us).
+		if (client.proc.sharedMux) {
+			await sendNotification(client, MUX_RESTART_METHOD, undefined, AbortSignal.timeout(2_000)).catch(() => {});
+		}
 		if (!(await shutdownClientInstance(client))) {
 			throw new Error(`Failed to restart ${serverName}: server process did not exit after kill`);
 		}

package/src/lsp/mux/daemon.ts

@@ -0,0 +1,348 @@
+/**
+ * Client half of the broker-shared LSP mux.
+ *
+ * `connectSharedLspTransport` ensures the per-project mux daemon is running
+ * under the daemon broker (same lifecycle as the shared Chromium: started on
+ * first use, stopped when the last omp process in the project exits), dials
+ * its socket, performs the `omp/muxConnect` handshake, and returns an
+ * {@link LspTransport} the ordinary LSP client machinery drives exactly like
+ * a locally spawned server. Every failure degrades to `null` so callers fall
+ * back to a process-local spawn.
+ */
+import * as net from "node:net";
+import * as os from "node:os";
+import * as path from "node:path";
+import { logger, ptree } from "@oh-my-pi/pi-utils";
+import { MessageFramer } from "../../jsonrpc/message-framing";
+import { daemonClientForProject } from "../../launch/client";
+import { describeQuietly, stopQuietly, waitReady } from "../../launch/ensure";
+import { daemonRuntimeDir } from "../../launch/paths";
+import { resolveWorkerSpawnCmd, SMOKE_TEST_TIMEOUT_MS, workerEnvFromParent } from "../../subprocess/worker-client";
+import type { LspJsonRpcRequest, LspJsonRpcResponse, LspTransport, LspWriteSink } from "../types";
+import {
+	LSP_MUX_DAEMON_NAME,
+	LSP_MUX_PROJECT_DIR_ENV,
+	LSP_MUX_READY_PATTERN,
+	LSP_MUX_SOCKET_ENV,
+	LSP_MUX_WORKER_ARG,
+	lspMuxEndpoint,
+	MUX_CONNECT_METHOD,
+	MUX_PING_METHOD,
+	MUX_PING_RESULT,
+	type MuxConnectParams,
+	type MuxConnectResult,
+} from "./protocol";
+
+const CONNECT_TIMEOUT_MS = 3_000;
+const HANDSHAKE_TIMEOUT_MS = 10_000;
+const PROBE_TIMEOUT_MS = 1_500;
+const READY_TIMEOUT_MS = 15_000;
+/** probe→describe→start rounds; bounds cross-process start races and wedged-mux replacement. */
+const ENSURE_ATTEMPTS = 3;
+
+/** Dial a mux endpoint (Unix socket or Windows named pipe) with a bounded connect. */
+function connectEndpoint(endpoint: string, timeoutMs: number): Promise<net.Socket> {
+	const { promise, resolve, reject } = Promise.withResolvers<net.Socket>();
+	const socket = net.connect(endpoint);
+	const timer = setTimeout(() => {
+		socket.destroy();
+		reject(new Error(`Timed out connecting to LSP mux at ${endpoint}`));
+	}, timeoutMs);
+	socket.once("connect", () => {
+		clearTimeout(timer);
+		socket.removeListener("error", onError);
+		resolve(socket);
+	});
+	const onError = (error: Error) => {
+		clearTimeout(timer);
+		reject(error);
+	};
+	socket.once("error", onError);
+	return promise;
+}
+
+/**
+ * Send one request on a fresh socket and await its response. Used only
+ * before the transport takes over the socket; leftover bytes after the
+ * response are returned so the transport's stream can replay them.
+ */
+function requestOnSocket(
+	socket: net.Socket,
+	request: LspJsonRpcRequest,
+	timeoutMs: number,
+): Promise<{ response: LspJsonRpcResponse; leftover: Buffer }> {
+	const { promise, resolve, reject } = Promise.withResolvers<{ response: LspJsonRpcResponse; leftover: Buffer }>();
+	const framer = new MessageFramer(Buffer.alloc(0));
+	const timer = setTimeout(() => {
+		cleanup();
+		reject(new Error(`LSP mux ${request.method} timed out`));
+	}, timeoutMs);
+	const onData = (chunk: Buffer) => {
+		framer.push(chunk);
+		for (const text of framer.drain(() => {})) {
+			let message: LspJsonRpcResponse;
+			try {
+				message = JSON.parse(text);
+			} catch (error) {
+				cleanup();
+				reject(error instanceof Error ? error : new Error(String(error)));
+				return;
+			}
+			if (message.id !== request.id) continue;
+			cleanup();
+			if (message.error) reject(new Error(`LSP mux ${request.method} failed: ${message.error.message}`));
+			else resolve({ response: message, leftover: framer.remainder() });
+			return;
+		}
+	};
+	const onClose = () => {
+		cleanup();
+		reject(new Error("LSP mux connection closed during handshake"));
+	};
+	const cleanup = () => {
+		clearTimeout(timer);
+		socket.removeListener("data", onData);
+		socket.removeListener("close", onClose);
+		socket.removeListener("error", onClose);
+	};
+	socket.on("data", onData);
+	socket.once("close", onClose);
+	socket.once("error", onClose);
+	const content = JSON.stringify(request);
+	socket.write(`Content-Length: ${Buffer.byteLength(content, "utf-8")}\r\n\r\n${content}`);
+	return promise;
+}
+
+/**
+ * Wrap a connected, handshaken mux socket as an {@link LspTransport}.
+ * Socket close maps to a clean "process exit" (code 0) so the LSP client's
+ * crash-recovery path treats a dead mux exactly like a dead server.
+ */
+function socketTransport(socket: net.Socket, leftover: Buffer, pid: number | undefined): LspTransport {
+	const exited = Promise.withResolvers<number>();
+	let exitCode: number | null = null;
+	let needDrain = false;
+	socket.on("error", () => {
+		// "close" always follows; swallowing here prevents an unhandled error event.
+	});
+	socket.once("close", () => {
+		exitCode = 0;
+		exited.resolve(0);
+	});
+	const stdout = new ReadableStream<Uint8Array>({
+		start(controller) {
+			if (leftover.byteLength > 0) controller.enqueue(new Uint8Array(leftover));
+			socket.on("data", (chunk: Buffer) =>
+				controller.enqueue(new Uint8Array(chunk.buffer, chunk.byteOffset, chunk.byteLength)),
+			);
+			socket.once("close", () => {
+				try {
+					controller.close();
+				} catch {
+					// already errored/closed
+				}
+			});
+		},
+	});
+	const stdin: LspWriteSink = {
+		write(data) {
+			needDrain = !socket.write(data);
+			return typeof data === "string" ? Buffer.byteLength(data, "utf-8") : data.byteLength;
+		},
+		flush() {
+			if (!needDrain || socket.destroyed) return;
+			return new Promise<void>(resolveDrain => {
+				const onDrain = () => {
+					needDrain = false;
+					socket.removeListener("close", onDrain);
+					resolveDrain();
+				};
+				socket.once("drain", onDrain);
+				socket.once("close", onDrain);
+			});
+		},
+	};
+	return {
+		stdin,
+		stdout,
+		exited: exited.promise,
+		get exitCode() {
+			return exitCode;
+		},
+		pid,
+		sharedMux: true,
+		kill() {
+			socket.destroy();
+		},
+		peekStderr() {
+			return "";
+		},
+	};
+}
+
+let nextHandshakeId = 1;
+
+/** Dial and handshake one server link; throws on any failure. */
+async function dialMuxServer(endpoint: string, params: MuxConnectParams): Promise<LspTransport> {
+	const socket = await connectEndpoint(endpoint, CONNECT_TIMEOUT_MS);
+	socket.setNoDelay(true);
+	try {
+		const { response, leftover } = await requestOnSocket(
+			socket,
+			{ jsonrpc: "2.0", id: `connect:${nextHandshakeId++}`, method: MUX_CONNECT_METHOD, params },
+			HANDSHAKE_TIMEOUT_MS,
+		);
+		const result = response.result as MuxConnectResult;
+		return socketTransport(socket, leftover, result.pid);
+	} catch (error) {
+		socket.destroy();
+		throw error;
+	}
+}
+
+/** True when a mux answers the ping handshake at `endpoint`. */
+async function probeMux(endpoint: string): Promise<boolean> {
+	try {
+		const socket = await connectEndpoint(endpoint, PROBE_TIMEOUT_MS);
+		try {
+			const { response } = await requestOnSocket(
+				socket,
+				{ jsonrpc: "2.0", id: "probe", method: MUX_PING_METHOD, params: null },
+				PROBE_TIMEOUT_MS,
+			);
+			return response.result === MUX_PING_RESULT;
+		} finally {
+			socket.destroy();
+		}
+	} catch {
+		return false;
+	}
+}
+
+/**
+ * Ensure the project's mux daemon is running under the broker and reachable.
+ * Returns its endpoint, or null when the shared path is unavailable.
+ */
+async function ensureLspMuxDaemon(projectDir: string, signal?: AbortSignal): Promise<string | null> {
+	const client = await daemonClientForProject(projectDir);
+	const endpoint = lspMuxEndpoint(client.projectDir, daemonRuntimeDir(client.projectDir));
+	// The broker connection doubles as the presence lease keeping the daemon alive.
+	await client.request({ op: "ping" }, signal);
+	if (await probeMux(endpoint)) return endpoint;
+	const spawn = resolveWorkerSpawnCmd(LSP_MUX_WORKER_ARG);
+	for (let attempt = 0; attempt < ENSURE_ATTEMPTS; attempt++) {
+		signal?.throwIfAborted();
+		// A concurrent start may have won since the last round; adopt it.
+		if (await probeMux(endpoint)) return endpoint;
+		const existing = await describeQuietly(client, LSP_MUX_DAEMON_NAME, "LSP mux", signal);
+		if (existing && existing.state !== "exited" && existing.state !== "failed") {
+			if (existing.readyAt === undefined) {
+				await waitReady(client, LSP_MUX_DAEMON_NAME, "LSP mux", signal, READY_TIMEOUT_MS);
+			}
+			if (await probeMux(endpoint)) return endpoint;
+			// Live record but nothing listening: replace the wedged daemon.
+			await stopQuietly(client, LSP_MUX_DAEMON_NAME, "LSP mux", signal);
+			continue;
+		}
+		try {
+			const started = await client.request(
+				{
+					op: "start",
+					spec: {
+						name: LSP_MUX_DAEMON_NAME,
+						application: spawn.cmd[0]!,
+						args: spawn.cmd.slice(1),
+						env: {
+							[LSP_MUX_SOCKET_ENV]: endpoint,
+							[LSP_MUX_PROJECT_DIR_ENV]: client.projectDir,
+						},
+						cwd: spawn.cwd ?? client.projectDir,
+						pty: false,
+						ready: { log: LSP_MUX_READY_PATTERN, timeoutMs: READY_TIMEOUT_MS },
+						restart: "no",
+						persist: false,
+						detached: false,
+					},
+				},
+				signal,
+			);
+			if (started.op !== "start") continue;
+			if (await probeMux(endpoint)) return endpoint;
+			await stopQuietly(client, LSP_MUX_DAEMON_NAME, "LSP mux", signal);
+		} catch (error) {
+			signal?.throwIfAborted();
+			// Lost a cross-process start race; the next round adopts the winner.
+			logger.debug("LSP mux start contention", {
+				name: LSP_MUX_DAEMON_NAME,
+				error: error instanceof Error ? error.message : String(error),
+			});
+		}
+	}
+	return null;
+}
+
+/**
+ * Open a broker-shared transport for one language server, ensuring the mux
+ * daemon first. Returns null (after a debug log) when the shared path is
+ * unavailable so the caller falls back to a process-local spawn.
+ */
+export async function connectSharedLspTransport(opts: {
+	command: string;
+	args: string[];
+	cwd: string;
+	env?: Record<string, string>;
+	signal?: AbortSignal;
+}): Promise<LspTransport | null> {
+	try {
+		const endpoint = await ensureLspMuxDaemon(opts.cwd, opts.signal);
+		if (!endpoint) return null;
+		return await dialMuxServer(endpoint, {
+			command: opts.command,
+			args: opts.args,
+			cwd: opts.cwd,
+			env: opts.env,
+		});
+	} catch (error) {
+		if (opts.signal?.aborted) throw error;
+		logger.debug("Shared LSP transport unavailable; falling back to local spawn", {
+			command: opts.command,
+			cwd: opts.cwd,
+			error: error instanceof Error ? error.message : String(error),
+		});
+		return null;
+	}
+}
+
+/** Exercise worker-host mux startup and the ping handshake for distribution smoke tests. */
+export async function smokeTestLspMux(): Promise<void> {
+	const endpoint =
+		process.platform === "win32"
+			? `\\\\.\\pipe\\omp-lsp-mux-smoke-${process.pid.toString(16)}`
+			: path.join(os.tmpdir(), `omp-lsp-mux-smoke-${process.pid.toString(36)}.sock`);
+	const spawn = resolveWorkerSpawnCmd(LSP_MUX_WORKER_ARG);
+	const proc = ptree.spawn(spawn.cmd, {
+		cwd: spawn.cwd,
+		env: workerEnvFromParent({
+			[LSP_MUX_SOCKET_ENV]: endpoint,
+			[LSP_MUX_PROJECT_DIR_ENV]: process.cwd(),
+		}),
+	});
+	try {
+		const deadline = Date.now() + SMOKE_TEST_TIMEOUT_MS;
+		let alive = false;
+		while (Date.now() < deadline) {
+			if (proc.exitCode !== null) break;
+			if (await probeMux(endpoint)) {
+				alive = true;
+				break;
+			}
+			await Bun.sleep(200);
+		}
+		if (!alive) {
+			throw new Error(`lsp mux smoke failed: no ping response (${proc.peekStderr().slice(-500) || "no stderr"})`);
+		}
+	} finally {
+		proc.kill();
+		await proc.exited.catch(() => {});
+	}
+}

package/src/lsp/mux/protocol.ts

@@ -0,0 +1,96 @@
+/**
+ * Cross-process contract for the broker-owned LSP mux daemon.
+ *
+ * One mux daemon runs per project scope (launched through the same daemon
+ * broker that owns the shared Chromium and `hub start` processes). It spawns
+ * each language server once and multiplexes every omp instance in the project
+ * onto that single server over a local socket. The link speaks plain
+ * Content-Length-framed LSP JSON-RPC after a one-request handshake
+ * ({@link MUX_CONNECT_METHOD}); everything below is shared by the worker
+ * entry (`server.ts`), the client connector (`daemon.ts`), and tests.
+ */
+import * as path from "node:path";
+
+/** Hidden CLI selector used to re-enter the LSP mux worker. */
+export const LSP_MUX_WORKER_ARG = "__omp_worker_lsp_mux";
+
+/** Environment key carrying the socket endpoint the mux must listen on. */
+export const LSP_MUX_SOCKET_ENV = "OMP_LSP_MUX_SOCKET";
+
+/** Environment key carrying the canonical project directory the mux serves. */
+export const LSP_MUX_PROJECT_DIR_ENV = "OMP_LSP_MUX_PROJECT_DIR";
+
+/** Stable broker daemon name for the shared LSP mux. */
+export const LSP_MUX_DAEMON_NAME = "omp.lsp.mux";
+
+/** Broker readiness regex matched against the banner printed by the worker. */
+export const LSP_MUX_READY_PATTERN = String.raw`omp lsp mux listening on \S+`;
+
+/** Banner printed on stdout once the mux socket accepts connections. */
+export function lspMuxReadyBanner(endpoint: string): string {
+	return `omp lsp mux listening on ${endpoint}`;
+}
+
+/** Resolve the Unix socket or Windows named pipe for one project scope. */
+export function lspMuxEndpoint(projectDir: string, runtimeDir: string): string {
+	if (process.platform === "win32") {
+		const key = Bun.hash.wyhash(path.resolve(projectDir)).toString(16).padStart(16, "0");
+		return `\\\\.\\pipe\\omp-lsp-mux-${key}`;
+	}
+	return path.join(runtimeDir, "lsp-mux.sock");
+}
+
+// =============================================================================
+// Mux-private JSON-RPC methods
+// =============================================================================
+
+/**
+ * First (and only pre-LSP) request on a fresh connection: binds the link to
+ * one shared server instance, spawning it on first use. Params:
+ * {@link MuxConnectParams}, result: {@link MuxConnectResult}. After the
+ * response the link carries ordinary LSP traffic for that server.
+ */
+export const MUX_CONNECT_METHOD = "omp/muxConnect";
+
+/**
+ * Liveness probe answered with {@link MUX_PING_RESULT} without binding the
+ * link to a server. Used by the smoke probe and the ensure loop.
+ */
+export const MUX_PING_METHOD = "omp/muxPing";
+export const MUX_PING_RESULT = "pong";
+
+/**
+ * Notification on a bound link: kill the shared server process (all sessions
+ * on it are disconnected). Sent by `lsp reload` when the generic reload path
+ * decides the server is wedged — a plain per-session `shutdown`/`exit` is
+ * intercepted by the mux and would leave the wedged server running for
+ * every other instance.
+ */
+export const MUX_RESTART_METHOD = "omp/muxRestartServer";
+
+/** Handshake parameters identifying (and if needed spawning) a shared server. */
+export interface MuxConnectParams {
+	/** Executable to spawn (the client's `resolvedCommand ?? command`). */
+	command: string;
+	/** Arguments passed verbatim to the executable. */
+	args: string[];
+	/** Workspace root the server runs in; part of the server identity key. */
+	cwd: string;
+	/** Extra environment overlaid on the mux daemon's own env for the spawn. */
+	env?: Record<string, string>;
+}
+
+/** Handshake result. */
+export interface MuxConnectResult {
+	/** Server identity key inside the mux (`${command}:${cwd}`). */
+	key: string;
+	/** True when this handshake spawned the server process. */
+	spawned: boolean;
+	/** Pid of the shared server process. */
+	pid?: number;
+}
+
+/** Server identity key used by the mux registry. */
+export function muxServerKey(command: string, cwd: string): string {
+	return `${command}:${cwd}`;
+}