From 7c7a9b1ce45cd97bf69b1c26b47d69ba044ec830 Mon Sep 17 00:00:00 2001
From: "Artem.Bukhonov" <artem.bukhonov@jetbrains.com>
Date: Fri, 18 Jul 2025 22:58:20 +0200
Subject: [PATCH] [MCP Server] Track changes related only to a particular call

(cherry picked from commit ae72e42dba48c6c6a656c4a9bb1a45c92546fec8)

GitOrigin-RevId: 894355da9e7928ad0723846154922a0f03fffa1f
---
 .../mcpserver/McpCallAdditionalData.kt        | 10 ++++--
 .../mcpserver/impl/McpServerService.kt        | 35 +++++++++++++------
 2 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt b/plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt
index d628d0ffc20c..b04dafa67c03 100644
--- a/plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt
+++ b/plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt
@@ -7,12 +7,17 @@ import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 
 class McpCallAdditionalData(
+  val callId: Int,
   val clientInfo: ClientInfo,
   val project: Project?,
   val mcpToolDescriptor: McpToolDescriptor,
   val rawArguments: JsonObject,
   val meta: JsonObject
-)
+) {
+  override fun toString(): String {
+    return "McpCallAdditionalData(id=$callId, clientInfo=$clientInfo, toolName=${mcpToolDescriptor.name}"
+  }
+}
 
 class ClientInfo(val name: String, val version: String)
 
@@ -20,7 +25,8 @@ class McpCallAdditionalDataElement(val additionalData: McpCallAdditionalData) :
   companion object Key : CoroutineContext.Key<McpCallAdditionalDataElement>
 }
 
-val CoroutineContext.mcpCallAdditionalData: McpCallAdditionalData get() = get(McpCallAdditionalDataElement)?.additionalData ?: error("mcpCallAdditionalData called outside of a MCP call")
+val CoroutineContext.mcpCallAdditionalDataOrNull: McpCallAdditionalData? get() = get(McpCallAdditionalDataElement)?.additionalData
+val CoroutineContext.mcpCallAdditionalData: McpCallAdditionalData get() = mcpCallAdditionalDataOrNull ?: error("mcpCallAdditionalData called outside of a MCP call")
 
 /**
  * Returns information about the MCP client that is calling a tool.
diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/impl/McpServerService.kt b/plugins/mcp-server/src/com/intellij/mcpserver/impl/McpServerService.kt
index 3639c4490761..7e56c251ce39 100644
--- a/plugins/mcp-server/src/com/intellij/mcpserver/impl/McpServerService.kt
+++ b/plugins/mcp-server/src/com/intellij/mcpserver/impl/McpServerService.kt
@@ -1,5 +1,6 @@
 package com.intellij.mcpserver.impl
 
+import com.intellij.concurrency.currentThreadContext
 import com.intellij.mcpserver.*
 import com.intellij.mcpserver.impl.util.network.*
 import com.intellij.mcpserver.settings.McpServerSettings
@@ -45,6 +46,8 @@ import kotlinx.coroutines.flow.update
 import kotlinx.serialization.json.JsonPrimitive
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.CopyOnWriteArrayList
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.concurrent.atomics.ExperimentalAtomicApi
 import kotlin.coroutines.cancellation.CancellationException
 
 
@@ -58,6 +61,8 @@ class McpServerService(val cs: CoroutineScope) {
   }
 
   private val server = MutableStateFlow(startServerIfEnabled())
+  @OptIn(ExperimentalAtomicApi::class)
+  private val callId = AtomicInteger(0)
 
   val isRunning: Boolean
     get() = server.value != null
@@ -198,11 +203,25 @@ private fun McpTool.mcpToolToRegisteredTool(server: Server, projectPathFromIniti
 
       val vfsEvent = CopyOnWriteArrayList<VFileEvent>()
       val initialDocumentContents = ConcurrentHashMap<Document, String>()
+      val clientVersion = server.clientVersion ?: Implementation("Unknown MCP client", "Unknown version")
+
+      val additionalData = McpCallAdditionalData(
+        callId = callId.getAndAdd(1),
+        clientInfo = ClientInfo(clientVersion.name, clientVersion.version),
+        project = project,
+        mcpToolDescriptor = descriptor,
+        rawArguments = request.arguments,
+        meta = request._meta
+      )
 
       val callResult = coroutineScope {
 
         VirtualFileManager.getInstance().addAsyncFileListener(this, AsyncFileListener { events ->
-          vfsEvent.addAll(events)
+          val inHandlerData = currentThreadContext().mcpCallAdditionalDataOrNull
+          if (inHandlerData != null && inHandlerData.callId == additionalData.callId) {
+            logger.trace { "VFS changes detected for call: $inHandlerData" }
+            vfsEvent.addAll(events)
+          }
           // probably we have to read initial contents here
           // see comment below near `is VFileContentChangeEvent`
           return@AsyncFileListener object : AsyncFileListener.ChangeApplier {}
@@ -211,20 +230,16 @@ private fun McpTool.mcpToolToRegisteredTool(server: Server, projectPathFromIniti
         val documentListener = object : DocumentListener {
           // record content before any change
           override fun beforeDocumentChange(event: DocumentEvent) {
-            initialDocumentContents.computeIfAbsent(event.document) { event.document.text }
+            val inHandlerData = currentThreadContext().mcpCallAdditionalDataOrNull
+            if (inHandlerData != null && inHandlerData.callId == additionalData.callId) {
+              logger.trace { "Document changes detected for call: $inHandlerData" }
+              initialDocumentContents.computeIfAbsent(event.document) { event.document.text }
+            }
           }
         }
 
         EditorFactory.getInstance().eventMulticaster.addDocumentListener(documentListener, this.asDisposable())
-        val clientVersion = server.clientVersion ?: Implementation("Unknown MCP client", "Unknown version")
 
-        val additionalData = McpCallAdditionalData(
-          clientInfo = ClientInfo(clientVersion.name, clientVersion.version),
-          project = project,
-          mcpToolDescriptor = descriptor,
-          rawArguments = request.arguments,
-          meta = request._meta
-        )
         withContext(
           McpCallAdditionalDataElement(additionalData)
         ) {