Minor refactoring of VS Code companion extension code (#4761)

This commit is contained in:
Shreya Keshive 2025-07-24 01:32:55 -04:00 committed by GitHub
parent b1e0fb157b
commit 6380bfe35c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
8 changed files with 244 additions and 115 deletions

12
.vscode/launch.json vendored
View File

@ -30,6 +30,18 @@
"GEMINI_SANDBOX": "false"
}
},
{
"name": "Launch Companion VS Code Extension",
"type": "extensionHost",
"request": "launch",
"args": [
"--extensionDevelopmentPath=${workspaceFolder}/packages/vscode-ide-companion"
],
"outFiles": [
"${workspaceFolder}/packages/vscode-ide-companion/dist/**/*.js"
],
"preLaunchTask": "npm: build: vscode-ide-companion"
},
{
"name": "Attach",
"port": 9229,

9
.vscode/tasks.json vendored
View File

@ -11,6 +11,15 @@
"problemMatcher": [],
"label": "npm: build",
"detail": "scripts/build.sh"
},
{
"type": "npm",
"script": "build",
"path": "packages/vscode-ide-companion",
"group": "build",
"problemMatcher": [],
"label": "npm: build: vscode-ide-companion",
"detail": "npm run build -w packages/vscode-ide-companion"
}
]
}

View File

