diff options
| -rw-r--r-- | packages/core/src/core/client.test.ts | 361 | ||||
| -rw-r--r-- | packages/core/src/core/client.ts | 226 |
2 files changed, 509 insertions, 78 deletions
diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index ff901a8b..2e8d086f 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -201,6 +201,7 @@ describe('Gemini Client (client.ts)', () => { getUsageStatisticsEnabled: vi.fn().mockReturnValue(true), getIdeModeFeature: vi.fn().mockReturnValue(false), getIdeMode: vi.fn().mockReturnValue(true), + getDebugMode: vi.fn().mockReturnValue(false), getWorkspaceContext: vi.fn().mockReturnValue({ getDirectories: vi.fn().mockReturnValue(['/test/dir']), }), @@ -449,8 +450,7 @@ describe('Gemini Client (client.ts)', () => { const mockChat = { addHistory: vi.fn(), }; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - client['chat'] = mockChat as any; + client['chat'] = mockChat as unknown as GeminiChat; const newContent = { role: 'user', @@ -667,7 +667,7 @@ describe('Gemini Client (client.ts)', () => { }); describe('sendMessageStream', () => { - it('should include IDE context when ideModeFeature is enabled', async () => { + it('should include editor context when ideModeFeature is enabled', async () => { // Arrange vi.mocked(ideContext.getIdeContext).mockReturnValue({ workspaceState: { @@ -725,21 +725,30 @@ describe('Gemini Client (client.ts)', () => { // 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 +Here is the user's editor context as a JSON object. This is for your information only. +\`\`\`json +${JSON.stringify( + { + activeFile: { + path: '/path/to/active/file.ts', + cursor: { + line: 5, + character: 10, + }, + selectedText: 'hello', + }, + otherOpenFiles: ['/path/to/recent/file1.ts', '/path/to/recent/file2.ts'], + }, + null, + 2, +)} +\`\`\` `.trim(); - const expectedRequest = [{ text: expectedContext }, ...initialRequest]; - expect(mockTurnRunFn).toHaveBeenCalledWith( - expectedRequest, - expect.any(Object), - ); + const expectedRequest = [{ text: expectedContext }]; + expect(mockChat.addHistory).toHaveBeenCalledWith({ + role: 'user', + parts: expectedRequest, + }); }); it('should not add context if ideModeFeature is enabled but no open files', async () => { @@ -839,18 +848,29 @@ Here are some other files the user has open, with the most recent at the top: // 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 is the user's editor context as a JSON object. This is for your information only. +\`\`\`json +${JSON.stringify( + { + activeFile: { + path: '/path/to/active/file.ts', + cursor: { + line: 5, + character: 10, + }, + selectedText: 'hello', + }, + }, + null, + 2, +)} +\`\`\` `.trim(); - const expectedRequest = [{ text: expectedContext }, ...initialRequest]; - expect(mockTurnRunFn).toHaveBeenCalledWith( - expectedRequest, - expect.any(Object), - ); + const expectedRequest = [{ text: expectedContext }]; + expect(mockChat.addHistory).toHaveBeenCalledWith({ + role: 'user', + parts: expectedRequest, + }); }); it('should add context if ideModeFeature is enabled and there are open files but no active file', async () => { @@ -904,15 +924,22 @@ This is the selected text in the file: // 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 +Here is the user's editor context as a JSON object. This is for your information only. +\`\`\`json +${JSON.stringify( + { + otherOpenFiles: ['/path/to/recent/file1.ts', '/path/to/recent/file2.ts'], + }, + null, + 2, +)} +\`\`\` `.trim(); - const expectedRequest = [{ text: expectedContext }, ...initialRequest]; - expect(mockTurnRunFn).toHaveBeenCalledWith( - expectedRequest, - expect.any(Object), - ); + const expectedRequest = [{ text: expectedContext }]; + expect(mockChat.addHistory).toHaveBeenCalledWith({ + role: 'user', + parts: expectedRequest, + }); }); it('should return the turn instance after the stream is complete', async () => { @@ -1190,6 +1217,268 @@ Here are some files the user has open, with the most recent at the top: `${eventCount} events generated (properly bounded by MAX_TURNS)`, ); }); + + describe('Editor context delta', () => { + const mockStream = (async function* () { + yield { type: 'content', value: 'Hello' }; + })(); + + beforeEach(() => { + client['forceFullIdeContext'] = false; // Reset before each delta test + vi.spyOn(client, 'tryCompressChat').mockResolvedValue(null); + vi.spyOn(client['config'], 'getIdeModeFeature').mockReturnValue(true); + mockTurnRunFn.mockReturnValue(mockStream); + + const mockChat: Partial<GeminiChat> = { + addHistory: vi.fn(), + setHistory: vi.fn(), + sendMessage: vi.fn().mockResolvedValue({ text: 'summary' }), + // Assume history is not empty for delta checks + getHistory: vi + .fn() + .mockReturnValue([ + { role: 'user', parts: [{ text: 'previous message' }] }, + ]), + }; + client['chat'] = mockChat as GeminiChat; + + const mockGenerator: Partial<ContentGenerator> = { + countTokens: vi.fn().mockResolvedValue({ totalTokens: 0 }), + generateContent: mockGenerateContentFn, + }; + client['contentGenerator'] = mockGenerator as ContentGenerator; + }); + + const testCases = [ + { + description: 'sends delta when active file changes', + previousActiveFile: { + path: '/path/to/old/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + shouldSendContext: true, + }, + { + description: 'sends delta when cursor line changes', + previousActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 1, character: 10 }, + selectedText: 'hello', + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + shouldSendContext: true, + }, + { + description: 'sends delta when cursor character changes', + previousActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 1 }, + selectedText: 'hello', + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + shouldSendContext: true, + }, + { + description: 'sends delta when selected text changes', + previousActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'world', + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + shouldSendContext: true, + }, + { + description: 'sends delta when selected text is added', + previousActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + shouldSendContext: true, + }, + { + description: 'sends delta when selected text is removed', + previousActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + }, + shouldSendContext: true, + }, + { + description: 'does not send context when nothing changes', + previousActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + currentActiveFile: { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }, + shouldSendContext: false, + }, + ]; + + it.each(testCases)( + '$description', + async ({ + previousActiveFile, + currentActiveFile, + shouldSendContext, + }) => { + // Setup previous context + client['lastSentIdeContext'] = { + workspaceState: { + openFiles: [ + { + path: previousActiveFile.path, + cursor: previousActiveFile.cursor, + selectedText: previousActiveFile.selectedText, + isActive: true, + timestamp: Date.now() - 1000, + }, + ], + }, + }; + + // Setup current context + vi.mocked(ideContext.getIdeContext).mockReturnValue({ + workspaceState: { + openFiles: [ + { ...currentActiveFile, isActive: true, timestamp: Date.now() }, + ], + }, + }); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-id-delta', + ); + for await (const _ of stream) { + // consume stream + } + + const mockChat = client['chat'] as unknown as { + addHistory: (typeof vi)['fn']; + }; + + if (shouldSendContext) { + expect(mockChat.addHistory).toHaveBeenCalledWith( + expect.objectContaining({ + parts: expect.arrayContaining([ + expect.objectContaining({ + text: expect.stringContaining( + "Here is a summary of changes in the user's editor context", + ), + }), + ]), + }), + ); + } else { + expect(mockChat.addHistory).not.toHaveBeenCalled(); + } + }, + ); + + it('sends full context when history is cleared, even if editor state is unchanged', async () => { + const activeFile = { + path: '/path/to/active/file.ts', + cursor: { line: 5, character: 10 }, + selectedText: 'hello', + }; + + // Setup previous context + client['lastSentIdeContext'] = { + workspaceState: { + openFiles: [ + { + path: activeFile.path, + cursor: activeFile.cursor, + selectedText: activeFile.selectedText, + isActive: true, + timestamp: Date.now() - 1000, + }, + ], + }, + }; + + // Setup current context (same as previous) + vi.mocked(ideContext.getIdeContext).mockReturnValue({ + workspaceState: { + openFiles: [ + { ...activeFile, isActive: true, timestamp: Date.now() }, + ], + }, + }); + + // Make history empty + const mockChat = client['chat'] as unknown as { + getHistory: ReturnType<(typeof vi)['fn']>; + addHistory: ReturnType<(typeof vi)['fn']>; + }; + mockChat.getHistory.mockReturnValue([]); + + const stream = client.sendMessageStream( + [{ text: 'Hi' }], + new AbortController().signal, + 'prompt-id-history-cleared', + ); + for await (const _ of stream) { + // consume stream + } + + expect(mockChat.addHistory).toHaveBeenCalledWith( + expect.objectContaining({ + parts: expect.arrayContaining([ + expect.objectContaining({ + text: expect.stringContaining( + "Here is the user's editor context", + ), + }), + ]), + }), + ); + + // Also verify it's the full context, not a delta. + const call = mockChat.addHistory.mock.calls[0][0]; + const contextText = call.parts[0].text; + const contextJson = JSON.parse( + contextText.match(/```json\n(.*)\n```/s)![1], + ); + expect(contextJson).toHaveProperty('activeFile'); + expect(contextJson.activeFile.path).toBe('/path/to/active/file.ts'); + }); + }); }); describe('generateContent', () => { diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 13e60039..df3dbc4e 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -50,6 +50,7 @@ import { NextSpeakerCheckEvent, } from '../telemetry/types.js'; import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; +import { IdeContext, File } from '../ide/ideContext.js'; function isThinkingSupported(model: string) { if (model.startsWith('gemini-2.5')) return true; @@ -112,6 +113,8 @@ export class GeminiClient { private readonly loopDetector: LoopDetectionService; private lastPromptId: string; + private lastSentIdeContext: IdeContext | undefined; + private forceFullIdeContext = true; constructor(private config: Config) { if (config.getProxy()) { @@ -164,6 +167,7 @@ export class GeminiClient { setHistory(history: Content[]) { this.getChat().setHistory(history); + this.forceFullIdeContext = true; } async setTools(): Promise<void> { @@ -189,6 +193,7 @@ export class GeminiClient { } async startChat(extraHistory?: Content[]): Promise<GeminiChat> { + this.forceFullIdeContext = true; const envParts = await getEnvironmentContext(this.config); const toolRegistry = await this.config.getToolRegistry(); const toolDeclarations = toolRegistry.getFunctionDeclarations(); @@ -238,6 +243,174 @@ export class GeminiClient { } } + private getIdeContextParts(forceFullContext: boolean): { + contextParts: string[]; + newIdeContext: IdeContext | undefined; + } { + const currentIdeContext = ideContext.getIdeContext(); + if (!currentIdeContext) { + return { contextParts: [], newIdeContext: undefined }; + } + + if (forceFullContext || !this.lastSentIdeContext) { + // Send full context as JSON + const openFiles = currentIdeContext.workspaceState?.openFiles || []; + const activeFile = openFiles.find((f) => f.isActive); + const otherOpenFiles = openFiles + .filter((f) => !f.isActive) + .map((f) => f.path); + + const contextData: Record<string, unknown> = {}; + + if (activeFile) { + contextData.activeFile = { + path: activeFile.path, + cursor: activeFile.cursor + ? { + line: activeFile.cursor.line, + character: activeFile.cursor.character, + } + : undefined, + selectedText: activeFile.selectedText || undefined, + }; + } + + if (otherOpenFiles.length > 0) { + contextData.otherOpenFiles = otherOpenFiles; + } + + if (Object.keys(contextData).length === 0) { + return { contextParts: [], newIdeContext: currentIdeContext }; + } + + const jsonString = JSON.stringify(contextData, null, 2); + const contextParts = [ + "Here is the user's editor context as a JSON object. This is for your information only.", + '```json', + jsonString, + '```', + ]; + + if (this.config.getDebugMode()) { + console.log(contextParts.join('\n')); + } + return { + contextParts, + newIdeContext: currentIdeContext, + }; + } else { + // Calculate and send delta as JSON + const delta: Record<string, unknown> = {}; + const changes: Record<string, unknown> = {}; + + const lastFiles = new Map( + (this.lastSentIdeContext.workspaceState?.openFiles || []).map( + (f: File) => [f.path, f], + ), + ); + const currentFiles = new Map( + (currentIdeContext.workspaceState?.openFiles || []).map((f: File) => [ + f.path, + f, + ]), + ); + + const openedFiles: string[] = []; + for (const [path] of currentFiles.entries()) { + if (!lastFiles.has(path)) { + openedFiles.push(path); + } + } + if (openedFiles.length > 0) { + changes.filesOpened = openedFiles; + } + + const closedFiles: string[] = []; + for (const [path] of lastFiles.entries()) { + if (!currentFiles.has(path)) { + closedFiles.push(path); + } + } + if (closedFiles.length > 0) { + changes.filesClosed = closedFiles; + } + + const lastActiveFile = ( + this.lastSentIdeContext.workspaceState?.openFiles || [] + ).find((f: File) => f.isActive); + const currentActiveFile = ( + currentIdeContext.workspaceState?.openFiles || [] + ).find((f: File) => f.isActive); + + if (currentActiveFile) { + if (!lastActiveFile || lastActiveFile.path !== currentActiveFile.path) { + changes.activeFileChanged = { + path: currentActiveFile.path, + cursor: currentActiveFile.cursor + ? { + line: currentActiveFile.cursor.line, + character: currentActiveFile.cursor.character, + } + : undefined, + selectedText: currentActiveFile.selectedText || undefined, + }; + } else { + const lastCursor = lastActiveFile.cursor; + const currentCursor = currentActiveFile.cursor; + if ( + currentCursor && + (!lastCursor || + lastCursor.line !== currentCursor.line || + lastCursor.character !== currentCursor.character) + ) { + changes.cursorMoved = { + path: currentActiveFile.path, + cursor: { + line: currentCursor.line, + character: currentCursor.character, + }, + }; + } + + const lastSelectedText = lastActiveFile.selectedText || ''; + const currentSelectedText = currentActiveFile.selectedText || ''; + if (lastSelectedText !== currentSelectedText) { + changes.selectionChanged = { + path: currentActiveFile.path, + selectedText: currentSelectedText, + }; + } + } + } else if (lastActiveFile) { + changes.activeFileChanged = { + path: null, + previousPath: lastActiveFile.path, + }; + } + + if (Object.keys(changes).length === 0) { + return { contextParts: [], newIdeContext: currentIdeContext }; + } + + delta.changes = changes; + const jsonString = JSON.stringify(delta, null, 2); + const contextParts = [ + "Here is a summary of changes in the user's editor context, in JSON format. This is for your information only.", + '```json', + jsonString, + '```', + ]; + + if (this.config.getDebugMode()) { + console.log(contextParts.join('\n')); + } + return { + contextParts, + newIdeContext: currentIdeContext, + }; + } + } + async *sendMessageStream( request: PartListUnion, signal: AbortSignal, @@ -273,49 +446,17 @@ export class GeminiClient { } if (this.config.getIdeModeFeature() && this.config.getIdeMode()) { - const ideContextState = ideContext.getIdeContext(); - const openFiles = ideContextState?.workspaceState?.openFiles; - - if (openFiles && openFiles.length > 0) { - const contextParts: string[] = []; - const firstFile = openFiles[0]; - const activeFile = firstFile.isActive ? firstFile : undefined; - - if (activeFile) { - contextParts.push( - `This is the file that the user is looking at:\n- Path: ${activeFile.path}`, - ); - if (activeFile.cursor) { - contextParts.push( - `This is the cursor position in the file:\n- Cursor Position: Line ${activeFile.cursor.line}, Character ${activeFile.cursor.character}`, - ); - } - if (activeFile.selectedText) { - contextParts.push( - `This is the selected text in the file:\n- ${activeFile.selectedText}`, - ); - } - } - - const otherOpenFiles = activeFile ? openFiles.slice(1) : openFiles; - - if (otherOpenFiles.length > 0) { - const recentFiles = otherOpenFiles - .map((file) => `- ${file.path}`) - .join('\n'); - 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) { - request = [ - { text: contextParts.join('\n') }, - ...(Array.isArray(request) ? request : [request]), - ]; - } + const { contextParts, newIdeContext } = this.getIdeContextParts( + this.forceFullIdeContext || this.getHistory().length === 0, + ); + if (contextParts.length > 0) { + this.getChat().addHistory({ + role: 'user', + parts: [{ text: contextParts.join('\n') }], + }); } + this.lastSentIdeContext = newIdeContext; + this.forceFullIdeContext = false; } const turn = new Turn(this.getChat(), prompt_id); @@ -648,6 +789,7 @@ export class GeminiClient { }, ...historyToKeep, ]); + this.forceFullIdeContext = true; const { totalTokens: newTokenCount } = await this.getContentGenerator().countTokens({ |