[ide-mode] Use active files and selected text in user prompt (#4614)

5 files changed, 166 insertions, 81 deletions

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<GeminiChat> = {
+        addHistory: vi.fn(),
+        getHistory: vi.fn().mockReturnValue([]),
+      };
+      client['chat'] = mockChat as GeminiChat;
+
+      const mockGenerator: Partial<ContentGenerator> = {
+        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<typeof CursorSchema>;
 /**
  * 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<typeof ActiveFileSchema>;
+export type OpenFiles = z.infer<typeof OpenFilesSchema>;
 
 /**
- * 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<ActiveFileSubscriber>();
+  let openFilesContext: OpenFiles | undefined = undefined;
+  const subscribers = new Set<OpenFilesSubscriber>();
 
   /**
    * 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);
           },
         );
       }