Skip to content
This repository was archived by the owner on May 20, 2026. It is now read-only.

Commit 231a0f8

Browse files
lukkalukka
committed
lsp query tool for CCLI
1 parent 4995bb8 commit 231a0f8

3 files changed

Lines changed: 331 additions & 5 deletions

File tree

‎src/extension/chatSessions/copilotcli/vscode-node/test/runLspQuery.spec.ts‎

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,172 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { beforeEach, describe, expect, it, vi } from 'vitest';
7+
import { TestLogService } from '../../../../../platform/testing/common/testLogService';
8+
import { MockMcpServer } from './testHelpers';
9+
10+
const { mockExecuteCommand, mockWriteFile } = vi.hoisted(() => ({
11+
mockExecuteCommand: vi.fn(),
12+
mockWriteFile: vi.fn(),
13+
}));
14+
15+
vi.mock('vscode', () => {
16+
class MockPosition {
17+
constructor(public line: number, public character: number) { }
18+
}
19+
class MockUri {
20+
public fsPath: string;
21+
public scheme: string;
22+
constructor(public _str: string) {
23+
this.fsPath = _str.replace('file://', '');
24+
this.scheme = _str.startsWith('file:') ? 'file' : 'unknown';
25+
}
26+
toString() {
27+
return this._str;
28+
}
29+
}
30+
return {
31+
Position: MockPosition,
32+
Uri: {
33+
parse: (str: string) => new MockUri(str),
34+
file: (path: string) => new MockUri(`file://${path}`),
35+
},
36+
commands: {
37+
executeCommand: mockExecuteCommand,
38+
},
39+
};
40+
});
41+
42+
vi.mock('fs/promises', () => ({
43+
writeFile: mockWriteFile,
44+
}));
45+
46+
vi.mock('os', () => ({
47+
tmpdir: () => '/mock/tmpdir',
48+
}));
49+
50+
import { registerRunLspQueryTool } from '../tools/runLspQuery';
51+
52+
describe('runLspQuery tool', () => {
53+
const logger = new TestLogService();
54+
let server: MockMcpServer;
55+
56+
beforeEach(() => {
57+
vi.clearAllMocks();
58+
server = new MockMcpServer();
59+
registerRunLspQueryTool(server as unknown as import('@modelcontextprotocol/sdk/server/mcp.js').McpServer, logger);
60+
});
61+
62+
it('should register the run_lsp_query tool', () => {
63+
expect(server.hasToolRegistered('run_lsp_query')).toBe(true);
64+
});
65+
66+
it('should execute a generic LSP query correctly without position', async () => {
67+
mockExecuteCommand.mockResolvedValue([{ name: 'symbol1' }]);
68+
const handler = server.getToolHandler('run_lsp_query')!;
69+
70+
const result: any = await handler({
71+
command: 'vscode.executeDocumentSymbolProvider',
72+
uri: 'file:///test/file.ts',
73+
});
74+
75+
expect(mockExecuteCommand).toHaveBeenCalledWith('vscode.executeDocumentSymbolProvider', expect.objectContaining({ _str: 'file:///test/file.ts' }));
76+
expect(result.content[0].text).toContain('symbol1');
77+
expect(result.content[0].text).toContain('These are the results of executing vscode.executeDocumentSymbolProvider on file:///test/file.ts');
78+
});
79+
80+
it('should execute a positional LSP query correctly', async () => {
81+
mockExecuteCommand.mockResolvedValue({ start: { line: 0, character: 0 }, end: { line: 0, character: 10 } });
82+
const handler = server.getToolHandler('run_lsp_query')!;
83+
84+
const result: any = await handler({
85+
command: 'vscode.executeHoverProvider',
86+
uri: 'file:///test/file.ts',
87+
position: { line: 5, character: 10 }
88+
});
89+
90+
// Check the position argument
91+
const posArg = mockExecuteCommand.mock.calls[0][2];
92+
expect(posArg.line).toBe(5);
93+
expect(posArg.character).toBe(10);
94+
expect(result.content[0].text).toContain('These are the results of executing vscode.executeHoverProvider on file:///test/file.ts at line 5, character 10');
95+
});
96+
97+
it('should execute a workspace symbol query with query argument', async () => {
98+
mockExecuteCommand.mockResolvedValue([]);
99+
const handler = server.getToolHandler('run_lsp_query')!;
100+
101+
const result: any = await handler({
102+
command: 'vscode.executeWorkspaceSymbolProvider',
103+
uri: 'file:///test/file.ts',
104+
query: 'test_query'
105+
});
106+
107+
expect(mockExecuteCommand).toHaveBeenCalledWith('vscode.executeWorkspaceSymbolProvider', 'test_query');
108+
expect(result.content[0].text).toContain('with query "test_query"');
109+
});
110+
111+
it('should compact results by grouping them if array is returned with uri elements', async () => {
112+
const mockResult = [
113+
{ name: 'test1', uri: { fsPath: '/test/a.ts', toString: () => 'file:///test/a.ts' } },
114+
{ name: 'test2', location: { uri: { fsPath: '/test/a.ts', toString: () => 'file:///test/a.ts' } } },
115+
{ name: 'test3', targetUri: { fsPath: '/test/b.ts', toString: () => 'file:///test/b.ts' } }
116+
];
117+
mockExecuteCommand.mockResolvedValue(mockResult);
118+
119+
const handler = server.getToolHandler('run_lsp_query')!;
120+
const result: any = await handler({
121+
command: 'vscode.executeDefinitionProvider',
122+
uri: 'file:///test/file.ts'
123+
});
124+
125+
const text = result.content[0].text;
126+
expect(text).toContain('"/test/a.ts"');
127+
expect(text).toContain('"/test/b.ts"');
128+
expect(text).not.toContain('"uri"'); // Assert stripped URI objects inside compact struct
129+
});
130+
131+
it('should write to temporary file if result is very long', async () => {
132+
// Create a very large array > 50 chars to test limit arrays as well, but each object itself very long
133+
const mockResult = Array.from({ length: 60 }, () => ({ massive_field: 'a'.repeat(2000) }));
134+
mockExecuteCommand.mockResolvedValue(mockResult);
135+
136+
const handler = server.getToolHandler('run_lsp_query')!;
137+
const result: any = await handler({
138+
command: 'vscode.executeDefinitionProvider',
139+
uri: 'file:///test/file.ts'
140+
});
141+
142+
expect(mockWriteFile).toHaveBeenCalled();
143+
expect(result.content[0].text).toContain('The result is very long and has been saved to:');
144+
});
145+
146+
it('should gracefully handle circular references', async () => {
147+
const circularObj: any = { prop: 'value' };
148+
circularObj.self = circularObj;
149+
150+
mockExecuteCommand.mockResolvedValue([circularObj]);
151+
152+
const handler = server.getToolHandler('run_lsp_query')!;
153+
const result: any = await handler({
154+
command: 'vscode.executeDefinitionProvider',
155+
uri: 'file:///test/file.ts'
156+
});
157+
158+
expect(result.content[0].text).toContain('"[Circular]"');
159+
});
160+
161+
it('should catch errors and return them cleanly', async () => {
162+
mockExecuteCommand.mockRejectedValue(new Error('LSP Error!'));
163+
const handler = server.getToolHandler('run_lsp_query')!;
164+
165+
const result: any = await handler({
166+
command: 'vscode.executeDefinitionProvider',
167+
uri: 'file:///test/file.ts'
168+
});
169+
170+
expect(result.content[0].text).toContain('Error executing vscode.executeDefinitionProvider: LSP Error!');
171+
});
172+
});

