diff --git a/packages/cli/src/ui/App.test.tsx b/packages/cli/src/ui/App.test.tsx index 93230d1c..f35f8cb7 100644 --- a/packages/cli/src/ui/App.test.tsx +++ b/packages/cli/src/ui/App.test.tsx @@ -153,8 +153,8 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => { }); const ideContextMock = { - getOpenFilesContext: vi.fn(), - subscribeToOpenFiles: vi.fn(() => vi.fn()), // subscribe returns an unsubscribe function + getIdeContext: vi.fn(), + subscribeToIdeContext: vi.fn(() => vi.fn()), // subscribe returns an unsubscribe function }; return { @@ -277,7 +277,7 @@ describe('App UI', () => { // Ensure a theme is set so the theme dialog does not appear. mockSettings = createMockSettings({ workspace: { theme: 'Default' } }); - vi.mocked(ideContext.getOpenFilesContext).mockReturnValue(undefined); + vi.mocked(ideContext.getIdeContext).mockReturnValue(undefined); }); afterEach(() => { @@ -289,10 +289,17 @@ describe('App UI', () => { }); it('should display active file when available', async () => { - vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ - activeFile: '/path/to/my-file.ts', - recentOpenFiles: [{ filePath: '/path/to/my-file.ts', content: 'hello' }], - selectedText: 'hello', + vi.mocked(ideContext.getIdeContext).mockReturnValue({ + workspaceState: { + openFiles: [ + { + path: '/path/to/my-file.ts', + isActive: true, + selectedText: 'hello', + timestamp: 0, + }, + ], + }, }); const { lastFrame, unmount } = render( @@ -304,12 +311,14 @@ describe('App UI', () => { ); currentUnmount = unmount; await Promise.resolve(); - expect(lastFrame()).toContain('1 recent file (ctrl+e to view)'); + expect(lastFrame()).toContain('1 open file (ctrl+e to view)'); }); - it('should not display active file when not available', async () => { - vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ - activeFile: '', + it('should not display any files when not available', async () => { + vi.mocked(ideContext.getIdeContext).mockReturnValue({ + workspaceState: { + openFiles: [], + }, }); const { lastFrame, unmount } = render( @@ -324,11 +333,54 @@ describe('App UI', () => { expect(lastFrame()).not.toContain('Open File'); }); + it('should display active file and other open files', async () => { + vi.mocked(ideContext.getIdeContext).mockReturnValue({ + workspaceState: { + openFiles: [ + { + path: '/path/to/my-file.ts', + isActive: true, + selectedText: 'hello', + timestamp: 0, + }, + { + path: '/path/to/another-file.ts', + isActive: false, + timestamp: 1, + }, + { + path: '/path/to/third-file.ts', + isActive: false, + timestamp: 2, + }, + ], + }, + }); + + const { lastFrame, unmount } = render( + , + ); + currentUnmount = unmount; + await Promise.resolve(); + expect(lastFrame()).toContain('3 open files (ctrl+e to view)'); + }); + it('should display active file and other context', async () => { - vi.mocked(ideContext.getOpenFilesContext).mockReturnValue({ - activeFile: '/path/to/my-file.ts', - recentOpenFiles: [{ filePath: '/path/to/my-file.ts', content: 'hello' }], - selectedText: 'hello', + vi.mocked(ideContext.getIdeContext).mockReturnValue({ + workspaceState: { + openFiles: [ + { + path: '/path/to/my-file.ts', + isActive: true, + selectedText: 'hello', + timestamp: 0, + }, + ], + }, }); mockConfig.getGeminiMdFileCount.mockReturnValue(1); mockConfig.getAllGeminiMdFilenames.mockReturnValue(['GEMINI.md']); @@ -343,7 +395,7 @@ describe('App UI', () => { currentUnmount = unmount; await Promise.resolve(); expect(lastFrame()).toContain( - 'Using: 1 recent file (ctrl+e to view) | 1 GEMINI.md file', + 'Using: 1 open file (ctrl+e to view) | 1 GEMINI.md file', ); }); diff --git a/packages/cli/src/ui/App.tsx b/packages/cli/src/ui/App.tsx index 87a78ac6..aacf45d7 100644 --- a/packages/cli/src/ui/App.tsx +++ b/packages/cli/src/ui/App.tsx @@ -60,7 +60,7 @@ import { FlashFallbackEvent, logFlashFallback, AuthType, - type OpenFiles, + type IdeContext, ideContext, } from '@google/gemini-cli-core'; import { validateAuthMethod } from '../config/auth.js'; @@ -169,13 +169,15 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { const [modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError] = useState(false); const [userTier, setUserTier] = useState(undefined); - const [openFiles, setOpenFiles] = useState(); + const [ideContextState, setIdeContextState] = useState< + IdeContext | undefined + >(); const [isProcessing, setIsProcessing] = useState(false); useEffect(() => { - const unsubscribe = ideContext.subscribeToOpenFiles(setOpenFiles); + const unsubscribe = ideContext.subscribeToIdeContext(setIdeContextState); // Set the initial value - setOpenFiles(ideContext.getOpenFilesContext()); + setIdeContextState(ideContext.getIdeContext()); return unsubscribe; }, []); @@ -568,7 +570,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { if (Object.keys(mcpServers || {}).length > 0) { handleSlashCommand(newValue ? '/mcp desc' : '/mcp nodesc'); } - } else if (key.ctrl && input === 'e' && ideContext) { + } else if (key.ctrl && input === 'e' && ideContextState) { setShowIDEContextDetail((prev) => !prev); } else if (key.ctrl && (input === 'c' || input === 'C')) { handleExit(ctrlCPressedOnce, setCtrlCPressedOnce, ctrlCTimerRef); @@ -943,7 +945,7 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => { ) : ( { {showIDEContextDetail && ( - + )} {showErrorDetails && ( diff --git a/packages/cli/src/ui/components/ContextSummaryDisplay.tsx b/packages/cli/src/ui/components/ContextSummaryDisplay.tsx index b166056a..78a19f0d 100644 --- a/packages/cli/src/ui/components/ContextSummaryDisplay.tsx +++ b/packages/cli/src/ui/components/ContextSummaryDisplay.tsx @@ -7,7 +7,7 @@ import React from 'react'; import { Text } from 'ink'; import { Colors } from '../colors.js'; -import { type OpenFiles, type MCPServerConfig } from '@google/gemini-cli-core'; +import { type IdeContext, type MCPServerConfig } from '@google/gemini-cli-core'; interface ContextSummaryDisplayProps { geminiMdFileCount: number; @@ -15,7 +15,7 @@ interface ContextSummaryDisplayProps { mcpServers?: Record; blockedMcpServers?: Array<{ name: string; extensionName: string }>; showToolDescriptions?: boolean; - openFiles?: OpenFiles; + ideContext?: IdeContext; } export const ContextSummaryDisplay: React.FC = ({ @@ -24,26 +24,28 @@ export const ContextSummaryDisplay: React.FC = ({ mcpServers, blockedMcpServers, showToolDescriptions, - openFiles, + ideContext, }) => { const mcpServerCount = Object.keys(mcpServers || {}).length; const blockedMcpServerCount = blockedMcpServers?.length || 0; + const openFileCount = ideContext?.workspaceState?.openFiles?.length ?? 0; if ( geminiMdFileCount === 0 && mcpServerCount === 0 && blockedMcpServerCount === 0 && - (openFiles?.recentOpenFiles?.length ?? 0) === 0 + openFileCount === 0 ) { return ; // Render an empty space to reserve height } - const recentFilesText = (() => { - const count = openFiles?.recentOpenFiles?.length ?? 0; - if (count === 0) { + const openFilesText = (() => { + if (openFileCount === 0) { return ''; } - return `${count} recent file${count > 1 ? 's' : ''} (ctrl+e to view)`; + return `${openFileCount} open file${ + openFileCount > 1 ? 's' : '' + } (ctrl+e to view)`; })(); const geminiMdText = (() => { @@ -81,8 +83,8 @@ export const ContextSummaryDisplay: React.FC = ({ let summaryText = 'Using: '; const summaryParts = []; - if (recentFilesText) { - summaryParts.push(recentFilesText); + if (openFilesText) { + summaryParts.push(openFilesText); } if (geminiMdText) { summaryParts.push(geminiMdText); diff --git a/packages/cli/src/ui/components/IDEContextDetailDisplay.tsx b/packages/cli/src/ui/components/IDEContextDetailDisplay.tsx index 8d4fb2c9..f535c40a 100644 --- a/packages/cli/src/ui/components/IDEContextDetailDisplay.tsx +++ b/packages/cli/src/ui/components/IDEContextDetailDisplay.tsx @@ -5,25 +5,21 @@ */ import { Box, Text } from 'ink'; -import { type OpenFiles } from '@google/gemini-cli-core'; +import { type File, type IdeContext } from '@google/gemini-cli-core'; import { Colors } from '../colors.js'; import path from 'node:path'; interface IDEContextDetailDisplayProps { - openFiles: OpenFiles | undefined; + ideContext: IdeContext | undefined; } export function IDEContextDetailDisplay({ - openFiles, + ideContext, }: IDEContextDetailDisplayProps) { - if ( - !openFiles || - !openFiles.recentOpenFiles || - openFiles.recentOpenFiles.length === 0 - ) { + const openFiles = ideContext?.workspaceState?.openFiles; + if (!openFiles || openFiles.length === 0) { return null; } - const recentFiles = openFiles.recentOpenFiles || []; return ( IDE Context (ctrl+e to toggle) - {recentFiles.length > 0 && ( + {openFiles.length > 0 && ( - Recent files: - {recentFiles.map((file) => ( - - - {path.basename(file.filePath)} - {file.filePath === openFiles.activeFile ? ' (active)' : ''} + Open files: + {openFiles.map((file: File) => ( + + - {path.basename(file.path)} + {file.isActive ? ' (active)' : ''} ))} 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,26 @@ 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); @@ -689,15 +701,188 @@ 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: +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 = { + 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.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); + + 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.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 + `.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 = { + 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.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; +describe('ideContext', () => { + describe('createIdeContextStore', () => { + let ideContext: ReturnType; - 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 set and retrieve the active file context', () => { - const testFile = { - activeFile: '/path/to/test/file.ts', - selectedText: '1234', - }; - - ideContext.setOpenFilesContext(testFile); - - const activeFile = ideContext.getOpenFilesContext(); - 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); - - const secondFile = { - activeFile: '/path/to/second.py', - cursor: { line: 20, character: 30 }, - }; - ideContext.setOpenFilesContext(secondFile); - - const activeFile = ideContext.getOpenFilesContext(); - 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 notify subscribers when active file context changes', () => { - const subscriber1 = vi.fn(); - const subscriber2 = vi.fn(); - - ideContext.subscribeToOpenFiles(subscriber1); - ideContext.subscribeToOpenFiles(subscriber2); - - const testFile = { - activeFile: '/path/to/subscribed.ts', - cursor: { line: 15, character: 25 }, - }; - ideContext.setOpenFilesContext(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); - - 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.subscribeToOpenFiles(subscriber1); - ideContext.subscribeToOpenFiles(subscriber2); - - ideContext.setOpenFilesContext({ - activeFile: '/path/to/file1.txt', - selectedText: '1234', + beforeEach(() => { + // Create a fresh, isolated instance for each test + ideContext = createIdeContextStore(); }); - expect(subscriber1).toHaveBeenCalledTimes(1); - expect(subscriber2).toHaveBeenCalledTimes(1); - unsubscribe1(); - - ideContext.setOpenFilesContext({ - activeFile: '/path/to/file2.txt', - selectedText: '1234', + it('should return undefined initially for ide context', () => { + expect(ideContext.getIdeContext()).toBeUndefined(); + }); + + 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.setIdeContext(testFile); + + const activeFile = ideContext.getIdeContext(); + expect(activeFile).toEqual(testFile); + }); + + 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 = { + workspaceState: { + openFiles: [ + { + path: '/path/to/second.py', + isActive: true, + cursor: { line: 20, character: 30 }, + timestamp: 0, + }, + ], + }, + }; + ideContext.setIdeContext(secondFile); + + const activeFile = ideContext.getIdeContext(); + expect(activeFile).toEqual(secondFile); + }); + + 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 ide context changes', () => { + const subscriber1 = vi.fn(); + const subscriber2 = vi.fn(); + + ideContext.subscribeToIdeContext(subscriber1); + ideContext.subscribeToIdeContext(subscriber2); + + 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); + + // 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); + }); + + it('should stop notifying a subscriber after unsubscribe', () => { + const subscriber1 = vi.fn(); + const subscriber2 = vi.fn(); + + const unsubscribe1 = ideContext.subscribeToIdeContext(subscriber1); + ideContext.subscribeToIdeContext(subscriber2); + + ideContext.setIdeContext({ + workspaceState: { + openFiles: [ + { + path: '/path/to/file1.txt', + isActive: true, + selectedText: '1234', + timestamp: 0, + }, + ], + }, + }); + expect(subscriber1).toHaveBeenCalledTimes(1); + expect(subscriber2).toHaveBeenCalledTimes(1); + + unsubscribe1(); + + 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); + }); + + it('should clear the ide context', () => { + const testFile = { + workspaceState: { + openFiles: [ + { + path: '/path/to/test/file.ts', + isActive: true, + selectedText: '1234', + timestamp: 0, + }, + ], + }, + }; + + 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); + }); - const activeFile = ideContext.getOpenFilesContext(); - expect(activeFile).toEqual(testFile); + it('should fail validation if path is missing', () => { + const file = { + timestamp: 12345, + }; + const result = FileSchema.safeParse(file); + expect(result.success).toBe(false); + }); + + 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 type Cursor = z.infer; - -/** - * Zod schema for validating an active file context from the IDE. - */ -export const OpenFilesSchema = z.object({ - activeFile: z.string(), +export const FileSchema = z.object({ + path: z.string(), + timestamp: z.number(), + isActive: z.boolean().optional(), selectedText: z.string().optional(), - cursor: CursorSchema.optional(), - recentOpenFiles: z - .array( - z.object({ - filePath: z.string(), - timestamp: z.number(), - }), - ) + cursor: z + .object({ + line: z.number(), + character: z.number(), + }) .optional(), }); -export type OpenFiles = z.infer; +export type File = z.infer; + +export const IdeContextSchema = z.object({ + workspaceState: z + .object({ + openFiles: z.array(FileSchema).optional(), + }) + .optional(), +}); +export type IdeContext = z.infer; /** - * 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(); + let ideContextState: IdeContext | undefined = undefined; + const subscribers = new Set(); /** - * 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, }; } diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index f47463ba..df8e160b 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -20,16 +20,17 @@ const MCP_SESSION_ID_HEADER = 'mcp-session-id'; const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT'; const MAX_SELECTED_TEXT_LENGTH = 16384; // 16 KiB limit -function sendOpenFilesChangedNotification( +function sendIdeContextUpdateNotification( transport: StreamableHTTPServerTransport, log: (message: string) => void, recentFilesManager: RecentFilesManager, ) { const editor = vscode.window.activeTextEditor; - const filePath = + const activeFile = editor && editor.document.uri.scheme === 'file' ? editor.document.uri.fsPath - : ''; + : undefined; + const selection = editor?.selection; const cursor = selection ? { @@ -38,25 +39,37 @@ function sendOpenFilesChangedNotification( character: selection.active.character, } : undefined; + let selectedText = editor?.document.getText(selection) ?? undefined; if (selectedText && selectedText.length > MAX_SELECTED_TEXT_LENGTH) { selectedText = selectedText.substring(0, MAX_SELECTED_TEXT_LENGTH) + '... [TRUNCATED]'; } + + const openFiles = recentFilesManager.recentFiles.map((file) => { + const isActive = file.filePath === activeFile; + return { + path: file.filePath, + timestamp: file.timestamp, + isActive, + ...(isActive && { + cursor, + selectedText, + }), + }; + }); + const notification: JSONRPCNotification = { jsonrpc: '2.0', - method: 'ide/openFilesChanged', + method: 'ide/contextUpdate', params: { - activeFile: filePath, - recentOpenFiles: recentFilesManager.recentFiles.filter( - (file) => file.filePath !== filePath, - ), - cursor, - selectedText, + workspaceState: { + openFiles, + }, }, }; log( - `Sending active file changed notification: ${JSON.stringify( + `Sending IDE context update notification: ${JSON.stringify( notification, null, 2, @@ -87,7 +100,7 @@ export class IDEServer { const recentFilesManager = new RecentFilesManager(context); const onDidChangeSubscription = recentFilesManager.onDidChange(() => { for (const transport of Object.values(transports)) { - sendOpenFilesChangedNotification( + sendIdeContextUpdateNotification( transport, this.log.bind(this), recentFilesManager, @@ -191,7 +204,7 @@ export class IDEServer { } if (!sessionsWithInitialNotification.has(sessionId)) { - sendOpenFilesChangedNotification( + sendIdeContextUpdateNotification( transport, this.log.bind(this), recentFilesManager,