diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index ceab46b1..0712d810 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -151,8 +151,8 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }); const ideContextMock = { - getActiveFileContext: vi.fn(), - subscribeToActiveFile: vi.fn(() => vi.fn()), // subscribe returns an unsubscribe function + getOpenFilesContext: vi.fn(), + subscribeToOpenFiles: vi.fn(() => vi.fn()), // subscribe returns an unsubscribe function }; return { @@ -267,7 +267,7 @@ describe('App UI', () => { // Ensure a theme is set so the theme dialog does not appear. mockSettings = createMockSettings({ workspace: { theme: 'Default' } }); - vi.mocked(ideContext.getActiveFileContext).mockReturnValue(undefined); + vi.mocked(ideContext.getOpenFilesContext).mockReturnValue(undefined); }); afterEach(() => { @@ -279,10 +279,9 @@ describe('App UI', () => { }); it('should display active file when available', async () => { - vi.mocked(ideContext.getActiveFileContext).mockReturnValue({ - filePath: '/path/to/my-file.ts', - content: 'const a = 1;', - cursor: 0, + vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ + activeFile: '/path/to/my-file.ts', + selectedText: 'hello', }); const { lastFrame, unmount } = render( @@ -298,10 +297,8 @@ describe('App UI', () => { }); it('should not display active file when not available', async () => { - vi.mocked(ideContext.getActiveFileContext).mockReturnValue({ - filePath: '', - content: '', - cursor: 0, + vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ + activeFile: '', }); const { lastFrame, unmount } = render( @@ -317,10 +314,9 @@ describe('App UI', () => { }); it('should display active file and other context', async () => { - vi.mocked(ideContext.getActiveFileContext).mockReturnValue({ - filePath: '/path/to/my-file.ts', - content: 'const a = 1;', - cursor: 0, + vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ + activeFile: '/path/to/my-file.ts', + selectedText: 'hello', }); mockConfig.getGeminiMdFileCount.mockReturnValue(1); diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index c6e6bd43..12838710 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -58,7 +58,7 @@ import { FlashFallbackEvent, logFlashFallback, AuthType, - type ActiveFile, + type OpenFiles, ideContext, } from '@google/gemini-cli-core'; import { validateAuthMethod } from '../config/auth.js'; @@ -160,12 +160,12 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { const [modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError] = useState(false); const [userTier, setUserTier] = useState(undefined); - const [activeFile, setActiveFile] = useState(); + const [openFiles, setOpenFiles] = useState(); useEffect(() => { - const unsubscribe = ideContext.subscribeToActiveFile(setActiveFile); + const unsubscribe = ideContext.subscribeToOpenFiles(setOpenFiles); // Set the initial value - setActiveFile(ideContext.getActiveFileContext()); + setOpenFiles(ideContext.getOpenFilesContext()); return unsubscribe; }, []); @@ -880,7 +880,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { ) : ( ; blockedMcpServers?: Array<{ name: string; extensionName: string }>; showToolDescriptions?: boolean; - activeFile?: ActiveFile; + openFiles?: OpenFiles; } export const ContextSummaryDisplay: React.FC = ({ @@ -25,7 +25,7 @@ export const ContextSummaryDisplay: React.FC = ({ mcpServers, blockedMcpServers, showToolDescriptions, - activeFile, + openFiles, }) => { const mcpServerCount = Object.keys(mcpServers || {}).length; const blockedMcpServerCount = blockedMcpServers?.length || 0; @@ -34,16 +34,16 @@ export const ContextSummaryDisplay: React.FC = ({ geminiMdFileCount === 0 && mcpServerCount === 0 && blockedMcpServerCount === 0 && - !activeFile?.filePath + !openFiles?.activeFile ) { return ; // Render an empty space to reserve height } const activeFileText = (() => { - if (!activeFile?.filePath) { + if (!openFiles?.activeFile) { return ''; } - return `Open File (${path.basename(activeFile.filePath)})`; + return `Open File (${path.basename(openFiles.activeFile)})`; })(); const geminiMdText = (() => { diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index faeaa01e..2d75637c 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -23,6 +23,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { setSimulate429 } from '../utils/testUtils.js'; import { tokenLimit } from './tokenLimits.js'; +import { ideContext } from '../services/ideContext.js'; // --- Mocks --- const mockChatCreateFn = vi.fn(); @@ -71,6 +72,7 @@ vi.mock('../telemetry/index.js', () => ({ logApiResponse: vi.fn(), logApiError: vi.fn(), })); +vi.mock('../services/ideContext.js'); describe('findIndexAfterFraction', () => { const history: Content[] = [ @@ -642,6 +644,69 @@ describe('Gemini Client (client.ts)', () => { }); describe('sendMessageStream', () => { + it('should include IDE context when ideMode is enabled', async () => { + // Arrange + vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ + activeFile: '/path/to/active/file.ts', + selectedText: 'hello', + cursor: { line: 5, character: 10 }, + recentOpenFiles: [ + { filePath: '/path/to/recent/file1.ts', timestamp: Date.now() }, + { filePath: '/path/to/recent/file2.ts', timestamp: Date.now() }, + ], + }); + + vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true); + + const mockStream = (async function* () { + yield { type: 'content', value: 'Hello' }; + })(); + mockTurnRunFn.mockReturnValue(mockStream); + + const mockChat: Partial = { + addHistory: vi.fn(), + getHistory: vi.fn().mockReturnValue([]), + }; + client['chat'] = mockChat as GeminiChat; + + const mockGenerator: Partial = { + countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }), + generateContent: mockGenerateContentFn, + }; + client['contentGenerator'] = mockGenerator as ContentGenerator; + + const initialRequest = [{ text: 'Hi' }]; + + // Act + const stream = client.sendMessageStream( + initialRequest, + new AbortController().signal, + 'prompt-id-ide', + ); + for await (const _ of stream) { + // consume stream + } + + // Assert + expect(ideContext.getOpenFilesContext).toHaveBeenCalled(); + const expectedContext = ` +This is the file that the user was most recently looking at: +- Path: /path/to/active/file.ts +This is the cursor position in the file: +- Cursor Position: Line 5, Character 10 +This is the selected text in the active file: +- hello +Here are files the user has recently opened, with the most recent at the top: +- /path/to/recent/file1.ts +- /path/to/recent/file2.ts + `.trim(); + const expectedRequest = [{ text: expectedContext }, ...initialRequest]; + expect(mockTurnRunFn).toHaveBeenCalledWith( + expectedRequest, + expect.any(Object), + ); + }); + it('should return the turn instance after the stream is complete', async () => { // Arrange const mockStream = (async function* () { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 61195e2f..aadc446c 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -311,20 +311,40 @@ export class GeminiClient { } if (this.config.getIdeMode()) { - const activeFile = ideContext.getActiveFileContext(); - if (activeFile?.filePath) { - let context = ` -This is the file that the user was most recently looking at: -- Path: ${activeFile.filePath}`; - if (activeFile.cursor) { - context += ` -This is the cursor position in the file: -- Cursor Position: Line ${activeFile.cursor.line}, Character ${activeFile.cursor.character}`; + const openFiles = ideContext.getOpenFilesContext(); + if (openFiles) { + const contextParts: string[] = []; + if (openFiles.activeFile) { + contextParts.push( + `This is the file that the user was most recently looking at:\n- Path: ${openFiles.activeFile}`, + ); + if (openFiles.cursor) { + contextParts.push( + `This is the cursor position in the file:\n- Cursor Position: Line ${openFiles.cursor.line}, Character ${openFiles.cursor.character}`, + ); + } + if (openFiles.selectedText) { + contextParts.push( + `This is the selected text in the active file:\n- ${openFiles.selectedText}`, + ); + } + } + + if (openFiles.recentOpenFiles && openFiles.recentOpenFiles.length > 0) { + const recentFiles = openFiles.recentOpenFiles + .map((file) => `- ${file.filePath}`) + .join('\n'); + contextParts.push( + `Here are files the user has recently opened, with the most recent at the top:\n${recentFiles}`, + ); + } + + if (contextParts.length > 0) { + request = [ + { text: contextParts.join('\n') }, + ...(Array.isArray(request) ? request : [request]), + ]; } - request = [ - { text: context }, - ...(Array.isArray(request) ? request : [request]), - ]; } } diff --git a/packages/core/src/services/ideContext.test.ts b/packages/core/src/services/ideContext.test.ts index 9aa4c013..1cb09c53 100644 --- a/packages/core/src/services/ideContext.test.ts +++ b/packages/core/src/services/ideContext.test.ts @@ -16,59 +16,59 @@ describe('ideContext - Active File', () => { }); it('should return undefined initially for active file context', () => { - expect(ideContext.getActiveFileContext()).toBeUndefined(); + expect(ideContext.getOpenFilesContext()).toBeUndefined(); }); it('should set and retrieve the active file context', () => { const testFile = { - filePath: '/path/to/test/file.ts', - cursor: { line: 5, character: 10 }, + activeFile: '/path/to/test/file.ts', + selectedText: '1234', }; - ideContext.setActiveFileContext(testFile); + ideContext.setOpenFilesContext(testFile); - const activeFile = ideContext.getActiveFileContext(); + const activeFile = ideContext.getOpenFilesContext(); expect(activeFile).toEqual(testFile); }); it('should update the active file context when called multiple times', () => { const firstFile = { - filePath: '/path/to/first.js', - cursor: { line: 1, character: 1 }, + activeFile: '/path/to/first.js', + selectedText: '1234', }; - ideContext.setActiveFileContext(firstFile); + ideContext.setOpenFilesContext(firstFile); const secondFile = { - filePath: '/path/to/second.py', + activeFile: '/path/to/second.py', cursor: { line: 20, character: 30 }, }; - ideContext.setActiveFileContext(secondFile); + ideContext.setOpenFilesContext(secondFile); - const activeFile = ideContext.getActiveFileContext(); + const activeFile = ideContext.getOpenFilesContext(); expect(activeFile).toEqual(secondFile); }); it('should handle empty string for file path', () => { const testFile = { - filePath: '', - cursor: { line: 0, character: 0 }, + activeFile: '', + selectedText: '1234', }; - ideContext.setActiveFileContext(testFile); - expect(ideContext.getActiveFileContext()).toEqual(testFile); + ideContext.setOpenFilesContext(testFile); + expect(ideContext.getOpenFilesContext()).toEqual(testFile); }); it('should notify subscribers when active file context changes', () => { const subscriber1 = vi.fn(); const subscriber2 = vi.fn(); - ideContext.subscribeToActiveFile(subscriber1); - ideContext.subscribeToActiveFile(subscriber2); + ideContext.subscribeToOpenFiles(subscriber1); + ideContext.subscribeToOpenFiles(subscriber2); const testFile = { - filePath: '/path/to/subscribed.ts', + activeFile: '/path/to/subscribed.ts', cursor: { line: 15, character: 25 }, }; - ideContext.setActiveFileContext(testFile); + ideContext.setOpenFilesContext(testFile); expect(subscriber1).toHaveBeenCalledTimes(1); expect(subscriber1).toHaveBeenCalledWith(testFile); @@ -77,10 +77,10 @@ describe('ideContext - Active File', () => { // Test with another update const newFile = { - filePath: '/path/to/new.js', - cursor: { line: 1, character: 1 }, + activeFile: '/path/to/new.js', + selectedText: '1234', }; - ideContext.setActiveFileContext(newFile); + ideContext.setOpenFilesContext(newFile); expect(subscriber1).toHaveBeenCalledTimes(2); expect(subscriber1).toHaveBeenCalledWith(newFile); @@ -92,21 +92,21 @@ describe('ideContext - Active File', () => { const subscriber1 = vi.fn(); const subscriber2 = vi.fn(); - const unsubscribe1 = ideContext.subscribeToActiveFile(subscriber1); - ideContext.subscribeToActiveFile(subscriber2); + const unsubscribe1 = ideContext.subscribeToOpenFiles(subscriber1); + ideContext.subscribeToOpenFiles(subscriber2); - ideContext.setActiveFileContext({ - filePath: '/path/to/file1.txt', - cursor: { line: 1, character: 1 }, + ideContext.setOpenFilesContext({ + activeFile: '/path/to/file1.txt', + selectedText: '1234', }); expect(subscriber1).toHaveBeenCalledTimes(1); expect(subscriber2).toHaveBeenCalledTimes(1); unsubscribe1(); - ideContext.setActiveFileContext({ - filePath: '/path/to/file2.txt', - cursor: { line: 2, character: 2 }, + ideContext.setOpenFilesContext({ + activeFile: '/path/to/file2.txt', + selectedText: '1234', }); expect(subscriber1).toHaveBeenCalledTimes(1); // Should not be called again expect(subscriber2).toHaveBeenCalledTimes(2); @@ -114,27 +114,27 @@ describe('ideContext - Active File', () => { it('should allow the cursor to be optional', () => { const testFile = { - filePath: '/path/to/test/file.ts', + activeFile: '/path/to/test/file.ts', }; - ideContext.setActiveFileContext(testFile); + ideContext.setOpenFilesContext(testFile); - const activeFile = ideContext.getActiveFileContext(); + const activeFile = ideContext.getOpenFilesContext(); expect(activeFile).toEqual(testFile); }); it('should clear the active file context', () => { const testFile = { - filePath: '/path/to/test/file.ts', - cursor: { line: 5, character: 10 }, + activeFile: '/path/to/test/file.ts', + selectedText: '1234', }; - ideContext.setActiveFileContext(testFile); + ideContext.setOpenFilesContext(testFile); - expect(ideContext.getActiveFileContext()).toEqual(testFile); + expect(ideContext.getOpenFilesContext()).toEqual(testFile); - ideContext.clearActiveFileContext(); + ideContext.clearOpenFilesContext(); - expect(ideContext.getActiveFileContext()).toBeUndefined(); + expect(ideContext.getOpenFilesContext()).toBeUndefined(); }); }); diff --git a/packages/core/src/services/ideContext.ts b/packages/core/src/services/ideContext.ts index 0aab1e8d..349bff59 100644 --- a/packages/core/src/services/ideContext.ts +++ b/packages/core/src/services/ideContext.ts @@ -10,7 +10,6 @@ import { z } from 'zod'; * The reserved server name for the IDE's MCP server. */ export const IDE_SERVER_NAME = '_ide_server'; - /** * Zod schema for validating a cursor position. */ @@ -23,8 +22,9 @@ export type Cursor = z.infer; /** * Zod schema for validating an active file context from the IDE. */ -export const ActiveFileSchema = z.object({ - filePath: z.string(), +export const OpenFilesSchema = z.object({ + activeFile: z.string(), + selectedText: z.string().optional(), cursor: CursorSchema.optional(), recentOpenFiles: z .array( @@ -35,17 +35,17 @@ export const ActiveFileSchema = z.object({ ) .optional(), }); -export type ActiveFile = z.infer; +export type OpenFiles = z.infer; /** - * Zod schema for validating the 'ide/activeFileChanged' notification from the IDE. + * Zod schema for validating the 'ide/openFilesChanged' notification from the IDE. */ -export const ActiveFileNotificationSchema = z.object({ - method: z.literal('ide/activeFileChanged'), - params: ActiveFileSchema, +export const OpenFilesNotificationSchema = z.object({ + method: z.literal('ide/openFilesChanged'), + params: OpenFilesSchema, }); -type ActiveFileSubscriber = (activeFile: ActiveFile | undefined) => void; +type OpenFilesSubscriber = (openFiles: OpenFiles | undefined) => void; /** * Creates a new store for managing the IDE's active file context. @@ -55,41 +55,41 @@ type ActiveFileSubscriber = (activeFile: ActiveFile | undefined) => void; * @returns An object with methods to interact with the active file context. */ export function createIdeContextStore() { - let activeFileContext: ActiveFile | undefined = undefined; - const subscribers = new Set(); + let openFilesContext: OpenFiles | undefined = undefined; + const subscribers = new Set(); /** * Notifies all registered subscribers about the current active file context. */ function notifySubscribers(): void { for (const subscriber of subscribers) { - subscriber(activeFileContext); + subscriber(openFilesContext); } } /** * Sets the active file context and notifies all registered subscribers of the change. - * @param newActiveFile The new active file context from the IDE. + * @param newOpenFiles The new active file context from the IDE. */ - function setActiveFileContext(newActiveFile: ActiveFile): void { - activeFileContext = newActiveFile; + function setOpenFilesContext(newOpenFiles: OpenFiles): void { + openFilesContext = newOpenFiles; notifySubscribers(); } /** * Clears the active file context and notifies all registered subscribers of the change. */ - function clearActiveFileContext(): void { - activeFileContext = undefined; + function clearOpenFilesContext(): void { + openFilesContext = undefined; notifySubscribers(); } /** * Retrieves the current active file context. - * @returns The `ActiveFile` object if a file is active, otherwise `undefined`. + * @returns The `OpenFiles` object if a file is active, otherwise `undefined`. */ - function getActiveFileContext(): ActiveFile | undefined { - return activeFileContext; + function getOpenFilesContext(): OpenFiles | undefined { + return openFilesContext; } /** @@ -101,7 +101,7 @@ export function createIdeContextStore() { * @param subscriber The function to be called when the active file context changes. * @returns A function that, when called, will unsubscribe the provided subscriber. */ - function subscribeToActiveFile(subscriber: ActiveFileSubscriber): () => void { + function subscribeToOpenFiles(subscriber: OpenFilesSubscriber): () => void { subscribers.add(subscriber); return () => { subscribers.delete(subscriber); @@ -109,10 +109,10 @@ export function createIdeContextStore() { } return { - setActiveFileContext, - getActiveFileContext, - subscribeToActiveFile, - clearActiveFileContext, + setOpenFilesContext, + getOpenFilesContext, + subscribeToOpenFiles, + clearOpenFilesContext, }; } diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index cb191b0d..b1786af0 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -22,7 +22,7 @@ import { DiscoveredMCPTool } from './mcp-tool.js'; import { FunctionDeclaration, mcpToTool } from '@google/genai'; import { ToolRegistry } from './tool-registry.js'; import { - ActiveFileNotificationSchema, + OpenFilesNotificationSchema, IDE_SERVER_NAME, ideContext, } from '../services/ideContext.js'; @@ -217,15 +217,15 @@ export async function connectAndDiscover( console.error(`MCP ERROR (${mcpServerName}):`, error.toString()); updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED); if (mcpServerName === IDE_SERVER_NAME) { - ideContext.clearActiveFileContext(); + ideContext.clearOpenFilesContext(); } }; if (mcpServerName === IDE_SERVER_NAME) { mcpClient.setNotificationHandler( - ActiveFileNotificationSchema, + OpenFilesNotificationSchema, (notification) => { - ideContext.setActiveFileContext(notification.params); + ideContext.setOpenFilesContext(notification.params); }, ); } diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 37e07737..9111f349 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -19,7 +19,7 @@ import { RecentFilesManager } from './recent-files-manager.js'; const MCP_SESSION_ID_HEADER = 'mcp-session-id'; const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT'; -function sendActiveFileChangedNotification( +function sendOpenFilesChangedNotification( transport: StreamableHTTPServerTransport, logger: vscode.OutputChannel, recentFilesManager: RecentFilesManager, @@ -29,9 +29,9 @@ function sendActiveFileChangedNotification( logger.appendLine(`Sending active file changed notification: ${filePath}`); const notification: JSONRPCNotification = { jsonrpc: '2.0', - method: 'ide/activeFileChanged', + method: 'ide/openFilesChanged', params: { - filePath, + activeFile: filePath, recentOpenFiles: recentFilesManager.recentFiles, }, }; @@ -60,7 +60,7 @@ export class IDEServer { const recentFilesManager = new RecentFilesManager(context); const disposable = recentFilesManager.onDidChange(() => { for (const transport of Object.values(transports)) { - sendActiveFileChangedNotification( + sendOpenFilesChangedNotification( transport, this.logger, recentFilesManager, @@ -168,7 +168,7 @@ export class IDEServer { } if (!sessionsWithInitialNotification.has(sessionId)) { - sendActiveFileChangedNotification( + sendOpenFilesChangedNotification( transport, this.logger, recentFilesManager, @@ -224,7 +224,7 @@ const createMcpServer = () => { { capabilities: { logging: {} } }, ); server.registerTool( - 'getActiveFile', + 'getOpenFiles', { description: '(IDE Tool) Get the path of the file currently active in VS Code.',