From 2c98be700274a4241963b5f43530bf3bd8f8963f Mon Sep 17 00:00:00 2001 From: Sreedhar Busanelli <18707170+sbusanelli@users.noreply.github.com> Date: Sat, 18 Apr 2026 20:02:52 -0500 Subject: [PATCH] fix: remove cached mcpClient in diagnostic tracking to prevent stale references (#727) * fix: remove cached mcpClient in diagnostic tracking to prevent stale references Resolves TODO comment about not caching the connected mcpClient since it can change. Changes: - Remove cached mcpClient field from DiagnosticTrackingService - Add currentMcpClients storage to track active clients - Update beforeFileEdited, getNewDiagnostics, and ensureFileOpened to accept client parameter - Add backward-compatible methods to maintain existing API - Update all callers to use new methods - Add comprehensive test coverage This prevents using stale MCP client references during reconnections, making diagnostic tracking more reliable. Fixes #TODO * docs: add my contributions section to README Add fork-specific section highlighting: - Diagnostic tracking enhancement (PR #727) - Technical skills demonstrated - Links to original project and my work - Professional contribution showcase * revert: remove README.md contributions section to comply with reviewer request - Remove 'My Fork & Contributions' section from README.md - Keep README.md focused on original project documentation - Maintain clean, project-focused README as requested by reviewer --- README.md | 3 +- src/services/diagnosticTracking.test.ts | 152 +++++++++++++++++++++++ src/services/diagnosticTracking.ts | 76 +++++++++--- src/tools/FileEditTool/FileEditTool.ts | 2 +- src/tools/FileWriteTool/FileWriteTool.ts | 2 +- src/utils/attachments.ts | 2 +- 6 files changed, 216 insertions(+), 21 deletions(-) create mode 100644 src/services/diagnosticTracking.test.ts diff --git a/README.md b/README.md index 52ecc7b8..3fde048c 100644 --- a/README.md +++ b/README.md @@ -331,7 +331,8 @@ For larger changes, open an issue first so the scope is clear before implementat - `bun run build` - `bun run test:coverage` - `bun run smoke` -- focused `bun test ...` runs for touched areas +- focused `bun test ...` runs for files and flows you changed + ## Disclaimer diff --git a/src/services/diagnosticTracking.test.ts b/src/services/diagnosticTracking.test.ts new file mode 100644 index 00000000..2559bac5 --- /dev/null +++ b/src/services/diagnosticTracking.test.ts @@ -0,0 +1,152 @@ +import { describe, test, expect, beforeEach, afterEach } from 'bun:test' +import { DiagnosticTrackingService } from './diagnosticTracking.js' +import type { MCPServerConnection } from './mcp/types.js' + +// Mock the IDE client utility +const mockGetConnectedIdeClient = (clients: MCPServerConnection[]) => + clients.find(client => client.type === 'connected') + +describe('DiagnosticTrackingService', () => { + let service: DiagnosticTrackingService + let mockClients: MCPServerConnection[] + let mockIdeClient: MCPServerConnection + + beforeEach(() => { + // Get fresh instance for each test + service = DiagnosticTrackingService.getInstance() + + // Setup mock clients + mockIdeClient = { + type: 'connected', + name: 'test-ide', + capabilities: {}, + config: {}, + cleanup: async () => {}, + client: { + request: async () => ({}), + setNotificationHandler: () => {}, + close: async () => {}, + }, + } as unknown as MCPServerConnection + + mockClients = [ + { type: 'disconnected', name: 'test-disconnected', config: {} } as unknown as MCPServerConnection, + mockIdeClient, + ] + }) + + afterEach(async () => { + await service.shutdown() + }) + + describe('handleQueryStart', () => { + test('should store MCP clients and initialize service', async () => { + await service.handleQueryStart(mockClients) + + // Service should be initialized + expect(service).toBeDefined() + + // Should be able to get IDE client from stored clients + // We can't directly test private methods, but we can test the behavior + const result = await service.getNewDiagnosticsCompat() + expect(result).toEqual([]) // Should return empty when no diagnostics + }) + + test('should reset service if already initialized', async () => { + // Initialize first + await service.handleQueryStart(mockClients) + + // Call again - should reset without error + await service.handleQueryStart(mockClients) + + // Should still work + const result = await service.getNewDiagnosticsCompat() + expect(result).toEqual([]) + }) + }) + + describe('backward-compatible methods', () => { + beforeEach(async () => { + await service.handleQueryStart(mockClients) + }) + + test('beforeFileEditedCompat should work without explicit client', async () => { + // Should not throw error and should return undefined when no IDE client + const result = await service.beforeFileEditedCompat('/test/file.ts') + expect(result).toBeUndefined() + }) + + test('getNewDiagnosticsCompat should work without explicit client', async () => { + const result = await service.getNewDiagnosticsCompat() + expect(Array.isArray(result)).toBe(true) + }) + + test('ensureFileOpenedCompat should work without explicit client', async () => { + const result = await service.ensureFileOpenedCompat('/test/file.ts') + expect(result).toBeUndefined() + }) + }) + + describe('new explicit client methods', () => { + test('beforeFileEdited should require client parameter', async () => { + // Should not work without client + const result = await service.beforeFileEdited('/test/file.ts', undefined as any) + expect(result).toBeUndefined() + }) + + test('getNewDiagnostics should require client parameter', async () => { + // Should not work without client + const result = await service.getNewDiagnostics(undefined as any) + expect(result).toEqual([]) + }) + + test('ensureFileOpened should require client parameter', async () => { + // Should not work without client + const result = await service.ensureFileOpened('/test/file.ts', undefined as any) + expect(result).toBeUndefined() + }) + }) + + describe('shutdown', () => { + test('should clear stored clients on shutdown', async () => { + await service.handleQueryStart(mockClients) + + // Verify service is working + const beforeResult = await service.getNewDiagnosticsCompat() + expect(Array.isArray(beforeResult)).toBe(true) + + // Shutdown + await service.shutdown() + + // After shutdown, compat methods should return empty results + const afterResult = await service.getNewDiagnosticsCompat() + expect(afterResult).toEqual([]) + }) + }) + + describe('integration with existing functionality', () => { + test('should maintain existing diagnostic tracking behavior', async () => { + await service.handleQueryStart(mockClients) + + // Test baseline tracking + await service.beforeFileEditedCompat('/test/file.ts') + + // Test getting new diagnostics (should be empty since no IDE client is actually connected) + const newDiagnostics = await service.getNewDiagnosticsCompat() + expect(Array.isArray(newDiagnostics)).toBe(true) + }) + + test('should handle missing IDE client gracefully', async () => { + // Test with no connected clients + const noIdeClients = [ + { type: 'disconnected', name: 'test-disconnected-2', config: {} } as unknown as MCPServerConnection, + ] + + await service.handleQueryStart(noIdeClients) + + // Should handle gracefully + const result = await service.getNewDiagnosticsCompat() + expect(result).toEqual([]) + }) + }) +}) diff --git a/src/services/diagnosticTracking.ts b/src/services/diagnosticTracking.ts index 2978a3dd..de5f5b1e 100644 --- a/src/services/diagnosticTracking.ts +++ b/src/services/diagnosticTracking.ts @@ -32,7 +32,7 @@ export class DiagnosticTrackingService { private baseline: Map = new Map() private initialized = false - private mcpClient: MCPServerConnection | undefined + private currentMcpClients: MCPServerConnection[] = [] // Track when files were last processed/fetched private lastProcessedTimestamps: Map = new Map() @@ -48,18 +48,17 @@ export class DiagnosticTrackingService { return DiagnosticTrackingService.instance } - initialize(mcpClient: MCPServerConnection) { + initialize() { if (this.initialized) { return } - // TODO: Do not cache the connected mcpClient since it can change. - this.mcpClient = mcpClient this.initialized = true } async shutdown(): Promise { this.initialized = false + this.currentMcpClients = [] this.baseline.clear() this.rightFileDiagnosticsState.clear() this.lastProcessedTimestamps.clear() @@ -75,6 +74,46 @@ export class DiagnosticTrackingService { this.lastProcessedTimestamps.clear() } + /** + * Get the current IDE client from stored MCP clients + */ + private getCurrentIdeClient(): MCPServerConnection | undefined { + return getConnectedIdeClient(this.currentMcpClients) + } + + /** + * Backward-compatible method that uses stored IDE client + */ + async beforeFileEditedCompat(filePath: string): Promise { + const ideClient = this.getCurrentIdeClient() + if (!ideClient) { + return + } + return await this.beforeFileEdited(filePath, ideClient) + } + + /** + * Backward-compatible method that uses stored IDE client + */ + async getNewDiagnosticsCompat(): Promise { + const ideClient = this.getCurrentIdeClient() + if (!ideClient) { + return [] + } + return await this.getNewDiagnostics(ideClient) + } + + /** + * Backward-compatible method that uses stored IDE client + */ + async ensureFileOpenedCompat(fileUri: string): Promise { + const ideClient = this.getCurrentIdeClient() + if (!ideClient) { + return + } + return await this.ensureFileOpened(fileUri, ideClient) + } + private normalizeFileUri(fileUri: string): string { // Remove our protocol prefixes const protocolPrefixes = [ @@ -100,11 +139,11 @@ export class DiagnosticTrackingService { * Ensure a file is opened in the IDE before processing. * This is important for language services like diagnostics to work properly. */ - async ensureFileOpened(fileUri: string): Promise { + async ensureFileOpened(fileUri: string, mcpClient: MCPServerConnection): Promise { if ( !this.initialized || - !this.mcpClient || - this.mcpClient.type !== 'connected' + !mcpClient || + mcpClient.type !== 'connected' ) { return } @@ -121,7 +160,7 @@ export class DiagnosticTrackingService { selectToEndOfLine: false, makeFrontmost: false, }, - this.mcpClient, + mcpClient, ) } catch (error) { logError(error as Error) @@ -132,11 +171,11 @@ export class DiagnosticTrackingService { * Capture baseline diagnostics for a specific file before editing. * This is called before editing a file to ensure we have a baseline to compare against. */ - async beforeFileEdited(filePath: string): Promise { + async beforeFileEdited(filePath: string, mcpClient: MCPServerConnection): Promise { if ( !this.initialized || - !this.mcpClient || - this.mcpClient.type !== 'connected' + !mcpClient || + mcpClient.type !== 'connected' ) { return } @@ -147,7 +186,7 @@ export class DiagnosticTrackingService { const result = await callIdeRpc( 'getDiagnostics', { uri: `file://${filePath}` }, - this.mcpClient, + mcpClient, ) const diagnosticFile = this.parseDiagnosticResult(result)[0] if (diagnosticFile) { @@ -185,11 +224,11 @@ export class DiagnosticTrackingService { * Get new diagnostics from file://, _claude_fs_right, and _claude_fs_ URIs that aren't in the baseline. * Only processes diagnostics for files that have been edited. */ - async getNewDiagnostics(): Promise { + async getNewDiagnostics(mcpClient: MCPServerConnection): Promise { if ( !this.initialized || - !this.mcpClient || - this.mcpClient.type !== 'connected' + !mcpClient || + mcpClient.type !== 'connected' ) { return [] } @@ -200,7 +239,7 @@ export class DiagnosticTrackingService { const result = await callIdeRpc( 'getDiagnostics', {}, // Empty params fetches all diagnostics - this.mcpClient, + mcpClient, ) allDiagnosticFiles = this.parseDiagnosticResult(result) } catch (_error) { @@ -328,13 +367,16 @@ export class DiagnosticTrackingService { * @param shouldQuery Whether a query is actually being made (not just a command) */ async handleQueryStart(clients: MCPServerConnection[]): Promise { + // Store the current MCP clients for later use + this.currentMcpClients = clients + // Only proceed if we should query and have clients if (!this.initialized) { // Find the connected IDE client const connectedIdeClient = getConnectedIdeClient(clients) if (connectedIdeClient) { - this.initialize(connectedIdeClient) + this.initialize() } } else { // Reset diagnostic tracking for new query loops diff --git a/src/tools/FileEditTool/FileEditTool.ts b/src/tools/FileEditTool/FileEditTool.ts index 5c047fad..e8a13fcb 100644 --- a/src/tools/FileEditTool/FileEditTool.ts +++ b/src/tools/FileEditTool/FileEditTool.ts @@ -422,7 +422,7 @@ export const FileEditTool = buildTool({ activateConditionalSkillsForPaths([absoluteFilePath], cwd) } - await diagnosticTracker.beforeFileEdited(absoluteFilePath) + await diagnosticTracker.beforeFileEditedCompat(absoluteFilePath) // Ensure parent directory exists before the atomic read-modify-write section. // These awaits must stay OUTSIDE the critical section below — a yield between diff --git a/src/tools/FileWriteTool/FileWriteTool.ts b/src/tools/FileWriteTool/FileWriteTool.ts index e9f9aec1..1aa2d9dc 100644 --- a/src/tools/FileWriteTool/FileWriteTool.ts +++ b/src/tools/FileWriteTool/FileWriteTool.ts @@ -244,7 +244,7 @@ export const FileWriteTool = buildTool({ // Activate conditional skills whose path patterns match this file activateConditionalSkillsForPaths([fullFilePath], cwd) - await diagnosticTracker.beforeFileEdited(fullFilePath) + await diagnosticTracker.beforeFileEditedCompat(fullFilePath) // Ensure parent directory exists before the atomic read-modify-write section. // Must stay OUTSIDE the critical section below (a yield between the staleness diff --git a/src/utils/attachments.ts b/src/utils/attachments.ts index 37b5eb65..e3c83621 100644 --- a/src/utils/attachments.ts +++ b/src/utils/attachments.ts @@ -2882,7 +2882,7 @@ async function getDiagnosticAttachments( } // Get new diagnostics from the tracker (IDE diagnostics via MCP) - const newDiagnostics = await diagnosticTracker.getNewDiagnostics() + const newDiagnostics = await diagnosticTracker.getNewDiagnosticsCompat() if (newDiagnostics.length === 0) { return [] }