[MCP Server] Refactor some data like project, client info, descriptor and arguments to a data class add pass it as a parameter into listeners

(cherry picked from commit e3f6e3764c8b3237662c93aaf9e2f4c9782372f0)

GitOrigin-RevId: 08ae6cdac9476c916df41bd09de0f7d303f80cff

plugins/mcp-server/src/com/intellij/mcpserver/ClientInfo.kt

@@ -1,13 +0,0 @@
-package com.intellij.mcpserver
-
-import kotlin.coroutines.CoroutineContext
-
-class ClientInfo(val name: String, val version: String)
-
-internal class ClientInfoElement(val info: ClientInfo?) : CoroutineContext.Element {
-  companion object Key : CoroutineContext.Key<ClientInfoElement>
-  override val key: CoroutineContext.Key<*> = Key
-}
-
-val CoroutineContext.clientInfoOrNull: ClientInfo? get() = get(ClientInfoElement.Key)?.info
-val CoroutineContext.clientInfo: ClientInfo get() = get(ClientInfoElement.Key)?.info ?: ClientInfo("Unknown client", "Unknown version")

plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt

@@ -1,14 +1,38 @@
 package com.intellij.mcpserver
 
 import com.intellij.openapi.project.Project
+import kotlinx.serialization.json.JsonObject
 import org.jetbrains.ide.RestService.Companion.getLastFocusedOrOpenedProject
 import kotlin.coroutines.AbstractCoroutineContextElement
 import kotlin.coroutines.CoroutineContext
 
-data class ProjectContextElement(val project: Project?) : AbstractCoroutineContextElement(Key) {
-  companion object Key : CoroutineContext.Key<ProjectContextElement>
+class McpCallAdditionalData(
+  val clientInfo: ClientInfo,
+  val project: Project?,
+  val mcpToolDescriptor: McpToolDescriptor,
+  val rawArguments: JsonObject,
+  val meta: JsonObject
+)
+
+class ClientInfo(val name: String, val version: String)
+
+class McpCallAdditionalDataElement(val additionalData: McpCallAdditionalData) : AbstractCoroutineContextElement(Key) {
+  companion object Key : CoroutineContext.Key<McpCallAdditionalDataElement>
 }
 
+val CoroutineContext.mcpCallAdditionalData: McpCallAdditionalData get() = get(McpCallAdditionalDataElement)?.additionalData ?: error("mcpCallAdditionalData called outside of a MCP call")
+
+/**
+ * Returns information about the MCP client that is calling a tool.
+ */
+val CoroutineContext.clientInfo: ClientInfo get() = mcpCallAdditionalData.clientInfo
+
+
+/**
+ * Returns information about the MCP tool that is called.
+ */
+val CoroutineContext.currentToolDescriptor: McpToolDescriptor get() = mcpCallAdditionalData.mcpToolDescriptor
+
 /**
  * MCP tool can resolve a project with this extension property. In the case of running some MCP clients (like Claude) by IJ infrastructure
  * the project path can be specified by some ways (env or headers), then it can be resolved when calling MCP tool
@@ -28,8 +52,8 @@ val CoroutineContext.project: Project
  * The same as [projectOrNull], but allows to specify whether to look for any/last focused project or take only the one from the context element
  */
 fun CoroutineContext.getProjectOrNull(lookForAnyProject: Boolean): Project? {
-  val projectFromContext = this[ProjectContextElement.Key]?.project
+  val projectFromContext = mcpCallAdditionalData.project
   if (projectFromContext != null) return projectFromContext
   if (!lookForAnyProject) return null
   return getLastFocusedOrOpenedProject()
-}
+}

plugins/mcp-server/src/com/intellij/mcpserver/McpToolDescriptorElement.kt

@@ -1,12 +0,0 @@
-package com.intellij.mcpserver
-
-import kotlin.coroutines.AbstractCoroutineContextElement
-import kotlin.coroutines.CoroutineContext
-
-data class McpToolDescriptorElement(val mcpToolDescriptor: McpToolDescriptor?) : AbstractCoroutineContextElement(Key) {
-  companion object Key : CoroutineContext.Key<McpToolDescriptorElement>
-}
-
-val CoroutineContext.currentToolDescriptorOrNull: McpToolDescriptor? get() = get(McpToolDescriptorElement)?.mcpToolDescriptor
-val CoroutineContext.currentToolDescriptor: McpToolDescriptor get() = get(McpToolDescriptorElement)?.mcpToolDescriptor
-                                                                       ?: error("currentToolDescriptor called outside of a MCP tool call")

