[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
This commit is contained in:
Artem.Bukhonov
2025-07-18 21:32:54 +02:00
committed by intellij-monorepo-bot
parent 2140b3bad9
commit ef284060e5
7 changed files with 119 additions and 116 deletions

View File

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

View File

@@ -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,7 +52,7 @@ 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()

View File

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

View File

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

View File

@@ -216,22 +216,26 @@ 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 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 {
application.messageBus.syncPublisher(ToolCallListener.TOPIC).beforeMcpToolCall(this@mcpToolToRegisteredTool.descriptor)
application.messageBus.syncPublisher(ToolCallListener.TOPIC).beforeMcpToolCall(this@mcpToolToRegisteredTool.descriptor, additionalData)
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)
}
val result = this@mcpToolToRegisteredTool.call(request.arguments)
logger.trace { "Tool call successful '${this@mcpToolToRegisteredTool.descriptor.name}'. Result: ${result.content.joinToString("\n") { it.toString() }}" }
try {
@@ -285,30 +289,31 @@ private fun McpTool.mcpToolToRegisteredTool(server: Server, projectPathFromIniti
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)
application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, null, additionalData)
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)
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)
application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, mcpException, additionalData)
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)
application.messageBus.syncPublisher(ToolCallListener.TOPIC).afterMcpToolCall(this@mcpToolToRegisteredTool.descriptor, sideEffectEvents, t, additionalData)
McpToolCallResult.error(errorMessage)
}
finally {
McpServerCounterUsagesCollector.reportMcpCall(descriptor)
}
}
}
val contents = callResult.content.map { content ->
when (content) {

View File

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

View File

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