feat(ide ext): Write workspace path to port file (#6659)
This commit is contained in:
Shreya Keshive
GitHub
parent
6aff66f501
commit
80ff3cd25e
|
|
@ -1,211 +0,0 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import * as vscode from 'vscode';
|
||||
import * as path from 'path';
|
||||
import { activate } from './extension.js';
|
||||
|
||||
vi.mock('vscode', () => ({
|
||||
window: {
|
||||
createOutputChannel: vi.fn(() => ({
|
||||
appendLine: vi.fn(),
|
||||
})),
|
||||
showInformationMessage: vi.fn(),
|
||||
createTerminal: vi.fn(() => ({
|
||||
show: vi.fn(),
|
||||
sendText: vi.fn(),
|
||||
})),
|
||||
onDidChangeActiveTextEditor: vi.fn(),
|
||||
activeTextEditor: undefined,
|
||||
tabGroups: {
|
||||
all: [],
|
||||
close: vi.fn(),
|
||||
},
|
||||
showTextDocument: vi.fn(),
|
||||
},
|
||||
workspace: {
|
||||
workspaceFolders: [],
|
||||
onDidCloseTextDocument: vi.fn(),
|
||||
registerTextDocumentContentProvider: vi.fn(),
|
||||
onDidChangeWorkspaceFolders: vi.fn(),
|
||||
},
|
||||
commands: {
|
||||
registerCommand: vi.fn(),
|
||||
executeCommand: vi.fn(),
|
||||
},
|
||||
Uri: {
|
||||
joinPath: vi.fn(),
|
||||
file: (path: string) => ({ fsPath: path }),
|
||||
},
|
||||
ExtensionMode: {
|
||||
Development: 1,
|
||||
Production: 2,
|
||||
},
|
||||
EventEmitter: vi.fn(() => ({
|
||||
event: vi.fn(),
|
||||
fire: vi.fn(),
|
||||
dispose: vi.fn(),
|
||||
})),
|
||||
}));
|
||||
|
||||
describe('activate with multiple folders', () => {
|
||||
let context: vscode.ExtensionContext;
|
||||
let onDidChangeWorkspaceFoldersCallback: (
|
||||
e: vscode.WorkspaceFoldersChangeEvent,
|
||||
) => void;
|
||||
|
||||
beforeEach(() => {
|
||||
context = {
|
||||
subscriptions: [],
|
||||
environmentVariableCollection: {
|
||||
replace: vi.fn(),
|
||||
},
|
||||
globalState: {
|
||||
get: vi.fn().mockReturnValue(true),
|
||||
update: vi.fn(),
|
||||
},
|
||||
extensionUri: {
|
||||
fsPath: '/path/to/extension',
|
||||
},
|
||||
} as unknown as vscode.ExtensionContext;
|
||||
|
||||
vi.mocked(vscode.workspace.onDidChangeWorkspaceFolders).mockImplementation(
|
||||
(callback) => {
|
||||
onDidChangeWorkspaceFoldersCallback = callback;
|
||||
return { dispose: vi.fn() };
|
||||
},
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it('should set a single folder path', async () => {
|
||||
const workspaceFoldersSpy = vi.spyOn(
|
||||
vscode.workspace,
|
||||
'workspaceFolders',
|
||||
'get',
|
||||
);
|
||||
workspaceFoldersSpy.mockReturnValue([
|
||||
{ uri: { fsPath: '/foo/bar' } },
|
||||
] as vscode.WorkspaceFolder[]);
|
||||
|
||||
await activate(context);
|
||||
|
||||
expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'/foo/bar',
|
||||
);
|
||||
});
|
||||
|
||||
it('should set multiple folder paths, separated by OS-specific path delimiter', async () => {
|
||||
const workspaceFoldersSpy = vi.spyOn(
|
||||
vscode.workspace,
|
||||
'workspaceFolders',
|
||||
'get',
|
||||
);
|
||||
workspaceFoldersSpy.mockReturnValue([
|
||||
{ uri: { fsPath: '/foo/bar' } },
|
||||
{ uri: { fsPath: '/baz/qux' } },
|
||||
] as vscode.WorkspaceFolder[]);
|
||||
|
||||
await activate(context);
|
||||
|
||||
expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
['/foo/bar', '/baz/qux'].join(path.delimiter),
|
||||
);
|
||||
});
|
||||
|
||||
it('should set an empty string if no folders are open', async () => {
|
||||
const workspaceFoldersSpy = vi.spyOn(
|
||||
vscode.workspace,
|
||||
'workspaceFolders',
|
||||
'get',
|
||||
);
|
||||
workspaceFoldersSpy.mockReturnValue([]);
|
||||
|
||||
await activate(context);
|
||||
|
||||
expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'',
|
||||
);
|
||||
});
|
||||
|
||||
it('should update the path when workspace folders change', async () => {
|
||||
const workspaceFoldersSpy = vi.spyOn(
|
||||
vscode.workspace,
|
||||
'workspaceFolders',
|
||||
'get',
|
||||
);
|
||||
workspaceFoldersSpy.mockReturnValue([
|
||||
{ uri: { fsPath: '/foo/bar' } },
|
||||
] as vscode.WorkspaceFolder[]);
|
||||
|
||||
await activate(context);
|
||||
|
||||
expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'/foo/bar',
|
||||
);
|
||||
|
||||
// Simulate adding a folder
|
||||
workspaceFoldersSpy.mockReturnValue([
|
||||
{ uri: { fsPath: '/foo/bar' } },
|
||||
{ uri: { fsPath: '/baz/qux' } },
|
||||
] as vscode.WorkspaceFolder[]);
|
||||
onDidChangeWorkspaceFoldersCallback({
|
||||
added: [{ uri: { fsPath: '/baz/qux' } } as vscode.WorkspaceFolder],
|
||||
removed: [],
|
||||
});
|
||||
|
||||
expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
['/foo/bar', '/baz/qux'].join(path.delimiter),
|
||||
);
|
||||
|
||||
// Simulate removing a folder
|
||||
workspaceFoldersSpy.mockReturnValue([
|
||||
{ uri: { fsPath: '/baz/qux' } },
|
||||
] as vscode.WorkspaceFolder[]);
|
||||
onDidChangeWorkspaceFoldersCallback({
|
||||
added: [],
|
||||
removed: [{ uri: { fsPath: '/foo/bar' } } as vscode.WorkspaceFolder],
|
||||
});
|
||||
|
||||
expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'/baz/qux',
|
||||
);
|
||||
});
|
||||
|
||||
it.skipIf(process.platform !== 'win32')(
|
||||
'should handle windows paths',
|
||||
async () => {
|
||||
const workspaceFoldersSpy = vi.spyOn(
|
||||
vscode.workspace,
|
||||
'workspaceFolders',
|
||||
'get',
|
||||
);
|
||||
workspaceFoldersSpy.mockReturnValue([
|
||||
{ uri: { fsPath: 'c:/foo/bar' } },
|
||||
{ uri: { fsPath: 'd:/baz/qux' } },
|
||||
] as vscode.WorkspaceFolder[]);
|
||||
|
||||
await activate(context);
|
||||
|
||||
expect(
|
||||
context.environmentVariableCollection.replace,
|
||||
).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'c:/foo/bar;d:/baz/qux',
|
||||
);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
|
@ -5,13 +5,11 @@
|
|||
*/
|
||||
|
||||
import * as vscode from 'vscode';
|
||||
import * as path from 'path';
|
||||
import { IDEServer } from './ide-server.js';
|
||||
import { DiffContentProvider, DiffManager } from './diff-manager.js';
|
||||
import { createLogger } from './utils/logger.js';
|
||||
|
||||
const INFO_MESSAGE_SHOWN_KEY = 'geminiCliInfoMessageShown';
|
||||
const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH';
|
||||
export const DIFF_SCHEME = 'gemini-diff';
|
||||
|
||||
let ideServer: IDEServer;
|
||||
|
|
@ -19,31 +17,11 @@ let logger: vscode.OutputChannel;
|
|||
|
||||
let log: (message: string) => void = () => {};
|
||||
|
||||
function updateWorkspacePath(context: vscode.ExtensionContext) {
|
||||
const workspaceFolders = vscode.workspace.workspaceFolders;
|
||||
if (workspaceFolders && workspaceFolders.length > 0) {
|
||||
const workspacePaths = workspaceFolders
|
||||
.map((folder) => folder.uri.fsPath)
|
||||
.join(path.delimiter);
|
||||
context.environmentVariableCollection.replace(
|
||||
IDE_WORKSPACE_PATH_ENV_VAR,
|
||||
workspacePaths,
|
||||
);
|
||||
} else {
|
||||
context.environmentVariableCollection.replace(
|
||||
IDE_WORKSPACE_PATH_ENV_VAR,
|
||||
'',
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
export async function activate(context: vscode.ExtensionContext) {
|
||||
logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion');
|
||||
log = createLogger(context, logger);
|
||||
log('Extension activated');
|
||||
|
||||
updateWorkspacePath(context);
|
||||
|
||||
const diffContentProvider = new DiffContentProvider();
|
||||
const diffManager = new DiffManager(log, diffContentProvider);
|
||||
|
||||
|
|
@ -94,7 +72,7 @@ export async function activate(context: vscode.ExtensionContext) {
|
|||
|
||||
context.subscriptions.push(
|
||||
vscode.workspace.onDidChangeWorkspaceFolders(() => {
|
||||
updateWorkspacePath(context);
|
||||
ideServer.updateWorkspacePath();
|
||||
}),
|
||||
vscode.commands.registerCommand('gemini-cli.runGeminiCLI', async () => {
|
||||
const workspaceFolders = vscode.workspace.workspaceFolders;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,287 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import type * as vscode from 'vscode';
|
||||
import * as fs from 'node:fs/promises';
|
||||
import * as os from 'node:os';
|
||||
import * as path from 'node:path';
|
||||
import { IDEServer } from './ide-server.js';
|
||||
import { DiffManager } from './diff-manager.js';
|
||||
|
||||
const mocks = vi.hoisted(() => ({
|
||||
diffManager: {
|
||||
onDidChange: vi.fn(() => ({ dispose: vi.fn() })),
|
||||
} as unknown as DiffManager,
|
||||
}));
|
||||
|
||||
vi.mock('node:fs/promises', () => ({
|
||||
writeFile: vi.fn(() => Promise.resolve(undefined)),
|
||||
unlink: vi.fn(() => Promise.resolve(undefined)),
|
||||
}));
|
||||
|
||||
vi.mock('node:os', async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof os>();
|
||||
return {
|
||||
...actual,
|
||||
tmpdir: vi.fn(() => '/tmp'),
|
||||
};
|
||||
});
|
||||
|
||||
const vscodeMock = vi.hoisted(() => ({
|
||||
workspace: {
|
||||
workspaceFolders: [
|
||||
{
|
||||
uri: {
|
||||
fsPath: '/test/workspace1',
|
||||
},
|
||||
},
|
||||
{
|
||||
uri: {
|
||||
fsPath: '/test/workspace2',
|
||||
},
|
||||
},
|
||||
],
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('vscode', () => vscodeMock);
|
||||
|
||||
vi.mock('./open-files-manager', () => {
|
||||
const OpenFilesManager = vi.fn();
|
||||
OpenFilesManager.prototype.onDidChange = vi.fn(() => ({ dispose: vi.fn() }));
|
||||
return { OpenFilesManager };
|
||||
});
|
||||
|
||||
describe('IDEServer', () => {
|
||||
let ideServer: IDEServer;
|
||||
let mockContext: vscode.ExtensionContext;
|
||||
let mockLog: (message: string) => void;
|
||||
|
||||
const getPortFromMock = (
|
||||
replaceMock: ReturnType<
|
||||
() => vscode.ExtensionContext['environmentVariableCollection']['replace']
|
||||
>,
|
||||
) => {
|
||||
const port = vi
|
||||
.mocked(replaceMock)
|
||||
.mock.calls.find((call) => call[0] === 'GEMINI_CLI_IDE_SERVER_PORT')?.[1];
|
||||
|
||||
if (port === undefined) {
|
||||
expect.fail('Port was not set');
|
||||
}
|
||||
return port;
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
mockLog = vi.fn();
|
||||
ideServer = new IDEServer(mockLog, mocks.diffManager);
|
||||
mockContext = {
|
||||
subscriptions: [],
|
||||
environmentVariableCollection: {
|
||||
replace: vi.fn(),
|
||||
clear: vi.fn(),
|
||||
},
|
||||
} as unknown as vscode.ExtensionContext;
|
||||
});
|
||||
|
||||
afterEach(async () => {
|
||||
await ideServer.stop();
|
||||
vi.restoreAllMocks();
|
||||
vscodeMock.workspace.workspaceFolders = [
|
||||
{ uri: { fsPath: '/test/workspace1' } },
|
||||
{ uri: { fsPath: '/test/workspace2' } },
|
||||
];
|
||||
});
|
||||
|
||||
it('should set environment variables and workspace path on start with multiple folders', async () => {
|
||||
await ideServer.start(mockContext);
|
||||
|
||||
const replaceMock = mockContext.environmentVariableCollection.replace;
|
||||
expect(replaceMock).toHaveBeenCalledTimes(2);
|
||||
|
||||
expect(replaceMock).toHaveBeenNthCalledWith(
|
||||
1,
|
||||
'GEMINI_CLI_IDE_SERVER_PORT',
|
||||
expect.any(String), // port is a number as a string
|
||||
);
|
||||
|
||||
const expectedWorkspacePaths = [
|
||||
'/test/workspace1',
|
||||
'/test/workspace2',
|
||||
].join(path.delimiter);
|
||||
|
||||
expect(replaceMock).toHaveBeenNthCalledWith(
|
||||
2,
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
expectedWorkspacePaths,
|
||||
);
|
||||
|
||||
const port = getPortFromMock(replaceMock);
|
||||
const expectedPortFile = path.join(
|
||||
'/tmp',
|
||||
`gemini-ide-server-${process.ppid}.json`,
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(
|
||||
expectedPortFile,
|
||||
JSON.stringify({
|
||||
port: parseInt(port, 10),
|
||||
workspacePath: expectedWorkspacePaths,
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should set a single folder path', async () => {
|
||||
vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/foo/bar' } }];
|
||||
|
||||
await ideServer.start(mockContext);
|
||||
const replaceMock = mockContext.environmentVariableCollection.replace;
|
||||
|
||||
expect(replaceMock).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'/foo/bar',
|
||||
);
|
||||
|
||||
const port = getPortFromMock(replaceMock);
|
||||
const expectedPortFile = path.join(
|
||||
'/tmp',
|
||||
`gemini-ide-server-${process.ppid}.json`,
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(
|
||||
expectedPortFile,
|
||||
JSON.stringify({
|
||||
port: parseInt(port, 10),
|
||||
workspacePath: '/foo/bar',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should set an empty string if no folders are open', async () => {
|
||||
vscodeMock.workspace.workspaceFolders = [];
|
||||
|
||||
await ideServer.start(mockContext);
|
||||
const replaceMock = mockContext.environmentVariableCollection.replace;
|
||||
|
||||
expect(replaceMock).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'',
|
||||
);
|
||||
|
||||
const port = getPortFromMock(replaceMock);
|
||||
const expectedPortFile = path.join(
|
||||
'/tmp',
|
||||
`gemini-ide-server-${process.ppid}.json`,
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(
|
||||
expectedPortFile,
|
||||
JSON.stringify({
|
||||
port: parseInt(port, 10),
|
||||
workspacePath: '',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should update the path when workspace folders change', async () => {
|
||||
vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/foo/bar' } }];
|
||||
await ideServer.start(mockContext);
|
||||
const replaceMock = mockContext.environmentVariableCollection.replace;
|
||||
|
||||
expect(replaceMock).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'/foo/bar',
|
||||
);
|
||||
|
||||
// Simulate adding a folder
|
||||
vscodeMock.workspace.workspaceFolders = [
|
||||
{ uri: { fsPath: '/foo/bar' } },
|
||||
{ uri: { fsPath: '/baz/qux' } },
|
||||
];
|
||||
await ideServer.updateWorkspacePath();
|
||||
|
||||
const expectedWorkspacePaths = ['/foo/bar', '/baz/qux'].join(
|
||||
path.delimiter,
|
||||
);
|
||||
expect(replaceMock).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
expectedWorkspacePaths,
|
||||
);
|
||||
|
||||
const port = getPortFromMock(replaceMock);
|
||||
const expectedPortFile = path.join(
|
||||
'/tmp',
|
||||
`gemini-ide-server-${process.ppid}.json`,
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(
|
||||
expectedPortFile,
|
||||
JSON.stringify({
|
||||
port: parseInt(port, 10),
|
||||
workspacePath: expectedWorkspacePaths,
|
||||
}),
|
||||
);
|
||||
|
||||
// Simulate removing a folder
|
||||
vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/baz/qux' } }];
|
||||
await ideServer.updateWorkspacePath();
|
||||
|
||||
expect(replaceMock).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
'/baz/qux',
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(
|
||||
expectedPortFile,
|
||||
JSON.stringify({
|
||||
port: parseInt(port, 10),
|
||||
workspacePath: '/baz/qux',
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should clear env vars and delete port file on stop', async () => {
|
||||
await ideServer.start(mockContext);
|
||||
const portFile = path.join(
|
||||
'/tmp',
|
||||
`gemini-ide-server-${process.ppid}.json`,
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(portFile, expect.any(String));
|
||||
|
||||
await ideServer.stop();
|
||||
|
||||
expect(mockContext.environmentVariableCollection.clear).toHaveBeenCalled();
|
||||
expect(fs.unlink).toHaveBeenCalledWith(portFile);
|
||||
});
|
||||
|
||||
it.skipIf(process.platform !== 'win32')(
|
||||
'should handle windows paths',
|
||||
async () => {
|
||||
vscodeMock.workspace.workspaceFolders = [
|
||||
{ uri: { fsPath: 'c:\\foo\\bar' } },
|
||||
{ uri: { fsPath: 'd:\\baz\\qux' } },
|
||||
];
|
||||
|
||||
await ideServer.start(mockContext);
|
||||
const replaceMock = mockContext.environmentVariableCollection.replace;
|
||||
const expectedWorkspacePaths = 'c:\\foo\\bar;d:\\baz\\qux';
|
||||
|
||||
expect(replaceMock).toHaveBeenCalledWith(
|
||||
'GEMINI_CLI_IDE_WORKSPACE_PATH',
|
||||
expectedWorkspacePaths,
|
||||
);
|
||||
|
||||
const port = getPortFromMock(replaceMock);
|
||||
const expectedPortFile = path.join(
|
||||
'/tmp',
|
||||
`gemini-ide-server-${process.ppid}.json`,
|
||||
);
|
||||
expect(fs.writeFile).toHaveBeenCalledWith(
|
||||
expectedPortFile,
|
||||
JSON.stringify({
|
||||
port: parseInt(port, 10),
|
||||
workspacePath: expectedWorkspacePaths,
|
||||
}),
|
||||
);
|
||||
},
|
||||
);
|
||||
});
|
||||
|
|
@ -21,6 +21,37 @@ import { OpenFilesManager } from './open-files-manager.js';
|
|||
|
||||
const MCP_SESSION_ID_HEADER = 'mcp-session-id';
|
||||
const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT';
|
||||
const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH';
|
||||
|
||||
function writePortAndWorkspace(
|
||||
context: vscode.ExtensionContext,
|
||||
port: number,
|
||||
portFile: string,
|
||||
log: (message: string) => void,
|
||||
): Promise<void> {
|
||||
const workspaceFolders = vscode.workspace.workspaceFolders;
|
||||
const workspacePath =
|
||||
workspaceFolders && workspaceFolders.length > 0
|
||||
? workspaceFolders.map((folder) => folder.uri.fsPath).join(path.delimiter)
|
||||
: '';
|
||||
|
||||
context.environmentVariableCollection.replace(
|
||||
IDE_SERVER_PORT_ENV_VAR,
|
||||
port.toString(),
|
||||
);
|
||||
context.environmentVariableCollection.replace(
|
||||
IDE_WORKSPACE_PATH_ENV_VAR,
|
||||
workspacePath,
|
||||
);
|
||||
|
||||
log(`Writing port file to: ${portFile}`);
|
||||
return fs
|
||||
.writeFile(portFile, JSON.stringify({ port, workspacePath }))
|
||||
.catch((err) => {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
log(`Failed to write port to file: ${message}`);
|
||||
});
|
||||
}
|
||||
|
||||
function sendIdeContextUpdateNotification(
|
||||
transport: StreamableHTTPServerTransport,
|
||||
|
|
@ -50,6 +81,7 @@ export class IDEServer {
|
|||
private context: vscode.ExtensionContext | undefined;
|
||||
private log: (message: string) => void;
|
||||
private portFile: string;
|
||||
private port: number | undefined;
|
||||
diffManager: DiffManager;
|
||||
|
||||
constructor(log: (message: string) => void, diffManager: DiffManager) {
|
||||
|
|
@ -61,158 +93,170 @@ export class IDEServer {
|
|||
);
|
||||
}
|
||||
|
||||
async start(context: vscode.ExtensionContext) {
|
||||
this.context = context;
|
||||
const sessionsWithInitialNotification = new Set<string>();
|
||||
const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
|
||||
{};
|
||||
start(context: vscode.ExtensionContext): Promise<void> {
|
||||
return new Promise((resolve) => {
|
||||
this.context = context;
|
||||
const sessionsWithInitialNotification = new Set<string>();
|
||||
const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
|
||||
{};
|
||||
|
||||
const app = express();
|
||||
app.use(express.json());
|
||||
const mcpServer = createMcpServer(this.diffManager);
|
||||
const app = express();
|
||||
app.use(express.json());
|
||||
const mcpServer = createMcpServer(this.diffManager);
|
||||
|
||||
const openFilesManager = new OpenFilesManager(context);
|
||||
const onDidChangeSubscription = openFilesManager.onDidChange(() => {
|
||||
for (const transport of Object.values(transports)) {
|
||||
sendIdeContextUpdateNotification(
|
||||
transport,
|
||||
this.log.bind(this),
|
||||
openFilesManager,
|
||||
);
|
||||
}
|
||||
});
|
||||
context.subscriptions.push(onDidChangeSubscription);
|
||||
const onDidChangeDiffSubscription = this.diffManager.onDidChange(
|
||||
(notification) => {
|
||||
const openFilesManager = new OpenFilesManager(context);
|
||||
const onDidChangeSubscription = openFilesManager.onDidChange(() => {
|
||||
for (const transport of Object.values(transports)) {
|
||||
transport.send(notification);
|
||||
sendIdeContextUpdateNotification(
|
||||
transport,
|
||||
this.log.bind(this),
|
||||
openFilesManager,
|
||||
);
|
||||
}
|
||||
},
|
||||
);
|
||||
context.subscriptions.push(onDidChangeDiffSubscription);
|
||||
});
|
||||
context.subscriptions.push(onDidChangeSubscription);
|
||||
const onDidChangeDiffSubscription = this.diffManager.onDidChange(
|
||||
(notification) => {
|
||||
for (const transport of Object.values(transports)) {
|
||||
transport.send(notification);
|
||||
}
|
||||
},
|
||||
);
|
||||
context.subscriptions.push(onDidChangeDiffSubscription);
|
||||
|
||||
app.post('/mcp', async (req: Request, res: Response) => {
|
||||
const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
|
||||
| string
|
||||
| undefined;
|
||||
let transport: StreamableHTTPServerTransport;
|
||||
app.post('/mcp', async (req: Request, res: Response) => {
|
||||
const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
|
||||
| string
|
||||
| undefined;
|
||||
let transport: StreamableHTTPServerTransport;
|
||||
|
||||
if (sessionId && transports[sessionId]) {
|
||||
transport = transports[sessionId];
|
||||
} else if (!sessionId && isInitializeRequest(req.body)) {
|
||||
transport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => randomUUID(),
|
||||
onsessioninitialized: (newSessionId) => {
|
||||
this.log(`New session initialized: ${newSessionId}`);
|
||||
transports[newSessionId] = transport;
|
||||
},
|
||||
});
|
||||
const keepAlive = setInterval(() => {
|
||||
try {
|
||||
transport.send({ jsonrpc: '2.0', method: 'ping' });
|
||||
} catch (e) {
|
||||
this.log(
|
||||
'Failed to send keep-alive ping, cleaning up interval.' + e,
|
||||
);
|
||||
if (sessionId && transports[sessionId]) {
|
||||
transport = transports[sessionId];
|
||||
} else if (!sessionId && isInitializeRequest(req.body)) {
|
||||
transport = new StreamableHTTPServerTransport({
|
||||
sessionIdGenerator: () => randomUUID(),
|
||||
onsessioninitialized: (newSessionId) => {
|
||||
this.log(`New session initialized: ${newSessionId}`);
|
||||
transports[newSessionId] = transport;
|
||||
},
|
||||
});
|
||||
const keepAlive = setInterval(() => {
|
||||
try {
|
||||
transport.send({ jsonrpc: '2.0', method: 'ping' });
|
||||
} catch (e) {
|
||||
this.log(
|
||||
'Failed to send keep-alive ping, cleaning up interval.' + e,
|
||||
);
|
||||
clearInterval(keepAlive);
|
||||
}
|
||||
}, 60000); // 60 sec
|
||||
|
||||
transport.onclose = () => {
|
||||
clearInterval(keepAlive);
|
||||
}
|
||||
}, 60000); // 60 sec
|
||||
|
||||
transport.onclose = () => {
|
||||
clearInterval(keepAlive);
|
||||
if (transport.sessionId) {
|
||||
this.log(`Session closed: ${transport.sessionId}`);
|
||||
sessionsWithInitialNotification.delete(transport.sessionId);
|
||||
delete transports[transport.sessionId];
|
||||
}
|
||||
};
|
||||
mcpServer.connect(transport);
|
||||
} else {
|
||||
this.log(
|
||||
'Bad Request: No valid session ID provided for non-initialize request.',
|
||||
);
|
||||
res.status(400).json({
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
code: -32000,
|
||||
message:
|
||||
'Bad Request: No valid session ID provided for non-initialize request.',
|
||||
},
|
||||
id: null,
|
||||
});
|
||||
return;
|
||||
}
|
||||
|
||||
try {
|
||||
await transport.handleRequest(req, res, req.body);
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
this.log(`Error handling MCP request: ${errorMessage}`);
|
||||
if (!res.headersSent) {
|
||||
res.status(500).json({
|
||||
jsonrpc: '2.0' as const,
|
||||
if (transport.sessionId) {
|
||||
this.log(`Session closed: ${transport.sessionId}`);
|
||||
sessionsWithInitialNotification.delete(transport.sessionId);
|
||||
delete transports[transport.sessionId];
|
||||
}
|
||||
};
|
||||
mcpServer.connect(transport);
|
||||
} else {
|
||||
this.log(
|
||||
'Bad Request: No valid session ID provided for non-initialize request.',
|
||||
);
|
||||
res.status(400).json({
|
||||
jsonrpc: '2.0',
|
||||
error: {
|
||||
code: -32603,
|
||||
message: 'Internal server error',
|
||||
code: -32000,
|
||||
message:
|
||||
'Bad Request: No valid session ID provided for non-initialize request.',
|
||||
},
|
||||
id: null,
|
||||
});
|
||||
return;
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
const handleSessionRequest = async (req: Request, res: Response) => {
|
||||
const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
|
||||
| string
|
||||
| undefined;
|
||||
if (!sessionId || !transports[sessionId]) {
|
||||
this.log('Invalid or missing session ID');
|
||||
res.status(400).send('Invalid or missing session ID');
|
||||
return;
|
||||
}
|
||||
|
||||
const transport = transports[sessionId];
|
||||
try {
|
||||
await transport.handleRequest(req, res);
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
this.log(`Error handling session request: ${errorMessage}`);
|
||||
if (!res.headersSent) {
|
||||
res.status(400).send('Bad Request');
|
||||
try {
|
||||
await transport.handleRequest(req, res, req.body);
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
this.log(`Error handling MCP request: ${errorMessage}`);
|
||||
if (!res.headersSent) {
|
||||
res.status(500).json({
|
||||
jsonrpc: '2.0' as const,
|
||||
error: {
|
||||
code: -32603,
|
||||
message: 'Internal server error',
|
||||
},
|
||||
id: null,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
if (!sessionsWithInitialNotification.has(sessionId)) {
|
||||
sendIdeContextUpdateNotification(
|
||||
transport,
|
||||
this.log.bind(this),
|
||||
openFilesManager,
|
||||
);
|
||||
sessionsWithInitialNotification.add(sessionId);
|
||||
}
|
||||
};
|
||||
const handleSessionRequest = async (req: Request, res: Response) => {
|
||||
const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
|
||||
| string
|
||||
| undefined;
|
||||
if (!sessionId || !transports[sessionId]) {
|
||||
this.log('Invalid or missing session ID');
|
||||
res.status(400).send('Invalid or missing session ID');
|
||||
return;
|
||||
}
|
||||
|
||||
app.get('/mcp', handleSessionRequest);
|
||||
const transport = transports[sessionId];
|
||||
try {
|
||||
await transport.handleRequest(req, res);
|
||||
} catch (error) {
|
||||
const errorMessage =
|
||||
error instanceof Error ? error.message : 'Unknown error';
|
||||
this.log(`Error handling session request: ${errorMessage}`);
|
||||
if (!res.headersSent) {
|
||||
res.status(400).send('Bad Request');
|
||||
}
|
||||
}
|
||||
|
||||
this.server = app.listen(0, () => {
|
||||
const address = (this.server as HTTPServer).address();
|
||||
if (address && typeof address !== 'string') {
|
||||
const port = address.port;
|
||||
context.environmentVariableCollection.replace(
|
||||
IDE_SERVER_PORT_ENV_VAR,
|
||||
port.toString(),
|
||||
);
|
||||
this.log(`IDE server listening on port ${port}`);
|
||||
fs.writeFile(this.portFile, JSON.stringify({ port })).catch((err) => {
|
||||
this.log(`Failed to write port to file: ${err}`);
|
||||
});
|
||||
this.log(this.portFile);
|
||||
}
|
||||
if (!sessionsWithInitialNotification.has(sessionId)) {
|
||||
sendIdeContextUpdateNotification(
|
||||
transport,
|
||||
this.log.bind(this),
|
||||
openFilesManager,
|
||||
);
|
||||
sessionsWithInitialNotification.add(sessionId);
|
||||
}
|
||||
};
|
||||
|
||||
app.get('/mcp', handleSessionRequest);
|
||||
|
||||
this.server = app.listen(0, async () => {
|
||||
const address = (this.server as HTTPServer).address();
|
||||
if (address && typeof address !== 'string') {
|
||||
this.port = address.port;
|
||||
this.log(`IDE server listening on port ${this.port}`);
|
||||
await writePortAndWorkspace(
|
||||
context,
|
||||
this.port,
|
||||
this.portFile,
|
||||
this.log,
|
||||
);
|
||||
}
|
||||
resolve();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async updateWorkspacePath(): Promise<void> {
|
||||
if (this.context && this.port) {
|
||||
await writePortAndWorkspace(
|
||||
this.context,
|
||||
this.port,
|
||||
this.portFile,
|
||||
this.log,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async stop(): Promise<void> {
|
||||
if (this.server) {
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
|
|
|
|||
Loading…
Reference in New Issue