@ -13,7 +13,7 @@ The Gemini CLI Companion extension seamlessly integrates [Gemini CLI](https://gi
To use this extension, you'll need:
- VS Code version 1.101.0 or newer
- Gemini CLI (installed separately) and running within the VS Code integrated terminal
- Gemini CLI (installed separately) running within the VS Code integrated terminal
# Terms of Service and Privacy Notice

View File

@ -6,30 +6,38 @@
import * as vscode from 'vscode';
import { IDEServer } from './ide-server';
import { createLogger } from './utils/logger';
let ideServer: IDEServer;
let logger: vscode.OutputChannel;
let log: (message: string) => void = () => {};
export async function activate(context: vscode.ExtensionContext) {
logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion');
logger.appendLine('Starting Gemini CLI IDE Companion server...');
ideServer = new IDEServer(logger);
log = createLogger(context, logger);
log('Extension activated');
ideServer = new IDEServer(log);
try {
await ideServer.start(context);
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
logger.appendLine(`Failed to start IDE server: ${message}`);
log(`Failed to start IDE server: ${message}`);
}
}
export function deactivate() {
if (ideServer) {
logger.appendLine('Deactivating Gemini CLI IDE Companion...');
return ideServer.stop().finally(() => {
export async function deactivate(): Promise<void> {
log('Extension deactivated');
try {
if (ideServer) {
await ideServer.stop();
}
} catch (err) {
const message = err instanceof Error ? err.message : String(err);
log(`Failed to stop IDE server during deactivation: ${message}`);
} finally {
if (logger) {
logger.dispose();
});
}
if (logger) {
logger.dispose();
}
}
}

View File

@ -21,30 +21,41 @@ const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT';
function sendOpenFilesChangedNotification(
transport: StreamableHTTPServerTransport,
logger: vscode.OutputChannel,
log: (message: string) => void,
recentFilesManager: RecentFilesManager,
) {
const editor = vscode.window.activeTextEditor;
const filePath = editor ? editor.document.uri.fsPath : '';
logger.appendLine(`Sending active file changed notification: ${filePath}`);
const filePath =
editor && editor.document.uri.scheme === 'file'
? editor.document.uri.fsPath
: '';
const notification: JSONRPCNotification = {
jsonrpc: '2.0',
method: 'ide/openFilesChanged',
params: {
activeFile: filePath,
recentOpenFiles: recentFilesManager.recentFiles,
recentOpenFiles: recentFilesManager.recentFiles.filter(
(file) => file.filePath !== filePath,
),
},
};
log(
`Sending active file changed notification: ${JSON.stringify(
notification,
null,
2,
)}`,
);
transport.send(notification);
}
export class IDEServer {
private server: HTTPServer | undefined;
private context: vscode.ExtensionContext | undefined;
private logger: vscode.OutputChannel;
private log: (message: string) => void;
constructor(logger: vscode.OutputChannel) {
this.logger = logger;
constructor(log: (message: string) => void) {
this.log = log;
}
async start(context: vscode.ExtensionContext) {
@ -62,7 +73,7 @@ export class IDEServer {
for (const transport of Object.values(transports)) {
sendOpenFilesChangedNotification(
transport,
this.logger,
this.log.bind(this),
recentFilesManager,
);
}
@ -81,7 +92,7 @@ export class IDEServer {
transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
onsessioninitialized: (newSessionId) => {
this.logger.appendLine(`New session initialized: ${newSessionId}`);
this.log(`New session initialized: ${newSessionId}`);
transports[newSessionId] = transport;
},
});
@ -90,26 +101,24 @@ export class IDEServer {
try {
transport.send({ jsonrpc: '2.0', method: 'ping' });
} catch (e) {
// If sending a ping fails, the connection is likely broken.
// Log the error and clear the interval to prevent further attempts.
this.logger.append(
this.log(
'Failed to send keep-alive ping, cleaning up interval.' + e,
);
clearInterval(keepAlive);
}
}, 60000); // Send ping every 60 seconds
}, 60000); // 60 sec
transport.onclose = () => {
clearInterval(keepAlive);
if (transport.sessionId) {
this.logger.appendLine(`Session closed: ${transport.sessionId}`);
this.log(`Session closed: ${transport.sessionId}`);
sessionsWithInitialNotification.delete(transport.sessionId);
delete transports[transport.sessionId];
}
};
mcpServer.connect(transport);
} else {
this.logger.appendLine(
this.log(
'Bad Request: No valid session ID provided for non-initialize request.',
);
res.status(400).json({
@ -129,7 +138,7 @@ export class IDEServer {
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : 'Unknown error';
this.logger.appendLine(`Error handling MCP request: ${errorMessage}`);
this.log(`Error handling MCP request: ${errorMessage}`);
if (!res.headersSent) {
res.status(500).json({
jsonrpc: '2.0' as const,
@ -148,7 +157,7 @@ export class IDEServer {
| string
| undefined;
if (!sessionId || !transports[sessionId]) {
this.logger.appendLine('Invalid or missing session ID');
this.log('Invalid or missing session ID');
res.status(400).send('Invalid or missing session ID');
return;
}
@ -159,9 +168,7 @@ export class IDEServer {
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : 'Unknown error';
this.logger.appendLine(
`Error handling session request: ${errorMessage}`,
);
this.log(`Error handling session request: ${errorMessage}`);
if (!res.headersSent) {
res.status(400).send('Bad Request');
}
@ -170,7 +177,7 @@ export class IDEServer {
if (!sessionsWithInitialNotification.has(sessionId)) {
sendOpenFilesChangedNotification(
transport,
this.logger,
this.log.bind(this),
recentFilesManager,
);
sessionsWithInitialNotification.add(sessionId);
@ -187,7 +194,7 @@ export class IDEServer {
IDE_SERVER_PORT_ENV_VAR,
port.toString(),
);
this.logger.appendLine(`IDE server listening on port ${port}`);
this.log(`IDE server listening on port ${port}`);
}
});
}
@ -197,12 +204,10 @@ export class IDEServer {
await new Promise<void>((resolve, reject) => {
this.server!.close((err?: Error) => {
if (err) {
this.logger.appendLine(
`Error shutting down IDE server: ${err.message}`,
);
this.log(`Error shutting down IDE server: ${err.message}`);
return reject(err);
}
this.logger.appendLine(`IDE server shut down`);
this.log(`IDE server shut down`);
resolve();
});
});

View File

@ -13,11 +13,19 @@ import {
} from './recent-files-manager.js';
vi.mock('vscode', () => ({
EventEmitter: vi.fn(() => ({
event: vi.fn(),
fire: vi.fn(),
dispose: vi.fn(),
})),
EventEmitter: vi.fn(() => {
const listeners: Array<(e: void) => unknown> = [];
return {
event: vi.fn((listener) => {
listeners.push(listener);
return { dispose: vi.fn() };
}),
fire: vi.fn(() => {
listeners.forEach((listener) => listener(undefined));
}),
dispose: vi.fn(),
};
}),
window: {
onDidChangeActiveTextEditor: vi.fn(),
},
@ -29,14 +37,48 @@ vi.mock('vscode', () => ({
Uri: {
file: (path: string) => ({
fsPath: path,
scheme: 'file',
}),
},
}));
describe('RecentFilesManager', () => {
let context: vscode.ExtensionContext;
let onDidChangeActiveTextEditorListener: (
editor: vscode.TextEditor | undefined,
) => void;
let onDidDeleteFilesListener: (e: vscode.FileDeleteEvent) => void;
let onDidCloseTextDocumentListener: (doc: vscode.TextDocument) => void;
let onDidRenameFilesListener: (e: vscode.FileRenameEvent) => void;
beforeEach(() => {
vi.useFakeTimers();
vi.mocked(vscode.window.onDidChangeActiveTextEditor).mockImplementation(
(listener) => {
onDidChangeActiveTextEditorListener = listener;
return { dispose: vi.fn() };
},
);
vi.mocked(vscode.workspace.onDidDeleteFiles).mockImplementation(
(listener) => {
onDidDeleteFilesListener = listener;
return { dispose: vi.fn() };
},
);
vi.mocked(vscode.workspace.onDidCloseTextDocument).mockImplementation(
(listener) => {
onDidCloseTextDocumentListener = listener;
return { dispose: vi.fn() };
},
);
vi.mocked(vscode.workspace.onDidRenameFiles).mockImplementation(
(listener) => {
onDidRenameFilesListener = listener;
return { dispose: vi.fn() };
},
);
context = {
subscriptions: [],
} as unknown as vscode.ExtensionContext;
@ -44,33 +86,40 @@ describe('RecentFilesManager', () => {
afterEach(() => {
vi.restoreAllMocks();
vi.useRealTimers();
});
it('adds a file to the list', () => {
const getUri = (path: string) =>
vscode.Uri.file(path) as unknown as vscode.Uri;
it('adds a file to the list', async () => {
const manager = new RecentFilesManager(context);
const uri = vscode.Uri.file('/test/file1.txt');
const uri = getUri('/test/file1.txt');
manager.add(uri);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file1.txt');
});
it('moves an existing file to the top', () => {
it('moves an existing file to the top', async () => {
const manager = new RecentFilesManager(context);
const uri1 = vscode.Uri.file('/test/file1.txt');
const uri2 = vscode.Uri.file('/test/file2.txt');
const uri1 = getUri('/test/file1.txt');
const uri2 = getUri('/test/file2.txt');
manager.add(uri1);
manager.add(uri2);
manager.add(uri1);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(2);
expect(manager.recentFiles[0].filePath).toBe('/test/file1.txt');
});
it('does not exceed the max number of files', () => {
it('does not exceed the max number of files', async () => {
const manager = new RecentFilesManager(context);
for (let i = 0; i < MAX_FILES + 5; i++) {
const uri = vscode.Uri.file(`/test/file${i}.txt`);
const uri = getUri(`/test/file${i}.txt`);
manager.add(uri);
}
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(MAX_FILES);
expect(manager.recentFiles[0].filePath).toBe(
`/test/file${MAX_FILES + 4}.txt`,
@ -78,134 +127,151 @@ describe('RecentFilesManager', () => {
expect(manager.recentFiles[MAX_FILES - 1].filePath).toBe(`/test/file5.txt`);
});
it('fires onDidChange when a file is added', () => {
it('fires onDidChange when a file is added', async () => {
const manager = new RecentFilesManager(context);
const spy = vi.spyOn(manager['onDidChangeEmitter'], 'fire');
const uri = vscode.Uri.file('/test/file1.txt');
const onDidChangeSpy = vi.fn();
manager.onDidChange(onDidChangeSpy);
const uri = getUri('/test/file1.txt');
manager.add(uri);
expect(spy).toHaveBeenCalled();
await vi.advanceTimersByTimeAsync(100);
expect(onDidChangeSpy).toHaveBeenCalled();
});
it('removes a file when it is closed', () => {
it('removes a file when it is closed', async () => {
const manager = new RecentFilesManager(context);
const uri = vscode.Uri.file('/test/file1.txt');
const uri = getUri('/test/file1.txt');
manager.add(uri);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
// Simulate closing the file
const closeHandler = vi.mocked(vscode.workspace.onDidCloseTextDocument).mock
.calls[0][0];
closeHandler({ uri } as vscode.TextDocument);
onDidCloseTextDocumentListener({ uri } as vscode.TextDocument);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(0);
});
it('fires onDidChange when a file is removed', () => {
it('fires onDidChange when a file is removed', async () => {
const manager = new RecentFilesManager(context);
const uri = vscode.Uri.file('/test/file1.txt');
const uri = getUri('/test/file1.txt');
manager.add(uri);
await vi.advanceTimersByTimeAsync(100);
const spy = vi.spyOn(manager['onDidChangeEmitter'], 'fire');
const closeHandler = vi.mocked(vscode.workspace.onDidCloseTextDocument).mock
.calls[0][0];
closeHandler({ uri } as vscode.TextDocument);
const onDidChangeSpy = vi.fn();
manager.onDidChange(onDidChangeSpy);
expect(spy).toHaveBeenCalled();
onDidCloseTextDocumentListener({ uri } as vscode.TextDocument);
await vi.advanceTimersByTimeAsync(100);
expect(onDidChangeSpy).toHaveBeenCalled();
});
it('removes a file when it is deleted', () => {
it('removes a file when it is deleted', async () => {
const manager = new RecentFilesManager(context);
const uri1 = vscode.Uri.file('/test/file1.txt');
const uri2 = vscode.Uri.file('/test/file2.txt');
const uri1 = getUri('/test/file1.txt');
const uri2 = getUri('/test/file2.txt');
manager.add(uri1);
manager.add(uri2);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(2);
// Simulate deleting a file
const deleteHandler = vi.mocked(vscode.workspace.onDidDeleteFiles).mock
.calls[0][0];
deleteHandler({ files: [uri1] });
onDidDeleteFilesListener({ files: [uri1] });
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file2.txt');
});
it('fires onDidChange when a file is deleted', () => {
it('fires onDidChange when a file is deleted', async () => {
const manager = new RecentFilesManager(context);
const uri = vscode.Uri.file('/test/file1.txt');
const uri = getUri('/test/file1.txt');
manager.add(uri);
await vi.advanceTimersByTimeAsync(100);
const spy = vi.spyOn(manager['onDidChangeEmitter'], 'fire');
const deleteHandler = vi.mocked(vscode.workspace.onDidDeleteFiles).mock
.calls[0][0];
deleteHandler({ files: [uri] });
const onDidChangeSpy = vi.fn();
manager.onDidChange(onDidChangeSpy);
expect(spy).toHaveBeenCalled();
onDidDeleteFilesListener({ files: [uri] });
await vi.advanceTimersByTimeAsync(100);
expect(onDidChangeSpy).toHaveBeenCalled();
});
it('removes multiple files when they are deleted', () => {
it('removes multiple files when they are deleted', async () => {
const manager = new RecentFilesManager(context);
const uri1 = vscode.Uri.file('/test/file1.txt');
const uri2 = vscode.Uri.file('/test/file2.txt');
const uri3 = vscode.Uri.file('/test/file3.txt');
const uri1 = getUri('/test/file1.txt');
const uri2 = getUri('/test/file2.txt');
const uri3 = getUri('/test/file3.txt');
manager.add(uri1);
manager.add(uri2);
manager.add(uri3);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(3);
// Simulate deleting multiple files
const deleteHandler = vi.mocked(vscode.workspace.onDidDeleteFiles).mock
.calls[0][0];
deleteHandler({ files: [uri1, uri3] });
onDidDeleteFilesListener({ files: [uri1, uri3] });
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file2.txt');
});
it('prunes files older than the max age', () => {
vi.useFakeTimers();
const manager = new RecentFilesManager(context);
const uri1 = vscode.Uri.file('/test/file1.txt');
const uri1 = getUri('/test/file1.txt');
manager.add(uri1);
// Advance time by more than the max age
const twoMinutesMs = (MAX_FILE_AGE_MINUTES + 1) * 60 * 1000;
vi.advanceTimersByTime(twoMinutesMs);
const uri2 = vscode.Uri.file('/test/file2.txt');
const uri2 = getUri('/test/file2.txt');
manager.add(uri2);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file2.txt');
vi.useRealTimers();
});
it('fires onDidChange only once when adding an existing file', () => {
it('fires onDidChange only once when adding an existing file', async () => {
const manager = new RecentFilesManager(context);
const uri = vscode.Uri.file('/test/file1.txt');
const uri = getUri('/test/file1.txt');
manager.add(uri);
await vi.advanceTimersByTimeAsync(100);
const onDidChangeSpy = vi.fn();
manager.onDidChange(onDidChangeSpy);
const spy = vi.spyOn(manager['onDidChangeEmitter'], 'fire');
manager.add(uri);
expect(spy).toHaveBeenCalledTimes(1);
await vi.advanceTimersByTimeAsync(100);
expect(onDidChangeSpy).toHaveBeenCalledTimes(1);
});
it('updates the file when it is renamed', () => {
it('updates the file when it is renamed', async () => {
const manager = new RecentFilesManager(context);
const oldUri = vscode.Uri.file('/test/file1.txt');
const newUri = vscode.Uri.file('/test/file2.txt');
const oldUri = getUri('/test/file1.txt');
const newUri = getUri('/test/file2.txt');
manager.add(oldUri);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file1.txt');
// Simulate renaming the file
const renameHandler = vi.mocked(vscode.workspace.onDidRenameFiles).mock
.calls[0][0];
renameHandler({ files: [{ oldUri, newUri }] });
onDidRenameFilesListener({ files: [{ oldUri, newUri }] });
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file2.txt');
});
it('adds a file when the active editor changes', async () => {
const manager = new RecentFilesManager(context);
const uri = getUri('/test/file1.txt');
onDidChangeActiveTextEditorListener({
document: { uri },
} as vscode.TextEditor);
await vi.advanceTimersByTimeAsync(100);
expect(manager.recentFiles).toHaveLength(1);
expect(manager.recentFiles[0].filePath).toBe('/test/file1.txt');
});
});

View File

@ -16,14 +16,14 @@ interface RecentFile {
/**
* Keeps track of the 10 most recently-opened files
* opened less than 5 ago. If a file is closed or deleted,
* it will be removed. If the length is maxxed out,
* the now-removed file will not be replaced by an older file.
* opened less than 5 min ago. If a file is closed or deleted,
* it will be removed. If the max length is reached, older files will get removed first.
*/
export class RecentFilesManager {
private readonly files: RecentFile[] = [];
private readonly onDidChangeEmitter = new vscode.EventEmitter<void>();
readonly onDidChange = this.onDidChangeEmitter.event;
private debounceTimer: NodeJS.Timeout | undefined;
constructor(private readonly context: vscode.ExtensionContext) {
const editorWatcher = vscode.window.onDidChangeActiveTextEditor(
@ -33,7 +33,7 @@ export class RecentFilesManager {
}
},
);
const fileWatcher = vscode.workspace.onDidDeleteFiles((event) => {
const deleteWatcher = vscode.workspace.onDidDeleteFiles((event) => {
for (const uri of event.files) {
this.remove(uri);
}
@ -49,12 +49,21 @@ export class RecentFilesManager {
});
context.subscriptions.push(
editorWatcher,
fileWatcher,
deleteWatcher,
closeWatcher,
renameWatcher,
);
}
private fireWithDebounce() {
if (this.debounceTimer) {
clearTimeout(this.debounceTimer);
}
this.debounceTimer = setTimeout(() => {
this.onDidChangeEmitter.fire();
}, 50); // 50ms
}
private remove(uri: vscode.Uri, fireEvent = true) {
const index = this.files.findIndex(
(file) => file.uri.fsPath === uri.fsPath,
@ -62,21 +71,23 @@ export class RecentFilesManager {
if (index !== -1) {
this.files.splice(index, 1);
if (fireEvent) {
this.onDidChangeEmitter.fire();
this.fireWithDebounce();
}
}
}
add(uri: vscode.Uri) {
// Remove if it already exists to avoid duplicates and move it to the top.
this.remove(uri, false);
if (uri.scheme !== 'file') {
return;
}
this.remove(uri, false);
this.files.unshift({ uri, timestamp: Date.now() });
if (this.files.length > MAX_FILES) {
this.files.pop();
}
this.onDidChangeEmitter.fire();
this.fireWithDebounce();
}
get recentFiles(): Array<{ filePath: string; timestamp: number }> {

View File

@ -0,0 +1,18 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import * as vscode from 'vscode';
export function createLogger(
context: vscode.ExtensionContext,
logger: vscode.OutputChannel,
) {
return (message: string) => {
if (context.extensionMode === vscode.ExtensionMode.Development) {
logger.appendLine(message);
}
};
}