From ef284060e520c7d167591f2462c439b99dae2f77 Mon Sep 17 00:00:00 2001 From: "Artem.Bukhonov" Date: Fri, 18 Jul 2025 21:32:54 +0200 Subject: [PATCH] [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 --- .../src/com/intellij/mcpserver/ClientInfo.kt | 13 -- ...extElement.kt => McpCallAdditionalData.kt} | 32 +++- .../mcpserver/McpToolDescriptorElement.kt | 12 -- .../intellij/mcpserver/ToolCallListener.kt | 8 +- .../mcpserver/impl/McpServerService.kt | 163 +++++++++--------- .../toolsets/terminal/TerminalToolset.kt | 2 +- .../toolsets/terminal/terminalToolsetUtil.kt | 5 +- 7 files changed, 119 insertions(+), 116 deletions(-) delete mode 100644 plugins/mcp-server/src/com/intellij/mcpserver/ClientInfo.kt rename plugins/mcp-server/src/com/intellij/mcpserver/{ProjectContextElement.kt => McpCallAdditionalData.kt} (54%) delete mode 100644 plugins/mcp-server/src/com/intellij/mcpserver/McpToolDescriptorElement.kt diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/ClientInfo.kt b/plugins/mcp-server/src/com/intellij/mcpserver/ClientInfo.kt deleted file mode 100644 index 027d27e4bd8b..000000000000 --- a/plugins/mcp-server/src/com/intellij/mcpserver/ClientInfo.kt +++ /dev/null @@ -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 - 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") \ No newline at end of file diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/ProjectContextElement.kt b/plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt similarity index 54% rename from plugins/mcp-server/src/com/intellij/mcpserver/ProjectContextElement.kt rename to plugins/mcp-server/src/com/intellij/mcpserver/McpCallAdditionalData.kt index dad4eecec319..d628d0ffc20c 100644 --- a/plugins/mcp-server/src/com/intellij/mcpserver/ProjectContextElement.kt +++ b/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 +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 } +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() -} +} \ No newline at end of file diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/McpToolDescriptorElement.kt b/plugins/mcp-server/src/com/intellij/mcpserver/McpToolDescriptorElement.kt deleted file mode 100644 index 944a7179facf..000000000000 --- a/plugins/mcp-server/src/com/intellij/mcpserver/McpToolDescriptorElement.kt +++ /dev/null @@ -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 -} - -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") diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/ToolCallListener.kt b/plugins/mcp-server/src/com/intellij/mcpserver/ToolCallListener.kt index 25c7f2c198eb..9b3ce9b271e0 100644 --- a/plugins/mcp-server/src/com/intellij/mcpserver/ToolCallListener.kt +++ b/plugins/mcp-server/src/com/intellij/mcpserver/ToolCallListener.kt @@ -12,11 +12,11 @@ interface ToolCallListener { val TOPIC: Topic = Topic(ToolCallListener::class.java) } - fun beforeMcpToolCall(mcpToolDescriptor: McpToolDescriptor) {} + fun beforeMcpToolCall(mcpToolDescriptor: McpToolDescriptor, additionalData: McpCallAdditionalData) {} - fun afterMcpToolCall(mcpToolDescriptor: McpToolDescriptor, events: List, error: Throwable?) {} + fun afterMcpToolCall(mcpToolDescriptor: McpToolDescriptor, events: List, 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) } \ No newline at end of file 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 7b3c7ef00ccf..3639c4490761 100644 --- a/plugins/mcp-server/src/com/intellij/mcpserver/impl/McpServerService.kt +++ b/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() - @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() + @Suppress("IncorrectCancellationExceptionHandling") try { - val processedChangedFiles = mutableSetOf() + 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() + + 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) } } diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/TerminalToolset.kt b/plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/TerminalToolset.kt index e790f40f2fde..cef156398b4c 100644 --- a/plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/TerminalToolset.kt +++ b/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, diff --git a/plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/terminalToolsetUtil.kt b/plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/terminalToolsetUtil.kt index 51cb87268852..67867204f7f7 100644 --- a/plugins/mcp-server/src/com/intellij/mcpserver/toolsets/terminal/terminalToolsetUtil.kt +++ b/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) {