From d1278762f8ff76c4a3488c2f4742aec13347aec6 Mon Sep 17 00:00:00 2001
From: roboomp <omp@can.ac>
Date: Thu, 13 Aug 2026 03:15:29 +0000
Subject: [PATCH] fix(lsp): deferred idle reuse until cleanup

Kept closing sessions attached to their server until orphan document closes and pending request cancellations are queued.

Added coverage that reopens the last orphan URI after reuse and verifies the replacement overlay remains active.

Fixes #8371
---
 packages/coding-agent/src/lsp/mux/server.ts |  6 +++--
 packages/coding-agent/test/lsp-mux.test.ts  | 26 ++++++++++++++-------
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/packages/coding-agent/src/lsp/mux/server.ts b/packages/coding-agent/src/lsp/mux/server.ts
index 6afb08c24..3297e1c36 100644
--- a/packages/coding-agent/src/lsp/mux/server.ts
+++ b/packages/coding-agent/src/lsp/mux/server.ts
@@ -643,21 +643,23 @@ export class LspMuxServer {
 		this.#sessions.delete(session);
 		const server = session.server;
 		if (server) {
-			server.sessions.delete(session);
+			let cleanup: Promise<void> | undefined;
 			for (const uri of session.openUris) {
 				server.documents.delete(uri);
-				await this.#writeServer(server, {
+				cleanup = this.#writeServer(server, {
 					jsonrpc: "2.0",
 					method: "textDocument/didClose",
 					params: { textDocument: { uri } },
 				});
 			}
+			await cleanup;
 			for (const [muxId, pending] of server.pending) {
 				if (pending.session !== session) continue;
 				pending.drop = true;
 				await this.#writeServer(server, { jsonrpc: "2.0", method: "$/cancelRequest", params: { id: muxId } });
 			}
 			server.initializeWaiters.delete(session);
+			server.sessions.delete(session);
 			if (server.sessions.size === 0 && !server.stopping) {
 				server.lingerTimer = setTimeout(() => {
 					if (server.sessions.size === 0) void this.#stopServer(server);
diff --git a/packages/coding-agent/test/lsp-mux.test.ts b/packages/coding-agent/test/lsp-mux.test.ts
index a8dc7438e..90d9c1cc3 100644
--- a/packages/coding-agent/test/lsp-mux.test.ts
+++ b/packages/coding-agent/test/lsp-mux.test.ts
@@ -365,21 +365,31 @@ describe("LspMuxServer", () => {
 	);
 
 	it.skipIf(process.platform === "win32")(
-		"closes orphaned documents before reusing an idle server",
+		"finishes orphan document closes before reusing a server",
 		async () => {
 			const first = await link();
 			await initialize(first.client);
-			const uri = "file:///orphan.ts";
-			first.client.notify("textDocument/didOpen", {
-				textDocument: { uri, languageId: "typescript", version: 1, text: "orphan" },
-			});
-			await pollUntil(async () => (await state(first.client)).didOpen[uri] === 1, "orphan didOpen");
+			const uris = Array.from({ length: 128 }, (_, index) => `file:///orphan-${index}.ts`);
+			for (const uri of uris) {
+				first.client.notify("textDocument/didOpen", {
+					textDocument: { uri, languageId: "typescript", version: 1, text: "orphan" },
+				});
+			}
+			await first.client.request("test/echo", { barrier: true });
+			const firstClosed = first.client.waitForClose();
 			first.client.destroy();
-			await pollUntil(() => Promise.resolve(server.sessionCount === 0), "orphan session close");
+			await firstClosed;
 
 			const second = await link();
 			expect(second.connected.spawned).toBe(false);
-			await pollUntil(async () => (await state(second.client)).didClose.includes(uri), "orphan didClose");
+			const uri = uris.at(-1);
+			expect(uri).toBeDefined();
+			await initialize(second.client);
+			second.client.notify("textDocument/didOpen", {
+				textDocument: { uri, languageId: "typescript", version: 1, text: "replacement" },
+			});
+			await second.client.request("test/echo", { barrier: true });
+			expect(await second.client.request<string | null>("test/documentText", { uri })).toBe("replacement");
 		},
 		10_000,
 	);