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

3 files changed, 316 insertions, 153 deletions

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,
   };
 }