Skip to content

Commit f73c554

Browse files
fix(agents): cancel pending bundle LSP requests (#104110)
Co-authored-by: Vincent Koc <[email protected]>
1 parent ba826be commit f73c554

2 files changed

Lines changed: 140 additions & 30 deletions

File tree

src/agents/agent-bundle-lsp-runtime.test.ts

Lines changed: 66 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@ class MockChildProcess extends EventEmitter {
4646
readonly stdout = new PassThrough();
4747
readonly stderr = new PassThrough();
4848
readonly stdin: Writable;
49+
readonly receivedMessages: Record<string, unknown>[] = [];
4950

5051
constructor(
5152
private readonly initializeResponsePrefix = "",
@@ -70,13 +71,26 @@ class MockChildProcess extends EventEmitter {
7071

7172
private respondToRequest(text: string): void {
7273
const body = parseWrittenLspBody(text);
73-
if (!body || typeof body.id !== "number" || typeof body.method !== "string") {
74+
if (!body) {
75+
return;
76+
}
77+
this.receivedMessages.push(body);
78+
if (typeof body.id !== "number" || typeof body.method !== "string") {
7479
return;
7580
}
7681
if (this.respondMethods && !this.respondMethods.has(body.method)) {
7782
return;
7883
}
79-
const result = body.method === "initialize" ? { capabilities: { hoverProvider: true } } : null;
84+
const result =
85+
body.method === "initialize"
86+
? {
87+
capabilities: {
88+
hoverProvider: true,
89+
definitionProvider: true,
90+
referencesProvider: true,
91+
},
92+
}
93+
: null;
8094
queueMicrotask(() => {
8195
this.stdout.write(
8296
`${this.initializeResponsePrefix}${encodeLspMessage({ jsonrpc: "2.0", id: body.id, result })}`,
@@ -240,6 +254,56 @@ describe("bundle LSP runtime", () => {
240254
await runtime.dispose();
241255
});
242256

257+
it.each([
258+
["lsp_hover_typescript", "textDocument/hover"],
259+
["lsp_definition_typescript", "textDocument/definition"],
260+
["lsp_references_typescript", "textDocument/references"],
261+
])("cancels pending %s requests when the tool signal aborts", async (toolName, method) => {
262+
configureSingleLspServer();
263+
const child = new MockChildProcess("", new Set(["initialize"]));
264+
spawnMock.mockReturnValue(child);
265+
const { createBundleLspToolRuntime } = await import("./agent-bundle-lsp-runtime.js");
266+
267+
const runtime = await createBundleLspToolRuntime({ workspaceDir: "/tmp/workspace" });
268+
const tool = runtime.tools.find((candidate) => candidate.name === toolName);
269+
if (!tool) {
270+
throw new Error(`expected ${toolName} tool`);
271+
}
272+
const controller = new AbortController();
273+
const request = tool.execute(
274+
"call-1",
275+
{
276+
uri: "file:///tmp/workspace/index.ts",
277+
line: 0,
278+
character: 0,
279+
},
280+
controller.signal,
281+
);
282+
const settled = request.then(
283+
() => "resolved",
284+
() => "rejected",
285+
);
286+
const lspRequest = child.receivedMessages.find((message) => message.method === method);
287+
288+
controller.abort(new Error("agent stopped"));
289+
290+
await expect(
291+
Promise.race([
292+
settled,
293+
new Promise((resolve) => {
294+
setTimeout(() => resolve("still pending"), 100);
295+
}),
296+
]),
297+
).resolves.toBe("rejected");
298+
expect(child.receivedMessages).toContainEqual({
299+
jsonrpc: "2.0",
300+
method: "$/cancelRequest",
301+
params: { id: lspRequest?.id },
302+
});
303+
304+
await runtime.dispose();
305+
});
306+
243307
it("keeps LSP framing aligned after multibyte messages in the same chunk", async () => {
244308
configureSingleLspServer();
245309
const prefix = encodeLspMessage({

src/agents/agent-bundle-lsp-runtime.ts

Lines changed: 74 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import { spawn, type ChildProcess } from "node:child_process";
33
import { normalizeOptionalLowercaseString } from "@openclaw/normalization-core/string-coerce";
44
import type { OpenClawConfig } from "../config/types.openclaw.js";
5+
import { createAbortError } from "../infra/abort-signal.js";
56
import { sanitizeHostExecEnv } from "../infra/host-env-security.js";
67
import { logDebug, logWarn } from "../logger.js";
78
import {
@@ -39,6 +40,7 @@ type PendingLspRequest = {
3940
resolve: (v: unknown) => void;
4041
reject: (e: Error) => void;
4142
timeout: ReturnType<typeof setTimeout>;
43+
dispose: () => void;
4244
};
4345

4446
type LspServerCapabilities = {
@@ -114,13 +116,22 @@ function rememberLspFailure(session: LspSession, error: Error): void {
114116
session.failure ??= error;
115117
}
116118

119+
function takePendingLspRequest(session: LspSession, id: number): PendingLspRequest | undefined {
120+
const pending = session.pendingRequests.get(id);
121+
if (!pending) {
122+
return undefined;
123+
}
124+
session.pendingRequests.delete(id);
125+
clearTimeout(pending.timeout);
126+
pending.dispose();
127+
return pending;
128+
}
129+
117130
function failLspSession(session: LspSession, error: Error): void {
118131
rememberLspFailure(session, error);
119-
for (const pending of session.pendingRequests.values()) {
120-
clearTimeout(pending.timeout);
121-
pending.reject(session.failure ?? error);
132+
for (const [id] of session.pendingRequests) {
133+
takePendingLspRequest(session, id)?.reject(session.failure ?? error);
122134
}
123-
session.pendingRequests.clear();
124135
}
125136

126137
function lspProcessExitError(
@@ -205,20 +216,49 @@ function parseLspMessages(buffer: Buffer): { messages: unknown[]; remaining: Buf
205216
return { messages, remaining };
206217
}
207218

208-
function sendRequest(session: LspSession, method: string, params?: unknown): Promise<unknown> {
219+
function lspAbortError(signal?: AbortSignal): Error {
220+
return signal?.reason instanceof Error
221+
? signal.reason
222+
: createAbortError("LSP request aborted", { cause: signal?.reason });
223+
}
224+
225+
function sendRequest(
226+
session: LspSession,
227+
method: string,
228+
params?: unknown,
229+
signal?: AbortSignal,
230+
): Promise<unknown> {
209231
if (session.failure) {
210232
return Promise.reject(session.failure);
211233
}
234+
if (signal?.aborted) {
235+
return Promise.reject(lspAbortError(signal));
236+
}
212237
const id = ++session.requestId;
213238
return new Promise((resolve, reject) => {
214239
const timeout = setTimeout(() => {
215-
if (session.pendingRequests.has(id)) {
216-
session.pendingRequests.delete(id);
217-
reject(new Error(`LSP request ${method} timed out`));
218-
}
240+
takePendingLspRequest(session, id)?.reject(new Error(`LSP request ${method} timed out`));
219241
}, 10_000);
220242
timeout.unref?.();
221-
session.pendingRequests.set(id, { resolve, reject, timeout });
243+
const onAbort = () => {
244+
const pending = takePendingLspRequest(session, id);
245+
if (!pending) {
246+
return;
247+
}
248+
// Bundle tools share the server process, so cancel only this request.
249+
try {
250+
session.process.stdin?.write(
251+
encodeLspMessage({ jsonrpc: "2.0", method: "$/cancelRequest", params: { id } }),
252+
"utf-8",
253+
);
254+
} catch {
255+
// Best-effort notification; the local tool promise must still settle.
256+
}
257+
pending.reject(lspAbortError(signal));
258+
};
259+
const dispose = () => signal?.removeEventListener("abort", onAbort);
260+
session.pendingRequests.set(id, { resolve, reject, timeout, dispose });
261+
signal?.addEventListener("abort", onAbort, { once: true });
222262
const message = { jsonrpc: "2.0", id, method, params };
223263
const encoded = encodeLspMessage(message);
224264
session.process.stdin?.write(encoded, "utf-8");
@@ -240,10 +280,8 @@ function handleIncomingData(session: LspSession, chunk: Buffer | string) {
240280
const record = msg as Record<string, unknown>;
241281

242282
if ("id" in record && typeof record.id === "number") {
243-
const pending = session.pendingRequests.get(record.id);
283+
const pending = takePendingLspRequest(session, record.id);
244284
if (pending) {
245-
session.pendingRequests.delete(record.id);
246-
clearTimeout(pending.timeout);
247285
if ("error" in record) {
248286
pending.reject(new Error(JSON.stringify(record.error)));
249287
} else {
@@ -316,11 +354,9 @@ async function disposeSession(session: LspSession) {
316354
// best-effort
317355
}
318356
}
319-
for (const [, pending] of session.pendingRequests) {
320-
clearTimeout(pending.timeout);
321-
pending.reject(new Error("LSP session disposed"));
357+
for (const [id] of session.pendingRequests) {
358+
takePendingLspRequest(session, id)?.reject(new Error("LSP session disposed"));
322359
}
323-
session.pendingRequests.clear();
324360
terminateLspProcessTree(session);
325361
}
326362

@@ -349,12 +385,17 @@ function createLspPositionTool(params: {
349385
},
350386
required: ["uri", "line", "character"],
351387
},
352-
execute: async (_toolCallId, input) => {
388+
execute: async (_toolCallId, input, signal) => {
353389
const position = input as LspPositionParams;
354-
const result = await sendRequest(params.session, params.method, {
355-
textDocument: { uri: position.uri },
356-
position: { line: position.line, character: position.character },
357-
});
390+
const result = await sendRequest(
391+
params.session,
392+
params.method,
393+
{
394+
textDocument: { uri: position.uri },
395+
position: { line: position.line, character: position.character },
396+
},
397+
signal,
398+
);
358399
return formatLspResult(params.session.serverName, params.resultLabel, result);
359400
},
360401
};
@@ -409,18 +450,23 @@ function buildLspTools(session: LspSession): AnyAgentTool[] {
409450
},
410451
required: ["uri", "line", "character"],
411452
},
412-
execute: async (_toolCallId, input) => {
453+
execute: async (_toolCallId, input, signal) => {
413454
const params = input as {
414455
uri: string;
415456
line: number;
416457
character: number;
417458
includeDeclaration?: boolean;
418459
};
419-
const result = await sendRequest(session, "textDocument/references", {
420-
textDocument: { uri: params.uri },
421-
position: { line: params.line, character: params.character },
422-
context: { includeDeclaration: params.includeDeclaration ?? true },
423-
});
460+
const result = await sendRequest(
461+
session,
462+
"textDocument/references",
463+
{
464+
textDocument: { uri: params.uri },
465+
position: { line: params.line, character: params.character },
466+
context: { includeDeclaration: params.includeDeclaration ?? true },
467+
},
468+
signal,
469+
);
424470
return formatLspResult(serverLabel, "references", result);
425471
},
426472
});

0 commit comments

Comments
 (0)