Skip to content

Commit c8f8a48

Browse files
phodalphodal
committed
refactor(mcp): replace CheckboxTree with JPanel for tool selection and improve loading behavior
1 parent 308758e commit c8f8a48

File tree

1 file changed

+103
-117
lines changed

1 file changed

+103
-117
lines changed

‎core/src/main/kotlin/cc/unitmesh/devti/mcp/ui/McpConfigPopup.kt‎

Lines changed: 103 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,6 @@ package cc.unitmesh.devti.mcp.ui
33
import cc.unitmesh.devti.settings.customize.customizeSetting
44
import com.intellij.openapi.project.Project
55
import com.intellij.openapi.ui.popup.JBPopupFactory
6-
import com.intellij.ui.CheckboxTree
7-
import com.intellij.ui.CheckedTreeNode
8-
import com.intellij.ui.SimpleTextAttributes
96
import com.intellij.ui.components.JBLoadingPanel
107
import com.intellij.ui.components.JBScrollPane
118
import com.intellij.ui.components.JBTextField
@@ -18,26 +15,10 @@ import java.awt.Dimension
1815
import java.awt.event.KeyAdapter
1916
import java.awt.event.KeyEvent
2017
import javax.swing.*
21-
import javax.swing.tree.DefaultTreeModel
22-
import javax.swing.tree.TreePath
2318
import com.intellij.openapi.application.invokeLater
2419
import io.modelcontextprotocol.kotlin.sdk.Tool
25-
26-
class ServerTreeNode(val serverName: String) : CheckedTreeNode(serverName) {
27-
init {
28-
allowsChildren = true
29-
}
30-
}
31-
32-
class ToolTreeNode(val serverName: String, val tool: Tool) : CheckedTreeNode(tool.name) {
33-
init {
34-
allowsChildren = false
35-
userObject = tool.name
36-
}
37-
38-
override fun toString(): String = tool.name
39-
}
40-
20+
import java.awt.Component
21+
import java.awt.Font
4122

