blackbox-cli-vscode-ide-companion 0.0.5

package/src/ide-server.test.ts

@@ -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 type * as os from 'node:os';
+import * as path from 'node:path';
+import { IDEServer } from './ide-server.js';
+import type { 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] === 'BLACKBOX_CODE_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,
+      'BLACKBOX_CODE_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,
+      'BLACKBOX_CODE_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(
+      'BLACKBOX_CODE_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(
+      'BLACKBOX_CODE_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(
+      'BLACKBOX_CODE_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(
+      'BLACKBOX_CODE_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(
+      'BLACKBOX_CODE_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(
+        'BLACKBOX_CODE_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,
+        }),
+      );
+    },
+  );
+});

package/src/ide-server.ts

@@ -0,0 +1,345 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import * as vscode from 'vscode';
+import { IdeContextNotificationSchema } from '@blackbox_ai/blackbox-cli-core';
+import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';
+import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
+import express, { type Request, type Response } from 'express';
+import { randomUUID } from 'node:crypto';
+import { type Server as HTTPServer } from 'node:http';
+import * as path from 'node:path';
+import * as fs from 'node:fs/promises';
+import * as os from 'node:os';
+import { z } from 'zod';
+import type { DiffManager } from './diff-manager.js';
+import { OpenFilesManager } from './open-files-manager.js';
+
+const MCP_SESSION_ID_HEADER = 'mcp-session-id';
+const IDE_SERVER_PORT_ENV_VAR = 'BLACKBOX_CODE_IDE_SERVER_PORT';
+const IDE_WORKSPACE_PATH_ENV_VAR = 'BLACKBOX_CODE_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,
+  log: (message: string) => void,
+  openFilesManager: OpenFilesManager,
+) {
+  const ideContext = openFilesManager.state;
+
+  const notification = IdeContextNotificationSchema.parse({
+    jsonrpc: '2.0',
+    method: 'ide/contextUpdate',
+    params: ideContext,
+  });
+
+  log(
+    `Sending IDE context update notification: ${JSON.stringify(
+      notification,
+      null,
+      2,
+    )}`,
+  );
+  transport.send(notification);
+}
+
+export class IDEServer {
+  private server: HTTPServer | undefined;
+  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) {
+    this.log = log;
+    this.diffManager = diffManager;
+    this.portFile = path.join(
+      os.tmpdir(),
+      `gemini-ide-server-${process.ppid}.json`,
+    );
+  }
+
+  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 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) => {
+          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;
+
+        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);
+            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,
+              error: {
+                code: -32603,
+                message: 'Internal server error',
+              },
+              id: null,
+            });
+          }
+        }
+      });
+
+      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');
+          }
+        }
+
+        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) => {
+        this.server!.close((err?: Error) => {
+          if (err) {
+            this.log(`Error shutting down IDE server: ${err.message}`);
+            return reject(err);
+          }
+          this.log(`IDE server shut down`);
+          resolve();
+        });
+      });
+      this.server = undefined;
+    }
+
+    if (this.context) {
+      this.context.environmentVariableCollection.clear();
+    }
+    try {
+      await fs.unlink(this.portFile);
+    } catch (_err) {
+      // Ignore errors if the file doesn't exist.
+    }
+  }
+}
+
+const createMcpServer = (diffManager: DiffManager) => {
+  const server = new McpServer(
+    {
+      name: 'blackbox-cli-companion-mcp-server',
+      version: '1.0.0',
+    },
+    { capabilities: { logging: {} } },
+  );
+  server.registerTool(
+    'openDiff',
+    {
+      description:
+        '(IDE Tool) Open a diff view to create or modify a file. Returns a notification once the diff has been accepted or rejcted.',
+      inputSchema: z.object({
+        filePath: z.string(),
+        // TODO(chrstn): determine if this should be required or not.
+        newContent: z.string().optional(),
+      }).shape,
+    },
+    async ({
+      filePath,
+      newContent,
+    }: {
+      filePath: string;
+      newContent?: string;
+    }) => {
+      await diffManager.showDiff(filePath, newContent ?? '');
+      return {
+        content: [
+          {
+            type: 'text',
+            text: `Showing diff for ${filePath}`,
+          },
+        ],
+      };
+    },
+  );
+  server.registerTool(
+    'closeDiff',
+    {
+      description: '(IDE Tool) Close an open diff view for a specific file.',
+      inputSchema: z.object({
+        filePath: z.string(),
+      }).shape,
+    },
+    async ({ filePath }: { filePath: string }) => {
+      const content = await diffManager.closeDiff(filePath);
+      const response = { content: content ?? undefined };
+      return {
+        content: [
+          {
+            type: 'text',
+            text: JSON.stringify(response),
+          },
+        ],
+      };
+    },
+  );
+  return server;
+};