‎src/extension/chatSessions/copilotcli/vscode-node/tools/index.ts‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,16 +4,17 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
7-
import { registerOpenDiffTool } from './openDiff';
7+
import { ILogger } from '../../../../../platform/log/common/logService';
8+
import { ICopilotCLISessionTracker } from '../copilotCLISessionTracker';
9+
import { DiffStateManager } from '../diffState';
10+
import { ReadonlyContentProvider } from '../readonlyContentProvider';
811
import { registerCloseDiffTool } from './closeDiff';
912
import { registerGetDiagnosticsTool } from './getDiagnostics';
1013
import { registerGetSelectionTool, SelectionState } from './getSelection';
1114
import { registerGetVscodeInfoTool } from './getVscodeInfo';
15+
import { registerOpenDiffTool } from './openDiff';
16+
import { registerRunLspQueryTool } from './runLspQuery';
1217
import { registerUpdateSessionNameTool } from './updateSessionName';
13-
import { ILogger } from '../../../../../platform/log/common/logService';
14-
import { DiffStateManager } from '../diffState';
15-
import { ReadonlyContentProvider } from '../readonlyContentProvider';
16-
import { ICopilotCLISessionTracker } from '../copilotCLISessionTracker';
1718

1819
export { getSelectionInfo, SelectionState } from './getSelection';
1920
export type { SelectionInfo } from './getSelection';
@@ -26,5 +27,6 @@ export function registerTools(server: McpServer, logger: ILogger, diffState: Dif
2627
registerCloseDiffTool(server, logger, diffState);
2728
registerGetDiagnosticsTool(server, logger);
2829
registerUpdateSessionNameTool(server, logger, sessionTracker, sessionId);
30+
registerRunLspQueryTool(server, logger);
2931
logger.debug('All MCP tools registered');
3032
}