4223
class McpConfigPopup {
4324
companion object {
@@ -48,6 +29,11 @@ class McpConfigPopup {
4829
}
4930
}
5031

32+
private val toolsPanel = JPanel().apply {
33+
layout = BoxLayout(this, BoxLayout.Y_AXIS)
34+
}
35+
private val toolCheckboxMap = mutableMapOf<Pair<String, String>, JCheckBox>() // Stores serverName/toolName to JCheckBox
36+
5137
private fun createAndShow(component: JComponent?, project: Project, configService: McpConfigService) {
5238
val mainPanel = JPanel(BorderLayout()).apply {
5339
preferredSize = Dimension(400, 500)
@@ -60,39 +46,11 @@ class McpConfigPopup {
6046
border = JBUI.Borders.empty(4)
6147
}
6248

63-
// Tree for tool selection
64-
val rootNode = CheckedTreeNode("MCP Tools")
65-
val treeModel = DefaultTreeModel(rootNode)
66-
val tree = CheckboxTree(object : CheckboxTree.CheckboxTreeCellRenderer() {
67-
override fun customizeRenderer(
68-
tree: JTree?,
69-
value: Any?,
70-
selected: Boolean,
71-
expanded: Boolean,
72-
leaf: Boolean,
73-
row: Int,
74-
hasFocus: Boolean
75-
) {
76-
if (value is ToolTreeNode) {
77-
textRenderer.append(value.tool.name, SimpleTextAttributes.REGULAR_ATTRIBUTES)
78-
value.tool.description?.let { desc ->
79-
if (desc.isNotEmpty()) {
80-
textRenderer.append(" - $desc", SimpleTextAttributes.GRAYED_ATTRIBUTES)
81-
}
82-
}
83-
} else if (value is ServerTreeNode) {
84-
textRenderer.append(value.serverName, SimpleTextAttributes.REGULAR_BOLD_ATTRIBUTES)
85-
} else {
86-
textRenderer.append(value.toString(), SimpleTextAttributes.REGULAR_ATTRIBUTES)
87-
}
88-
}
89-
}, rootNode).apply {
90-
isRootVisible = false
91-
showsRootHandles = true
92-
}
93-
49+
// Panel for tool selection (replacing CheckboxTree)
50+
// toolsPanel is now a class member, initialized above
51+
9452
val loadingPanel = JBLoadingPanel(BorderLayout(), project)
95-
loadingPanel.add(JBScrollPane(tree), BorderLayout.CENTER)
53+
loadingPanel.add(JBScrollPane(toolsPanel), BorderLayout.CENTER) // toolsPanel instead of tree
9654
loadingPanel.preferredSize = Dimension(380, 350)
9755

9856
var currentPopup: com.intellij.openapi.ui.popup.JBPopup? = null
@@ -104,15 +62,15 @@ class McpConfigPopup {
10462
val refreshButton = JButton("Refresh").apply {
10563
addActionListener {
10664
loadingPanel.startLoading()
107-
refreshToolsList(project, configService, rootNode, treeModel, tree, loadingPanel)
65+
refreshToolsList(project, configService, loadingPanel)
10866
}
10967
}
11068
add(refreshButton)
11169
add(Box.createHorizontalGlue())
11270

11371
val applyButton = JButton("Apply").apply {
11472
addActionListener {
115-
saveSelectedTools(tree, configService)
73+
saveSelectedTools(configService)
11674
currentPopup?.cancel()
11775
}
11876
}
@@ -134,12 +92,12 @@ class McpConfigPopup {
13492

13593
// Load tools asynchronously
13694
loadingPanel.startLoading()
137-
loadToolsIntoTree(project, configService, rootNode, treeModel, tree, loadingPanel)
95+
loadToolsIntoPanel(project, configService, loadingPanel)
13896

13997
// Search functionality
14098
searchField.addKeyListener(object : KeyAdapter() {
14199
override fun keyReleased(e: KeyEvent) {
142-
filterTree(tree, rootNode, searchField.text)
100+
filterToolsList(searchField.text)
143101
}
144102
})
145103

@@ -162,30 +120,27 @@ class McpConfigPopup {
162120
private fun refreshToolsList(
163121
project: Project,
164122
configService: McpConfigService,
165-
rootNode: CheckedTreeNode,
166-
treeModel: DefaultTreeModel,
167-
tree: CheckboxTree,
168123
loadingPanel: JBLoadingPanel
169124
) {
170-
rootNode.removeAllChildren()
171-
treeModel.reload()
172-
173-
loadToolsIntoTree(project, configService, rootNode, treeModel, tree, loadingPanel)
125+
toolsPanel.removeAll()
126+
toolCheckboxMap.clear()
127+
128+
loadToolsIntoPanel(project, configService, loadingPanel)
174129
}
175130

176-
private fun loadToolsIntoTree(
131+
private fun loadToolsIntoPanel(
177132
project: Project,
178133
configService: McpConfigService,
179-
rootNode: CheckedTreeNode,
180-
treeModel: DefaultTreeModel,
181-
tree: CheckboxTree,
182134
loadingPanel: JBLoadingPanel
183135
) {
184-
rootNode.removeAllChildren()
185-
val loadingNode = CheckedTreeNode("Loading tools...")
186-
rootNode.add(loadingNode)
187-
treeModel.reload(rootNode) // Reload to show the loading node
188-
expandAllNodes(tree)
136+
toolsPanel.removeAll() // Clear previous content
137+
toolCheckboxMap.clear()
138+
139+
val loadingLabel = JLabel("Loading tools...")
140+
loadingLabel.alignmentX = Component.LEFT_ALIGNMENT
141+
toolsPanel.add(loadingLabel)
142+
toolsPanel.revalidate()
143+
toolsPanel.repaint()
189144

190145
CoroutineScope(Dispatchers.IO).launch {
191146
try {
@@ -194,55 +149,64 @@ class McpConfigPopup {
194149
val selectedTools = configService.getSelectedTools()
195150

196151
invokeLater {
197-
rootNode.removeAllChildren() // Remove "Loading tools..." node
152+
toolsPanel.removeAll() // Remove "Loading tools..." label
198153

199154
if (allTools.isEmpty()) {
200-
val noToolsNode = CheckedTreeNode("No tools available.")
201-
rootNode.add(noToolsNode)
155+
val noToolsLabel = JLabel("No tools available.")
156+
noToolsLabel.alignmentX = Component.LEFT_ALIGNMENT
157+
toolsPanel.add(noToolsLabel)
202158
} else {
203159
allTools.forEach { (serverName, tools) ->
204-
val serverNode = ServerTreeNode(serverName)
205-
rootNode.add(serverNode)
160+
val serverLabel = JLabel(serverName).apply {
161+
font = font.deriveFont(Font.BOLD)
162+
border = JBUI.Borders.emptyTop(8)
163+
alignmentX = Component.LEFT_ALIGNMENT
164+
}
165+
toolsPanel.add(serverLabel)
206166

207167
if (tools.isEmpty()) {
208-
val noToolsForServerNode = CheckedTreeNode("No tools from this server.")
209-
serverNode.add(noToolsForServerNode)
168+
val noToolsForServerLabel = JLabel(" No tools from this server.").apply {
169+
alignmentX = Component.LEFT_ALIGNMENT
170+
}
171+
toolsPanel.add(noToolsForServerLabel)
210172
} else {
211173
tools.forEach { tool ->
212-
val toolNode = ToolTreeNode(serverName, tool)
213-
val isSelected = selectedTools[serverName]?.contains(tool.name) == true
214-
toolNode.isChecked = isSelected
215-
serverNode.add(toolNode)
174+
val checkBoxText = tool.description?.let { desc ->
175+
if (desc.isNotEmpty()) "${tool.name} - $desc" else tool.name
176+
} ?: tool.name
177+
val checkBox = JCheckBox(checkBoxText).apply {
178+
isSelected = selectedTools[serverName]?.contains(tool.name) == true
179+
alignmentX = Component.LEFT_ALIGNMENT
180+
border = JBUI.Borders.emptyLeft(10)
181+
}
182+
toolCheckboxMap[Pair(serverName, tool.name)] = checkBox
183+
toolsPanel.add(checkBox)
216184
}
217185
}
218186
}
219187
}
220-
221-
treeModel.nodeStructureChanged(rootNode) // Notify that rootNode's children changed
222-
expandAllNodes(tree)
223188

224189
loadingPanel.stopLoading()
225190

226-
tree.revalidate() // Ensure tree layout is updated
227-
tree.repaint()
191+
toolsPanel.revalidate()
192+
toolsPanel.repaint()
228193
loadingPanel.revalidate()
229194
loadingPanel.repaint()
230-
// Also revalidate and repaint parent in case its layout depends on loadingPanel
231195
(loadingPanel.parent as? JComponent)?.revalidate()
232196
(loadingPanel.parent as? JComponent)?.repaint()
233197
}
234198
} catch (e: Exception) {
235199
invokeLater {
236-
rootNode.removeAllChildren()
237-
val errorNode = CheckedTreeNode("Error loading tools: ${e.message}")
238-
rootNode.add(errorNode)
239-
treeModel.nodeStructureChanged(rootNode) // Notify change
240-
expandAllNodes(tree)
241-
200+
toolsPanel.removeAll()
201+
val errorLabel = JLabel("Error loading tools: ${e.message}").apply {
202+
alignmentX = Component.LEFT_ALIGNMENT
203+
}
204+
toolsPanel.add(errorLabel)
205+
242206
loadingPanel.stopLoading()
243207

244-
tree.revalidate()
245-
tree.repaint()
208+
toolsPanel.revalidate()
209+
toolsPanel.repaint()
246210
loadingPanel.revalidate()
247211
loadingPanel.repaint()
248212
(loadingPanel.parent as? JComponent)?.revalidate()
@@ -252,33 +216,55 @@ class McpConfigPopup {
252216
}
253217
}
254218

255-
private fun saveSelectedTools(tree: CheckboxTree, configService: McpConfigService) {
219+
private fun saveSelectedTools(configService: McpConfigService) {
256220
val selectedTools = mutableMapOf<String, MutableSet<String>>()
257221

258-
val root = tree.model.root as CheckedTreeNode
259-
for (i in 0 until root.childCount) {
260-
val serverNode = root.getChildAt(i) as ServerTreeNode
261-
val serverName = serverNode.serverName
262-
263-
for (j in 0 until serverNode.childCount) {
264-
val toolNode = serverNode.getChildAt(j) as ToolTreeNode
265-
if (toolNode.isChecked) {
266-
selectedTools.computeIfAbsent(serverName) { mutableSetOf() }
267-
.add(toolNode.tool.name)
268-
}
222+
toolCheckboxMap.forEach { (key, checkBox) ->
223+
val (serverName, toolName) = key
224+
if (checkBox.isSelected) {
225+
selectedTools.computeIfAbsent(serverName) { mutableSetOf() }
226+
.add(toolName)
269227
}
270228
}
271229

272230
configService.setSelectedTools(selectedTools)
273231
}
274232

275-
private fun filterTree(tree: CheckboxTree, rootNode: CheckedTreeNode, searchText: String) {
276-
tree.expandPath(TreePath(rootNode.path))
277-
}
278-
279-
private fun expandAllNodes(tree: CheckboxTree) {
280-
for (i in 0 until tree.rowCount) {
281-
tree.expandRow(i)
233+
private fun filterToolsList(searchText: String) {
234+
val lowerSearchText = searchText.lowercase().trim()
235+
var firstVisible: Component? = null
236+
237+
toolsPanel.components.forEach { component ->
238+
when (component) {
239+
is JCheckBox -> {
240+
val toolName = toolCheckboxMap.entries.find { it.value == component }?.key?.second ?: ""
241+
val toolDescription = component.text.substringAfter("$toolName - ", "").substringBeforeLast(" - $toolName", "")
242+
243+
val isVisible = toolName.lowercase().contains(lowerSearchText) ||
244+
toolDescription.lowercase().contains(lowerSearchText) ||
245+
lowerSearchText.isEmpty()
246+
component.isVisible = isVisible
247+
if (isVisible && firstVisible == null) {
248+
firstVisible = component
249+
}
250+
}
251+
is JLabel -> {
252+
// Server labels or status labels. For now, keep them visible or hide if all children are hidden.
253+
// This part can be enhanced to hide server labels if all its tools are hidden.
254+
// For simplicity, we'll keep them visible. If search is empty, all are visible.
255+
component.isVisible = true
256+
}
257+
}
258+
}
259+
// If there's a search term and some items are visible, try to scroll to the first visible item.
260+
if (lowerSearchText.isNotEmpty() && firstVisible != null) {
261+
val finalFirstVisible = firstVisible
262+
SwingUtilities.invokeLater {
263+
toolsPanel.scrollRectToVisible(finalFirstVisible.bounds)
264+
}
282265
}
266+
267+
toolsPanel.revalidate()
268+
toolsPanel.repaint()
283269
}
284270
}

0 commit comments

Comments
 (0)