plugins/mcp-server/src/com/intellij/mcpserver/ToolCallListener.kt

@@ -12,11 +12,11 @@ interface ToolCallListener {
     val TOPIC: Topic<ToolCallListener> = Topic(ToolCallListener::class.java)
   }
 
-  fun beforeMcpToolCall(mcpToolDescriptor: McpToolDescriptor) {}
+  fun beforeMcpToolCall(mcpToolDescriptor: McpToolDescriptor, additionalData: McpCallAdditionalData) {}
 
-  fun afterMcpToolCall(mcpToolDescriptor: McpToolDescriptor, events: List<McpToolSideEffectEvent>, error: Throwable?) {}
+  fun afterMcpToolCall(mcpToolDescriptor: McpToolDescriptor, events: List<McpToolSideEffectEvent>, error: Throwable?, additionalData: McpCallAdditionalData) {}
 
-  fun toolActivity(mcpToolDescriptor: McpToolDescriptor, @NlsContexts.Label toolActivityDescription: String) {}
+  fun toolActivity(mcpToolDescriptor: McpToolDescriptor, @NlsContexts.Label toolActivityDescription: String, additionalData: McpCallAdditionalData) {}
 }
 
 sealed interface McpToolSideEffectEvent
@@ -29,5 +29,5 @@ class FileMovedEvent(val file: VirtualFile, val oldParent: VirtualFile, val newP
 class FileContentChangeEvent(val file: VirtualFile, val oldContent: String?, val newContent: String) : FileEvent
 
 fun CoroutineContext.reportToolActivity(@NlsContexts.Label toolDescription: String) {
-  application.messageBus.syncPublisher(ToolCallListener.TOPIC).toolActivity(this.currentToolDescriptor, toolDescription)
+  application.messageBus.syncPublisher(ToolCallListener.TOPIC).toolActivity(this.currentToolDescriptor, toolDescription, this.mcpCallAdditionalData)
 }

plugins/mcp-server/src/com/intellij/mcpserver/impl/McpServerService.kt

@@ -216,97 +216,102 @@ private fun McpTool.mcpToolToRegisteredTool(server: Server, projectPathFromIniti
         }
 
         EditorFactory.getInstance().eventMulticaster.addDocumentListener(documentListener, this.asDisposable())
+        val clientVersion = server.clientVersion ?: Implementation("Unknown MCP client", "Unknown version")
 
-        val sideEffectEvents = mutableListOf<McpToolSideEffectEvent>()
-        @Suppress("IncorrectCancellationExceptionHandling")
-        try {
-          application.messageBus.syncPublisher(ToolCallListener.TOPIC).beforeMcpToolCall(this@mcpToolToRegisteredTool.descriptor)
-
-          logger.trace { "Start calling tool '${this@mcpToolToRegisteredTool.descriptor.name}'. Arguments: ${request.arguments}" }
-          val clientVersion = server.clientVersion ?: Implementation("Unknown client", "Unknown version")
-
-          val result = withContext(
-            ProjectContextElement(project) +
-            McpToolDescriptorElement(descriptor) +
-            ClientInfoElement(ClientInfo(clientVersion.name, clientVersion.version))
-            ) {
-            this@mcpToolToRegisteredTool.call(request.arguments)
-          }
-
-          logger.trace { "Tool call successful '${this@mcpToolToRegisteredTool.descriptor.name}'. Result: ${result.content.joinToString("\n") { it.toString() }}" }
+        val additionalData = McpCallAdditionalData(
+          clientInfo = ClientInfo(clientVersion.name, clientVersion.version),
+          project = project,
+          mcpToolDescriptor = descriptor,
+          rawArguments = request.arguments,
+          meta = request._meta
+        )
+        withContext(
+          McpCallAdditionalDataElement(additionalData)
+        ) {
+          val sideEffectEvents = mutableListOf<McpToolSideEffectEvent>()
+          @Suppress("IncorrectCancellationExceptionHandling")
           try {
-            val processedChangedFiles = mutableSetOf<VirtualFile>()
+            application.messageBus.syncPublisher(ToolCallListener.TOPIC).beforeMcpToolCall(this@mcpToolToRegisteredTool.descriptor, additionalData)
 
-            for ((doc, oldContent) in initialDocumentContents) {
-              val virtualFile = FileDocumentManager.getInstance().getFile(doc) ?: continue
-              val newContent = readAction { doc.text }
-              sideEffectEvents.add(FileContentChangeEvent(virtualFile, oldContent, newContent))
-              processedChangedFiles.add(virtualFile)
-            }
+            logger.trace { "Start calling tool '${this@mcpToolToRegisteredTool.descriptor.name}'. Arguments: ${request.arguments}" }
 
-            for (event in vfsEvent) {
-              when (event) {
-                is VFileMoveEvent -> {
-                  sideEffectEvents.add(FileMovedEvent(event.file, event.oldParent, event.newParent))
-                }
-                is VFileCreateEvent -> {
-                  val virtualFile = event.file ?: continue
-                  val newContent = readAction { FileDocumentManager.getInstance().getDocument(virtualFile)?.text } ?: continue
-                  sideEffectEvents.add(FileCreatedEvent(virtualFile, newContent))
-                }
-                is VFileDeleteEvent -> {
-                  val virtualFile = event.file
-                  val document = readAction { FileDocumentManager.getInstance().getDocument(virtualFile) } ?: continue
-                  val oldContent = initialDocumentContents[document]
-                  sideEffectEvents.add(FileDeletedEvent(virtualFile, oldContent))
-                }
-                is VFileCopyEvent -> {
-                  val createdFile = event.findCreatedFile() ?: continue
-                  val newContent = readAction { FileDocumentManager.getInstance().getDocument(createdFile)?.text } ?: continue
-                  sideEffectEvents.add(FileCreatedEvent(createdFile, newContent))
-                }
-                is VFileContentChangeEvent -> {
-                  // reported in documents loop
-                  if (processedChangedFiles.contains(event.file)) continue
-                  val virtualFile = event.file
-                  val newContent = readAction { FileDocumentManager.getInstance().getDocument(virtualFile)?.text } ?: continue
-                  // Important: there may be a case when file is changed via low level change (like File.replaceText).
-                  // in this case we don't track the old content, because it may be heavy, it requires loading the file in
-                  // AsyncFileListener above and decoding with encoding etc. The file can be binary etc.
-                  sideEffectEvents.add(FileContentChangeEvent(virtualFile, oldContent = null, newContent = newContent))
+            val result = this@mcpToolToRegisteredTool.call(request.arguments)
+
+            logger.trace { "Tool call successful '${this@mcpToolToRegisteredTool.descriptor.name}'. Result: ${result.content.joinToString("\n") { it.toString() }}" }
+            try {
+              val processedChangedFiles = mutableSetOf<VirtualFile>()
+
+              for ((doc, oldContent) in initialDocumentContents) {
+                val virtualFile = FileDocumentManager.getInstance().getFile(doc) ?: continue
+                val newContent = readAction { doc.text }
+                sideEffectEvents.add(FileContentChangeEvent(virtualFile, oldContent, newContent))
+                processedChangedFiles.add(virtualFile)
+              }
+
+              for (event in vfsEvent) {
+                when (event) {
+                  is VFileMoveEvent -> {
+                    sideEffectEvents.add(FileMovedEvent(event.file, event.oldParent, event.newParent))
+                  }
+                  is VFileCreateEvent -> {
+                    val virtualFile = event.file ?: continue
+                    val newContent = readAction { FileDocumentManager.getInstance().getDocument(virtualFile)?.text } ?: continue
+                    sideEffectEvents.add(FileCreatedEvent(virtualFile, newContent))
+                  }
+                  is VFileDeleteEvent -> {
+                    val virtualFile = event.file
+                    val document = readAction { FileDocumentManager.getInstance().getDocument(virtualFile) } ?: continue
+                    val oldContent = initialDocumentContents[document]
+                    sideEffectEvents.add(FileDeletedEvent(virtualFile, oldContent))
+                  }
+                  is VFileCopyEvent -> {
+                    val createdFile = event.findCreatedFile() ?: continue
+                    val newContent = readAction { FileDocumentManager.getInstance().getDocument(createdFile)?.text } ?: continue
+                    sideEffectEvents.add(FileCreatedEvent(createdFile, newContent))
+                  }
+                  is VFileContentChangeEvent -> {
+                    // reported in documents loop
+                    if (processedChangedFiles.contains(event.file)) continue
+                    val virtualFile = event.file
+                    val newContent = readAction { FileDocumentManager.getInstance().getDocument(virtualFile)?.text } ?: continue
+                    // Important: there may be a case when file is changed via low level change (like File.replaceText).
+                    // in this case we don't track the old content, because it may be heavy, it requires loading the file in
+                    // AsyncFileListener above and decoding with encoding etc. The file can be binary etc.
+                    sideEffectEvents.add(FileContentChangeEvent(virtualFile, oldContent = null, newContent = newContent))
+                  }
                 }
               }
-            }
 
+            }
+            catch (ce: CancellationException) {
+              throw ce
+            }
+            catch (t: Throwable) {
+              logger.error("Failed to process changed documents after calling MCP tool ${this@mcpToolToRegisteredTool.descriptor.name}", t)
+            }
+            application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, null, additionalData)
+            result
           }
           catch (ce: CancellationException) {
-            throw ce
+            val message = "MCP tool call has been cancelled: ${ce.message}"
+            logger.traceThrowable { CancellationException(message, ce) }
+            application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, ce, additionalData)
+            McpToolCallResult.error(message)
+          }
+          catch (mcpException: McpExpectedError) {
+            logger.traceThrowable { mcpException }
+            application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, mcpException, additionalData)
+            McpToolCallResult.error(mcpException.mcpErrorText)
           }
           catch (t: Throwable) {
-            logger.error("Failed to process changed documents after calling MCP tool ${this@mcpToolToRegisteredTool.descriptor.name}", t)
+            val errorMessage = "MCP tool call has been failed: ${t.message}"
+            logger.error(t)
+            application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, t, additionalData)
+            McpToolCallResult.error(errorMessage)
+          }
+          finally {
+            McpServerCounterUsagesCollector.reportMcpCall(descriptor)
           }
-          application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, null)
-          result
-        }
-        catch (ce: CancellationException) {
-          val message = "MCP tool call has been cancelled: ${ce.message}"
-          logger.traceThrowable { CancellationException(message, ce) }
-          application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, ce)
-          McpToolCallResult.error(message)
-        }
-        catch (mcpException: McpExpectedError) {
-          logger.traceThrowable { mcpException }
-          application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, mcpException)
-          McpToolCallResult.error(mcpException.mcpErrorText)
-        }
-        catch (t: Throwable) {
-          val errorMessage = "MCP tool call has been failed: ${t.message}"
-          logger.error(t)
-          application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, t)
-          McpToolCallResult.error(errorMessage)
-        }
-        finally {
-          McpServerCounterUsagesCollector.reportMcpCall(descriptor)
         }
     }
 

plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/TerminalToolset.kt

@@ -55,7 +55,7 @@ class TerminalToolset : McpToolset {
     val project = currentCoroutineContext().project
     checkUserConfirmationIfNeeded(McpServerBundle.message("label.do.you.want.to.execute.command.in.terminal"), command, project)
 
-    val id = currentCoroutineContext().clientInfoOrNull?.name ?: "mcp_session"
+    val id = currentCoroutineContext().clientInfo.name
     val window = ToolWindowManager.getInstance(project).getToolWindow(TerminalToolWindowFactory.TOOL_WINDOW_ID)
     return executeShellCommand(window = window,
                                project = project,

plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/terminalToolsetUtil.kt

@@ -6,8 +6,7 @@ import com.intellij.execution.process.ColoredProcessHandler
 import com.intellij.execution.process.ProcessEvent
 import com.intellij.execution.process.ProcessListener
 import com.intellij.execution.process.ProcessOutputTypes
-import com.intellij.mcpserver.McpServerBundle
-import com.intellij.mcpserver.clientInfoOrNull
+import com.intellij.mcpserver.clientInfo
 import com.intellij.mcpserver.mcpFail
 import com.intellij.mcpserver.toolsets.terminal.TerminalToolset.CommandExecutionResult
 import com.intellij.mcpserver.util.TruncateMode
@@ -87,7 +86,7 @@ suspend fun executeShellCommand(
       else {
         val executionConsole = TerminalExecutionConsole(project, processHandler).withConvertLfToCrlfForNonPtyProcess(true)
         @Suppress("HardCodedStringLiteral")
-        val displayName = currentCoroutineContext().clientInfoOrNull?.name ?: McpServerBundle.message ("mcp.general.terminal.tab.name")
+        val displayName = currentCoroutineContext().clientInfo.name
         val content = ContentFactory.getInstance().createContent(executionConsole.component, displayName, false)
         window.contentManager.addContent(content)
         Disposer.register(content) {