Updates schema, UX and prompt for IDE context (#5046)

5 files changed, 535 insertions, 179 deletions

diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 25ea9bc1..8c46d7f5 100644
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -647,14 +647,140 @@ 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.mocked(ideContext.getIdeContext).mockReturnValue({
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/active/file.ts',
+              timestamp: Date.now(),
+              isActive: true,
+              selectedText: 'hello',
+              cursor: { line: 5, character: 10 },
+            },
+            {
+              path: '/path/to/recent/file1.ts',
+              timestamp: Date.now(),
+            },
+            {
+              path: '/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.getIdeContext).toHaveBeenCalled();
+      const expectedContext = `
+This is the file that the user is 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 file:
+- hello
+Here are some other files the user has open, 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 not add context if ideMode is enabled but no open files', async () => {
+      // Arrange
+      vi.mocked(ideContext.getIdeContext).mockReturnValue({
+        workspaceState: {
+          openFiles: [],
+        },
+      });
+
+      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.getIdeContext).toHaveBeenCalled();
+      expect(mockTurnRunFn).toHaveBeenCalledWith(
+        initialRequest,
+        expect.any(Object),
+      );
+    });
+
+    it('should add context if ideMode is enabled and there is one active file', async () => {
+      // Arrange
+      vi.mocked(ideContext.getIdeContext).mockReturnValue({
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/active/file.ts',
+              timestamp: Date.now(),
+              isActive: true,
+              selectedText: 'hello',
+              cursor: { line: 5, character: 10 },
+            },
+          ],
+        },
       });
 
       vi.spyOn(client['config'], 'getIdeMode').mockReturnValue(true);
@@ -689,15 +815,74 @@ describe('Gemini Client (client.ts)', () => {
       }
 
       // Assert
-      expect(ideContext.getOpenFilesContext).toHaveBeenCalled();
+      expect(ideContext.getIdeContext).toHaveBeenCalled();
       const expectedContext = `
-This is the file that the user was most recently looking at:
+This is the file that the user is 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:
+This is the selected text in the file:
 - hello
-Here are files the user has recently opened, with the most recent at the top:
+      `.trim();
+      const expectedRequest = [{ text: expectedContext }, ...initialRequest];
+      expect(mockTurnRunFn).toHaveBeenCalledWith(
+        expectedRequest,
+        expect.any(Object),
+      );
+    });
+
+    it('should add context if ideMode is enabled and there are open files but no active file', async () => {
+      // Arrange
+      vi.mocked(ideContext.getIdeContext).mockReturnValue({
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/recent/file1.ts',
+              timestamp: Date.now(),
+            },
+            {
+              path: '/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.getIdeContext).toHaveBeenCalled();
+      const expectedContext = `
+Here are some files the user has open, with the most recent at the top:
 - /path/to/recent/file1.ts
 - /path/to/recent/file2.ts
       `.trim();
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index 77683a45..e58e7040 100644
--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -320,32 +320,40 @@ export class GeminiClient {
     }
 
     if (this.config.getIdeMode()) {
-      const openFiles = ideContext.getOpenFilesContext();
-      if (openFiles) {
+      const ideContextState = ideContext.getIdeContext();
+      const openFiles = ideContextState?.workspaceState?.openFiles;
+
+      if (openFiles && openFiles.length > 0) {
         const contextParts: string[] = [];
-        if (openFiles.activeFile) {
+        const firstFile = openFiles[0];
+        const activeFile = firstFile.isActive ? firstFile : undefined;
+
+        if (activeFile) {
           contextParts.push(
-            `This is the file that the user was most recently looking at:\n- Path: ${openFiles.activeFile}`,
+            `This is the file that the user is looking at:\n- Path: ${activeFile.path}`,
           );
-          if (openFiles.cursor) {
+          if (activeFile.cursor) {
             contextParts.push(
-              `This is the cursor position in the file:\n- Cursor Position: Line ${openFiles.cursor.line}, Character ${openFiles.cursor.character}`,
+              `This is the cursor position in the file:\n- Cursor Position: Line ${activeFile.cursor.line}, Character ${activeFile.cursor.character}`,
             );
           }
-          if (openFiles.selectedText) {
+          if (activeFile.selectedText) {
             contextParts.push(
-              `This is the selected text in the active file:\n- ${openFiles.selectedText}`,
+              `This is the selected text in the file:\n- ${activeFile.selectedText}`,
             );
           }
         }
 
-        if (openFiles.recentOpenFiles && openFiles.recentOpenFiles.length > 0) {
-          const recentFiles = openFiles.recentOpenFiles
-            .map((file) => `- ${file.filePath}`)
+        const otherOpenFiles = activeFile ? openFiles.slice(1) : openFiles;
+
+        if (otherOpenFiles.length > 0) {
+          const recentFiles = otherOpenFiles
+            .map((file) => `- ${file.path}`)
             .join('\n');
-          contextParts.push(
-            `Here are files the user has recently opened, with the most recent at the top:\n${recentFiles}`,
-          );
+          const heading = activeFile
+            ? `Here are some other files the user has open, with the most recent at the top:`
+            : `Here are some files the user has open, with the most recent at the top:`;
+          contextParts.push(`${heading}\n${recentFiles}`);
         }
 
         if (contextParts.length > 0) {
diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts
index 3f91f386..64264fd1 100644
--- a/packages/core/src/ide/ide-client.ts
+++ b/packages/core/src/ide/ide-client.ts
@@ -4,7 +4,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { ideContext, OpenFilesNotificationSchema } from '../ide/ideContext.js';
+import { ideContext, IdeContextNotificationSchema } from '../ide/ideContext.js';
 import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
 
@@ -77,20 +77,20 @@ export class IdeClient {
       await this.client.connect(transport);
 
       this.client.setNotificationHandler(
-        OpenFilesNotificationSchema,
+        IdeContextNotificationSchema,
         (notification) => {
-          ideContext.setOpenFilesContext(notification.params);
+          ideContext.setIdeContext(notification.params);
         },
       );
       this.client.onerror = (error) => {
         logger.debug('IDE MCP client error:', error);
         this.connectionStatus = IDEConnectionStatus.Disconnected;
-        ideContext.clearOpenFilesContext();
+        ideContext.clearIdeContext();
       };
       this.client.onclose = () => {
         logger.debug('IDE MCP client connection closed.');
         this.connectionStatus = IDEConnectionStatus.Disconnected;
-        ideContext.clearOpenFilesContext();
+        ideContext.clearIdeContext();
       };
 
       this.connectionStatus = IDEConnectionStatus.Connected;
diff --git a/packages/core/src/ide/ideContext.test.ts b/packages/core/src/ide/ideContext.test.ts
index 1cb09c53..7e01d3aa 100644
--- a/packages/core/src/ide/ideContext.test.ts
+++ b/packages/core/src/ide/ideContext.test.ts
@@ -5,136 +5,300 @@
  */
 
 import { describe, it, expect, beforeEach, vi } from 'vitest';
-import { createIdeContextStore } from './ideContext.js';
+import {
+  createIdeContextStore,
+  FileSchema,
+  IdeContextSchema,
+} from './ideContext.js';
 
-describe('ideContext - Active File', () => {
-  let ideContext: ReturnType<typeof createIdeContextStore>;
+describe('ideContext', () => {
+  describe('createIdeContextStore', () => {
+    let ideContext: ReturnType<typeof createIdeContextStore>;
 
-  beforeEach(() => {
-    // Create a fresh, isolated instance for each test
-    ideContext = createIdeContextStore();
-  });
+    beforeEach(() => {
+      // Create a fresh, isolated instance for each test
+      ideContext = createIdeContextStore();
+    });
 
-  it('should return undefined initially for active file context', () => {
-    expect(ideContext.getOpenFilesContext()).toBeUndefined();
-  });
+    it('should return undefined initially for ide context', () => {
+      expect(ideContext.getIdeContext()).toBeUndefined();
+    });
 
-  it('should set and retrieve the active file context', () => {
-    const testFile = {
-      activeFile: '/path/to/test/file.ts',
-      selectedText: '1234',
-    };
+    it('should set and retrieve the ide context', () => {
+      const testFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/test/file.ts',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      };
 
-    ideContext.setOpenFilesContext(testFile);
+      ideContext.setIdeContext(testFile);
 
-    const activeFile = ideContext.getOpenFilesContext();
-    expect(activeFile).toEqual(testFile);
-  });
+      const activeFile = ideContext.getIdeContext();
+      expect(activeFile).toEqual(testFile);
+    });
 
-  it('should update the active file context when called multiple times', () => {
-    const firstFile = {
-      activeFile: '/path/to/first.js',
-      selectedText: '1234',
-    };
-    ideContext.setOpenFilesContext(firstFile);
+    it('should update the ide context when called multiple times', () => {
+      const firstFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/first.js',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      };
+      ideContext.setIdeContext(firstFile);
 
-    const secondFile = {
-      activeFile: '/path/to/second.py',
-      cursor: { line: 20, character: 30 },
-    };
-    ideContext.setOpenFilesContext(secondFile);
+      const secondFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/second.py',
+              isActive: true,
+              cursor: { line: 20, character: 30 },
+              timestamp: 0,
+            },
+          ],
+        },
+      };
+      ideContext.setIdeContext(secondFile);
 
-    const activeFile = ideContext.getOpenFilesContext();
-    expect(activeFile).toEqual(secondFile);
-  });
+      const activeFile = ideContext.getIdeContext();
+      expect(activeFile).toEqual(secondFile);
+    });
 
-  it('should handle empty string for file path', () => {
-    const testFile = {
-      activeFile: '',
-      selectedText: '1234',
-    };
-    ideContext.setOpenFilesContext(testFile);
-    expect(ideContext.getOpenFilesContext()).toEqual(testFile);
-  });
+    it('should handle empty string for file path', () => {
+      const testFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      };
+      ideContext.setIdeContext(testFile);
+      expect(ideContext.getIdeContext()).toEqual(testFile);
+    });
 
-  it('should notify subscribers when active file context changes', () => {
-    const subscriber1 = vi.fn();
-    const subscriber2 = vi.fn();
+    it('should notify subscribers when ide context changes', () => {
+      const subscriber1 = vi.fn();
+      const subscriber2 = vi.fn();
 
-    ideContext.subscribeToOpenFiles(subscriber1);
-    ideContext.subscribeToOpenFiles(subscriber2);
+      ideContext.subscribeToIdeContext(subscriber1);
+      ideContext.subscribeToIdeContext(subscriber2);
 
-    const testFile = {
-      activeFile: '/path/to/subscribed.ts',
-      cursor: { line: 15, character: 25 },
-    };
-    ideContext.setOpenFilesContext(testFile);
+      const testFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/subscribed.ts',
+              isActive: true,
+              cursor: { line: 15, character: 25 },
+              timestamp: 0,
+            },
+          ],
+        },
+      };
+      ideContext.setIdeContext(testFile);
 
-    expect(subscriber1).toHaveBeenCalledTimes(1);
-    expect(subscriber1).toHaveBeenCalledWith(testFile);
-    expect(subscriber2).toHaveBeenCalledTimes(1);
-    expect(subscriber2).toHaveBeenCalledWith(testFile);
+      expect(subscriber1).toHaveBeenCalledTimes(1);
+      expect(subscriber1).toHaveBeenCalledWith(testFile);
+      expect(subscriber2).toHaveBeenCalledTimes(1);
+      expect(subscriber2).toHaveBeenCalledWith(testFile);
 
-    // Test with another update
-    const newFile = {
-      activeFile: '/path/to/new.js',
-      selectedText: '1234',
-    };
-    ideContext.setOpenFilesContext(newFile);
+      // Test with another update
+      const newFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/new.js',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      };
+      ideContext.setIdeContext(newFile);
 
-    expect(subscriber1).toHaveBeenCalledTimes(2);
-    expect(subscriber1).toHaveBeenCalledWith(newFile);
-    expect(subscriber2).toHaveBeenCalledTimes(2);
-    expect(subscriber2).toHaveBeenCalledWith(newFile);
-  });
+      expect(subscriber1).toHaveBeenCalledTimes(2);
+      expect(subscriber1).toHaveBeenCalledWith(newFile);
+      expect(subscriber2).toHaveBeenCalledTimes(2);
+      expect(subscriber2).toHaveBeenCalledWith(newFile);
+    });
+
+    it('should stop notifying a subscriber after unsubscribe', () => {
+      const subscriber1 = vi.fn();
+      const subscriber2 = vi.fn();
+
+      const unsubscribe1 = ideContext.subscribeToIdeContext(subscriber1);
+      ideContext.subscribeToIdeContext(subscriber2);
 
-  it('should stop notifying a subscriber after unsubscribe', () => {
-    const subscriber1 = vi.fn();
-    const subscriber2 = vi.fn();
+      ideContext.setIdeContext({
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/file1.txt',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      });
+      expect(subscriber1).toHaveBeenCalledTimes(1);
+      expect(subscriber2).toHaveBeenCalledTimes(1);
 
-    const unsubscribe1 = ideContext.subscribeToOpenFiles(subscriber1);
-    ideContext.subscribeToOpenFiles(subscriber2);
+      unsubscribe1();
 
-    ideContext.setOpenFilesContext({
-      activeFile: '/path/to/file1.txt',
-      selectedText: '1234',
+      ideContext.setIdeContext({
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/file2.txt',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      });
+      expect(subscriber1).toHaveBeenCalledTimes(1); // Should not be called again
+      expect(subscriber2).toHaveBeenCalledTimes(2);
     });
-    expect(subscriber1).toHaveBeenCalledTimes(1);
-    expect(subscriber2).toHaveBeenCalledTimes(1);
 
-    unsubscribe1();
+    it('should clear the ide context', () => {
+      const testFile = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/test/file.ts',
+              isActive: true,
+              selectedText: '1234',
+              timestamp: 0,
+            },
+          ],
+        },
+      };
 
-    ideContext.setOpenFilesContext({
-      activeFile: '/path/to/file2.txt',
-      selectedText: '1234',
+      ideContext.setIdeContext(testFile);
+
+      expect(ideContext.getIdeContext()).toEqual(testFile);
+
+      ideContext.clearIdeContext();
+
+      expect(ideContext.getIdeContext()).toBeUndefined();
     });
-    expect(subscriber1).toHaveBeenCalledTimes(1); // Should not be called again
-    expect(subscriber2).toHaveBeenCalledTimes(2);
   });
 
-  it('should allow the cursor to be optional', () => {
-    const testFile = {
-      activeFile: '/path/to/test/file.ts',
-    };
+  describe('FileSchema', () => {
+    it('should validate a file with only required fields', () => {
+      const file = {
+        path: '/path/to/file.ts',
+        timestamp: 12345,
+      };
+      const result = FileSchema.safeParse(file);
+      expect(result.success).toBe(true);
+    });
 
-    ideContext.setOpenFilesContext(testFile);
+    it('should validate a file with all fields', () => {
+      const file = {
+        path: '/path/to/file.ts',
+        timestamp: 12345,
+        isActive: true,
+        selectedText: 'const x = 1;',
+        cursor: {
+          line: 10,
+          character: 20,
+        },
+      };
+      const result = FileSchema.safeParse(file);
+      expect(result.success).toBe(true);
+    });
+
+    it('should fail validation if path is missing', () => {
+      const file = {
+        timestamp: 12345,
+      };
+      const result = FileSchema.safeParse(file);
+      expect(result.success).toBe(false);
+    });
 
-    const activeFile = ideContext.getOpenFilesContext();
-    expect(activeFile).toEqual(testFile);
+    it('should fail validation if timestamp is missing', () => {
+      const file = {
+        path: '/path/to/file.ts',
+      };
+      const result = FileSchema.safeParse(file);
+      expect(result.success).toBe(false);
+    });
   });
 
-  it('should clear the active file context', () => {
-    const testFile = {
-      activeFile: '/path/to/test/file.ts',
-      selectedText: '1234',
-    };
+  describe('IdeContextSchema', () => {
+    it('should validate an empty context', () => {
+      const context = {};
+      const result = IdeContextSchema.safeParse(context);
+      expect(result.success).toBe(true);
+    });
 
-    ideContext.setOpenFilesContext(testFile);
+    it('should validate a context with an empty workspaceState', () => {
+      const context = {
+        workspaceState: {},
+      };
+      const result = IdeContextSchema.safeParse(context);
+      expect(result.success).toBe(true);
+    });
 
-    expect(ideContext.getOpenFilesContext()).toEqual(testFile);
+    it('should validate a context with an empty openFiles array', () => {
+      const context = {
+        workspaceState: {
+          openFiles: [],
+        },
+      };
+      const result = IdeContextSchema.safeParse(context);
+      expect(result.success).toBe(true);
+    });
 
-    ideContext.clearOpenFilesContext();
+    it('should validate a context with a valid file', () => {
+      const context = {
+        workspaceState: {
+          openFiles: [
+            {
+              path: '/path/to/file.ts',
+              timestamp: 12345,
+            },
+          ],
+        },
+      };
+      const result = IdeContextSchema.safeParse(context);
+      expect(result.success).toBe(true);
+    });
 
-    expect(ideContext.getOpenFilesContext()).toBeUndefined();
+    it('should fail validation with an invalid file', () => {
+      const context = {
+        workspaceState: {
+          openFiles: [
+            {
+              timestamp: 12345, // path is missing
+            },
+          ],
+        },
+      };
+      const result = IdeContextSchema.safeParse(context);
+      expect(result.success).toBe(false);
+    });
   });
 });
diff --git a/packages/core/src/ide/ideContext.ts b/packages/core/src/ide/ideContext.ts
index bc7383a1..588e25ee 100644
--- a/packages/core/src/ide/ideContext.ts
+++ b/packages/core/src/ide/ideContext.ts
@@ -7,97 +7,96 @@
 import { z } from 'zod';
 
 /**
- * Zod schema for validating a cursor position.
+ * Zod schema for validating a file context from the IDE.
  */
-export const CursorSchema = z.object({
-  line: z.number(),
-  character: z.number(),
+export const FileSchema = z.object({
+  path: z.string(),
+  timestamp: z.number(),
+  isActive: z.boolean().optional(),
+  selectedText: z.string().optional(),
+  cursor: z
+    .object({
+      line: z.number(),
+      character: z.number(),
+    })
+    .optional(),
 });
-export type Cursor = z.infer<typeof CursorSchema>;
+export type File = z.infer<typeof FileSchema>;
 
-/**
- * Zod schema for validating an active file context from the IDE.
- */
-export const OpenFilesSchema = z.object({
-  activeFile: z.string(),
-  selectedText: z.string().optional(),
-  cursor: CursorSchema.optional(),
-  recentOpenFiles: z
-    .array(
-      z.object({
-        filePath: z.string(),
-        timestamp: z.number(),
-      }),
-    )
+export const IdeContextSchema = z.object({
+  workspaceState: z
+    .object({
+      openFiles: z.array(FileSchema).optional(),
+    })
     .optional(),
 });
-export type OpenFiles = z.infer<typeof OpenFilesSchema>;
+export type IdeContext = z.infer<typeof IdeContextSchema>;
 
 /**
- * Zod schema for validating the 'ide/openFilesChanged' notification from the IDE.
+ * Zod schema for validating the 'ide/contextUpdate' notification from the IDE.
  */
-export const OpenFilesNotificationSchema = z.object({
-  method: z.literal('ide/openFilesChanged'),
-  params: OpenFilesSchema,
+export const IdeContextNotificationSchema = z.object({
+  method: z.literal('ide/contextUpdate'),
+  params: IdeContextSchema,
 });
 
-type OpenFilesSubscriber = (openFiles: OpenFiles | undefined) => void;
+type IdeContextSubscriber = (ideContext: IdeContext | undefined) => void;
 
 /**
- * Creates a new store for managing the IDE's active file context.
+ * Creates a new store for managing the IDE's context.
  * This factory function encapsulates the state and logic, allowing for the creation
  * of isolated instances, which is particularly useful for testing.
  *
- * @returns An object with methods to interact with the active file context.
+ * @returns An object with methods to interact with the IDE context.
  */
 export function createIdeContextStore() {
-  let openFilesContext: OpenFiles | undefined = undefined;
-  const subscribers = new Set<OpenFilesSubscriber>();
+  let ideContextState: IdeContext | undefined = undefined;
+  const subscribers = new Set<IdeContextSubscriber>();
 
   /**
-   * Notifies all registered subscribers about the current active file context.
+   * Notifies all registered subscribers about the current IDE context.
    */
   function notifySubscribers(): void {
     for (const subscriber of subscribers) {
-      subscriber(openFilesContext);
+      subscriber(ideContextState);
     }
   }
 
   /**
-   * Sets the active file context and notifies all registered subscribers of the change.
-   * @param newOpenFiles The new active file context from the IDE.
+   * Sets the IDE context and notifies all registered subscribers of the change.
+   * @param newIdeContext The new IDE context from the IDE.
    */
-  function setOpenFilesContext(newOpenFiles: OpenFiles): void {
-    openFilesContext = newOpenFiles;
+  function setIdeContext(newIdeContext: IdeContext): void {
+    ideContextState = newIdeContext;
     notifySubscribers();
   }
 
   /**
-   * Clears the active file context and notifies all registered subscribers of the change.
+   * Clears the IDE context and notifies all registered subscribers of the change.
    */
-  function clearOpenFilesContext(): void {
-    openFilesContext = undefined;
+  function clearIdeContext(): void {
+    ideContextState = undefined;
     notifySubscribers();
   }
 
   /**
-   * Retrieves the current active file context.
-   * @returns The `OpenFiles` object if a file is active; otherwise, `undefined`.
+   * Retrieves the current IDE context.
+   * @returns The `IdeContext` object if a file is active; otherwise, `undefined`.
    */
-  function getOpenFilesContext(): OpenFiles | undefined {
-    return openFilesContext;
+  function getIdeContext(): IdeContext | undefined {
+    return ideContextState;
   }
 
   /**
-   * Subscribes to changes in the active file context.
+   * Subscribes to changes in the IDE context.
    *
-   * When the active file context changes, the provided `subscriber` function will be called.
+   * When the IDE context changes, the provided `subscriber` function will be called.
    * Note: The subscriber is not called with the current value upon subscription.
    *
-   * @param subscriber The function to be called when the active file context changes.
+   * @param subscriber The function to be called when the IDE context changes.
    * @returns A function that, when called, will unsubscribe the provided subscriber.
    */
-  function subscribeToOpenFiles(subscriber: OpenFilesSubscriber): () => void {
+  function subscribeToIdeContext(subscriber: IdeContextSubscriber): () => void {
     subscribers.add(subscriber);
     return () => {
       subscribers.delete(subscriber);
@@ -105,10 +104,10 @@ export function createIdeContextStore() {
   }
 
   return {
-    setOpenFilesContext,
-    getOpenFilesContext,
-    subscribeToOpenFiles,
-    clearOpenFilesContext,
+    setIdeContext,
+    getIdeContext,
+    subscribeToIdeContext,
+    clearIdeContext,
   };
 }