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

@@ -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
 

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([])
+    })
+  })
+})

src/services/diagnosticTracking.ts

@@ -32,7 +32,7 @@ export class DiagnosticTrackingService {
   private baseline: Map<string, Diagnostic[]> = new Map()
 
   private initialized = false
-  private mcpClient: MCPServerConnection | undefined
+  private currentMcpClients: MCPServerConnection[] = []
 
   // Track when files were last processed/fetched
   private lastProcessedTimestamps: Map<string, number> = 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<void> {
     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<void> {
+    const ideClient = this.getCurrentIdeClient()
+    if (!ideClient) {
+      return
+    }
+    return await this.beforeFileEdited(filePath, ideClient)
+  }
+
+  /**
+   * Backward-compatible method that uses stored IDE client
+   */
+  async getNewDiagnosticsCompat(): Promise<DiagnosticFile[]> {
+    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<void> {
+    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<void> {
+  async ensureFileOpened(fileUri: string, mcpClient: MCPServerConnection): Promise<void> {
     if (
       !this.initialized ||
-      !this.mcpClient ||
+      !mcpClient ||
-      this.mcpClient.type !== 'connected'
+      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<void> {
+  async beforeFileEdited(filePath: string, mcpClient: MCPServerConnection): Promise<void> {
     if (
       !this.initialized ||
-      !this.mcpClient ||
+      !mcpClient ||
-      this.mcpClient.type !== 'connected'
+      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<DiagnosticFile[]> {
+  async getNewDiagnostics(mcpClient: MCPServerConnection): Promise<DiagnosticFile[]> {
     if (
       !this.initialized ||
-      !this.mcpClient ||
+      !mcpClient ||
-      this.mcpClient.type !== 'connected'
+      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<void> {
+    // 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

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

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

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 []
   }