‎src/extension/chatSessions/copilotcli/vscode-node/tools/runLspQuery.ts‎

Lines changed: 152 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,152 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
7+
import * as fs from 'fs/promises';
8+
import * as os from 'os';
9+
import * as path from 'path';
10+
import * as vscode from 'vscode';
11+
import { z } from 'zod';
12+
import { ILogger } from '../../../../../platform/log/common/logService';
13+
import { makeTextResult } from './utils';
14+
15+
export function registerRunLspQueryTool(server: McpServer, logger: ILogger): void {
16+
const schema = {
17+
command: z.enum([
18+
'vscode.executeDocumentHighlights',
19+
'vscode.executeDocumentSymbolProvider',
20+
'vscode.executeFormatDocumentProvider',
21+
'vscode.executeFormatRangeProvider',
22+
'vscode.executeFormatOnTypeProvider',
23+
'vscode.executeDefinitionProvider',
24+
'vscode.executeTypeDefinitionProvider',
25+
'vscode.executeDeclarationProvider',
26+
'vscode.executeImplementationProvider',
27+
'vscode.executeReferenceProvider',
28+
'vscode.executeHoverProvider',
29+
'vscode.executeSelectionRangeProvider',
30+
'vscode.executeWorkspaceSymbolProvider',
31+
'vscode.prepareCallHierarchy',
32+
'vscode.provideIncomingCalls',
33+
'vscode.provideOutgoingCalls',
34+
'vscode.prepareRename',
35+
'vscode.executeDocumentRenameProvider',
36+
'vscode.executeLinkProvider',
37+
'vscode.provideDocumentSemanticTokensLegend',
38+
'vscode.provideDocumentSemanticTokens',
39+
'vscode.provideDocumentRangeSemanticTokensLegend'
40+
]).describe('The VS Code LSP command to execute'),
41+
uri: z.string().describe('The URI of the document'),
42+
position: z.object({
43+
line: z.number(),
44+
character: z.number()
45+
}).optional().describe('The position in the document (if required by the command)'),
46+
query: z.string().optional().describe('A string query (for workspace symbols)'),
47+
};
48+
server.registerTool(
49+
'run_lsp_query',
50+
{
51+
description: 'Execute a precise LSP query (e.g. definitions, references, hover) using built-in VS Code commands. Pass the command name and the appropriate arguments.',
52+
inputSchema: schema,
53+
},
54+
// @ts-ignore - TS2589: zod type instantiation too deep for server.tool() generics
55+
async (args: { command: string; uri: string; position?: { line: number; character: number }; query?: string }) => {
56+
const { command, uri, position, query } = args;
57+
logger.debug(`Executing LSP query: ${command} on ${uri}`);
58+
try {
59+
const decodedUri = decodeURIComponent(uri);
60+
const documentUri = decodedUri.startsWith('file:') ? vscode.Uri.parse(decodedUri) : vscode.Uri.file(decodedUri);
61+
const pos = position ? new vscode.Position(position.line, position.character) : undefined;
62+
let result: unknown;
63+
64+
if (command === 'vscode.executeWorkspaceSymbolProvider' && query !== undefined) {
65+
result = await vscode.commands.executeCommand(command, query);
66+
} else if (pos) {
67+
result = await vscode.commands.executeCommand(command, documentUri, pos);
68+
} else {
69+
result = await vscode.commands.executeCommand(command, documentUri);
70+
}
71+
72+
if (Array.isArray(result) && result.length > 0) {
73+
const grouped: Record<string, unknown[]> = {};
74+
const ungrouped: unknown[] = [];
75+
let hasUri = false;
76+
77+
for (const unknownItem of result) {
78+
// Cast to a flexible object temporarily since we just asserted it's in an array
79+
// and we're looking for specific properties
80+
const item = unknownItem as Record<string, unknown>;
81+
82+
if (!item || typeof item !== 'object') {
83+
ungrouped.push(item);
84+
continue;
85+
}
86+
87+
const uriObj = (item.uri || item.targetUri || (item.location as Record<string, unknown>)?.uri) as { fsPath?: string; toString?: () => string } | undefined;
88+
89+
if (uriObj && (uriObj.fsPath || typeof uriObj.toString === 'function')) {
90+
hasUri = true;
91+
const uriStr = uriObj.fsPath || uriObj.toString!();
92+
if (!grouped[uriStr]) {
93+
grouped[uriStr] = [];
94+
}
95+
const compactItem = { ...item };
96+
delete compactItem.uri;
97+
delete compactItem.targetUri;
98+
if (compactItem.location && typeof compactItem.location === 'object') {
99+
const locationObj = compactItem.location as Record<string, unknown>;
100+
if (locationObj.uri) {
101+
compactItem.location = { ...locationObj };
102+
delete (compactItem.location as Record<string, unknown>).uri;
103+
}
104+
}
105+
grouped[uriStr].push(compactItem);
106+
} else {
107+
ungrouped.push(item);
108+
}
109+
}
110+
111+
if (hasUri) {
112+
result = ungrouped.length > 0 ? { grouped, ungrouped } : grouped;
113+
}
114+
}
115+
116+
// Safe stringify handling circular references or overwhelming output
117+
const replacer = () => {
118+
const seen = new WeakSet();
119+
return (key: string, value: unknown) => {
120+
if (typeof value === 'object' && value !== null) {
121+
if (seen.has(value)) {
122+
return '[Circular]';
123+
}
124+
seen.add(value);
125+
}
126+
// Limit arrays
127+
if (Array.isArray(value) && value.length > 50) {
128+
return [...value.slice(0, 50), `... ${value.length - 50} more items`];
129+
}
130+
return value;
131+
};
132+
};
133+
134+
const resultString = JSON.stringify(result, replacer(), 2) ?? 'No results found.';
135+
const explanation = `These are the results of executing ${command} on ${uri}${position ? ` at line ${position.line}, character ${position.character}` : ''}${query ? ` with query "${query}"` : ''}.`;
136+
137+
if (resultString.length > 50000) {
138+
const tmpFile = path.join(os.tmpdir(), `lsp_query_result_${Date.now()}.json`);
139+
await fs.writeFile(tmpFile, resultString);
140+
logger.trace(`Returning saved file path for LSP query ${command} due to size`);
141+
return makeTextResult(`${explanation}\n\nThe result is very long and has been saved to: ${tmpFile}`);
142+
}
143+
144+
logger.trace(`Returning result for LSP query ${command}`);
145+
return makeTextResult(`${explanation}\n\n${resultString}`);
146+
} catch (error) {
147+
logger.error(`Error executing LSP query ${command}: ${error}`);
148+
return makeTextResult(`Error executing ${command}: ${error instanceof Error ? error.message : String(error)}`);
149+
}
150+
}
151+
);
152+
}

0 commit comments

Comments
 (0)