Skip to content

Commit 6155283

Browse files
committed
fix: prevent MCP server restarts during active tool executions
- Add tracking of active tool executions in McpHub - Prevent server restarts when tools are running - Update toggleToolAlwaysAllow to skip restart during tool execution - Update toggleToolEnabledForPrompt to skip restart during tool execution - Prevent toggleServerDisabled when tools are running - Add comprehensive tests for the new behavior Fixes #7189
1 parent 4222036 commit 6155283

File tree

2 files changed

+376
-14
lines changed

2 files changed

+376
-14
lines changed

src/services/mcp/McpHub.ts

Lines changed: 220 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,7 @@ export class McpHub {
151151
isConnecting: boolean = false
152152
private refCount: number = 0 // Reference counter for active clients
153153
private configChangeDebounceTimers: Map<string, NodeJS.Timeout> = new Map()
154+
private activeToolExecutions: Map<string, Set<string>> = new Map() // Track active tool executions per server
154155

155156
constructor(provider: ClineProvider) {
156157
this.providerRef = new WeakRef(provider)
@@ -1169,6 +1170,16 @@ export class McpHub {
11691170
}
11701171

11711172
async restartConnection(serverName: string, source?: "global" | "project"): Promise<void> {
1173+
// Check if there are active tool executions for this server
1174+
if (this.hasActiveToolExecutions(serverName, source)) {
1175+
console.log(`Skipping restart for ${serverName} - tools are currently executing`)
1176+
vscode.window.showWarningMessage(
1177+
t("mcp:errors.cannot_restart_tools_running", { serverName }) ||
1178+
`Cannot restart server "${serverName}" while tools are running. Please wait for tool execution to complete.`,
1179+
)
1180+
return
1181+
}
1182+
11721183
this.isConnecting = true
11731184

11741185
// Check if MCP is globally enabled
@@ -1357,6 +1368,20 @@ export class McpHub {
13571368
}
13581369

13591370
const serverSource = connection.server.source || "global"
1371+
1372+
// Check if there are active tool executions for this server
1373+
if (this.hasActiveToolExecutions(serverName, serverSource)) {
1374+
// Show a warning message and don't proceed with the toggle
1375+
vscode.window.showWarningMessage(
1376+
t("mcp:errors.cannot_toggle_server_tools_running", {
1377+
serverName,
1378+
action: disabled ? "disable" : "enable",
1379+
}) ||
1380+
`Cannot ${disabled ? "disable" : "enable"} server "${serverName}" while tools are running. Please wait for tool execution to complete.`,
1381+
)
1382+
return
1383+
}
1384+
13601385
// Update the server config in the appropriate file
13611386
await this.updateServerConfig(serverName, { disabled }, serverSource)
13621387

@@ -1603,19 +1628,69 @@ export class McpHub {
16031628
timeout = 60 * 1000
16041629
}
16051630

1606-
return await connection.client.request(
1607-
{
1608-
method: "tools/call",
1609-
params: {
1610-
name: toolName,
1611-
arguments: toolArguments,
1631+
// Track this tool execution as active
1632+
const serverKey = `${serverName}:${source || "global"}`
1633+
if (!this.activeToolExecutions.has(serverKey)) {
1634+
this.activeToolExecutions.set(serverKey, new Set())
1635+
}
1636+
const executionId = `${toolName}:${Date.now()}`
1637+
this.activeToolExecutions.get(serverKey)!.add(executionId)
1638+
1639+
try {
1640+
const result = await connection.client.request(
1641+
{
1642+
method: "tools/call",
1643+
params: {
1644+
name: toolName,
1645+
arguments: toolArguments,
1646+
},
16121647
},
1613-
},
1614-
CallToolResultSchema,
1615-
{
1616-
timeout,
1617-
},
1618-
)
1648+
CallToolResultSchema,
1649+
{
1650+
timeout,
1651+
},
1652+
)
1653+
1654+
// Remove from active executions on success
1655+
this.activeToolExecutions.get(serverKey)?.delete(executionId)
1656+
if (this.activeToolExecutions.get(serverKey)?.size === 0) {
1657+
this.activeToolExecutions.delete(serverKey)
1658+
}
1659+
1660+
return result
1661+
} catch (error) {
1662+
// Remove from active executions on error
1663+
this.activeToolExecutions.get(serverKey)?.delete(executionId)
1664+
if (this.activeToolExecutions.get(serverKey)?.size === 0) {
1665+
this.activeToolExecutions.delete(serverKey)
1666+
}
1667+
throw error
1668+
}
1669+
}
1670+
1671+
/**
1672+
* Check if any tools are currently executing for a specific server
1673+
* @param serverName The name of the server to check
1674+
* @param source The source of the server (global or project)
1675+
* @returns true if there are active tool executions, false otherwise
1676+
*/
1677+
private hasActiveToolExecutions(serverName: string, source?: "global" | "project"): boolean {
1678+
const serverKey = `${serverName}:${source || "global"}`
1679+
const activeTools = this.activeToolExecutions.get(serverKey)
1680+
return activeTools ? activeTools.size > 0 : false
1681+
}
1682+
1683+
/**
1684+
* Check if any tools are currently executing across all servers
1685+
* @returns true if there are any active tool executions, false otherwise
1686+
*/
1687+
private hasAnyActiveToolExecutions(): boolean {
1688+
for (const [, tools] of this.activeToolExecutions) {
1689+
if (tools.size > 0) {
1690+
return true
1691+
}
1692+
}
1693+
return false
16191694
}
16201695

16211696
/**
@@ -1703,7 +1778,15 @@ export class McpHub {
17031778
shouldAllow: boolean,
17041779
): Promise<void> {
17051780
try {
1706-
await this.updateServerToolList(serverName, source, toolName, "alwaysAllow", shouldAllow)
1781+
// Check if there are active tool executions for this server
1782+
if (this.hasActiveToolExecutions(serverName, source)) {
1783+
console.log(`Skipping server restart for ${serverName} - tools are currently executing`)
1784+
// Update the config file without triggering a restart
1785+
await this.updateServerToolListWithoutRestart(serverName, source, toolName, "alwaysAllow", shouldAllow)
1786+
} else {
1787+
// Normal flow - update and allow restart if needed
1788+
await this.updateServerToolList(serverName, source, toolName, "alwaysAllow", shouldAllow)
1789+
}
17071790
} catch (error) {
17081791
this.showErrorMessage(
17091792
`Failed to toggle always allow for tool "${toolName}" on server "${serverName}" with source "${source}"`,
@@ -1713,6 +1796,114 @@ export class McpHub {
17131796
}
17141797
}
17151798

1799+
/**
1800+
* Update server tool list without triggering a restart
1801+
* This is used when tools are actively running to prevent interruption
1802+
*/
1803+
private async updateServerToolListWithoutRestart(
1804+
serverName: string,
1805+
source: "global" | "project",
1806+
toolName: string,
1807+
listName: "alwaysAllow" | "disabledTools",
1808+
addTool: boolean,
1809+
): Promise<void> {
1810+
// Find the connection with matching name and source
1811+
const connection = this.findConnection(serverName, source)
1812+
1813+
if (!connection) {
1814+
throw new Error(`Server ${serverName} with source ${source} not found`)
1815+
}
1816+
1817+
// Determine the correct config path based on the source
1818+
let configPath: string
1819+
if (source === "project") {
1820+
// Get project MCP config path
1821+
const projectMcpPath = await this.getProjectMcpPath()
1822+
if (!projectMcpPath) {
1823+
throw new Error("Project MCP configuration file not found")
1824+
}
1825+
configPath = projectMcpPath
1826+
} else {
1827+
// Get global MCP settings path
1828+
configPath = await this.getMcpSettingsFilePath()
1829+
}
1830+
1831+
// Normalize path for cross-platform compatibility
1832+
const normalizedPath = process.platform === "win32" ? configPath.replace(/\\/g, "/") : configPath
1833+
1834+
// Read the appropriate config file
1835+
const content = await fs.readFile(normalizedPath, "utf-8")
1836+
const config = JSON.parse(content)
1837+
1838+
if (!config.mcpServers) {
1839+
config.mcpServers = {}
1840+
}
1841+
1842+
if (!config.mcpServers[serverName]) {
1843+
config.mcpServers[serverName] = {
1844+
type: "stdio",
1845+
command: "node",
1846+
args: [], // Default to an empty array; can be set later if needed
1847+
}
1848+
}
1849+
1850+
if (!config.mcpServers[serverName][listName]) {
1851+
config.mcpServers[serverName][listName] = []
1852+
}
1853+
1854+
const targetList = config.mcpServers[serverName][listName]
1855+
const toolIndex = targetList.indexOf(toolName)
1856+
1857+
if (addTool && toolIndex === -1) {
1858+
targetList.push(toolName)
1859+
} else if (!addTool && toolIndex !== -1) {
1860+
targetList.splice(toolIndex, 1)
1861+
}
1862+
1863+
// Write the config file directly without triggering file watcher
1864+
// We'll temporarily disable the watcher to prevent restart
1865+
const watcherKey = source === "project" ? this.projectMcpWatcher : this.settingsWatcher
1866+
const wasWatcherActive = !!watcherKey
1867+
1868+
if (wasWatcherActive && source === "project" && this.projectMcpWatcher) {
1869+
this.projectMcpWatcher.dispose()
1870+
this.projectMcpWatcher = undefined
1871+
} else if (wasWatcherActive && source === "global" && this.settingsWatcher) {
1872+
this.settingsWatcher.dispose()
1873+
this.settingsWatcher = undefined
1874+
}
1875+
1876+
await fs.writeFile(normalizedPath, JSON.stringify(config, null, 2))
1877+
1878+
// Update the in-memory tool state without restarting
1879+
if (connection) {
1880+
// Update the tool's alwaysAllow or enabledForPrompt status in memory
1881+
const tools = connection.server.tools
1882+
if (tools) {
1883+
const tool = tools.find((t) => t.name === toolName)
1884+
if (tool) {
1885+
if (listName === "alwaysAllow") {
1886+
tool.alwaysAllow = addTool
1887+
} else if (listName === "disabledTools") {
1888+
tool.enabledForPrompt = !addTool
1889+
}
1890+
}
1891+
}
1892+
await this.notifyWebviewOfServerChanges()
1893+
}
1894+
1895+
// Re-enable the watcher after a short delay
1896+
if (wasWatcherActive) {
1897+
setTimeout(async () => {
1898+
if (source === "project") {
1899+
await this.watchProjectMcpFile()
1900+
} else {
1901+
await this.watchMcpSettingsFile()
1902+
}
1903+
}, 1000)
1904+
}
1905+
}
1906+
17161907
async toggleToolEnabledForPrompt(
17171908
serverName: string,
17181909
source: "global" | "project",
@@ -1723,7 +1914,22 @@ export class McpHub {
17231914
// When isEnabled is true, we want to remove the tool from the disabledTools list.
17241915
// When isEnabled is false, we want to add the tool to the disabledTools list.
17251916
const addToolToDisabledList = !isEnabled
1726-
await this.updateServerToolList(serverName, source, toolName, "disabledTools", addToolToDisabledList)
1917+
1918+
// Check if there are active tool executions for this server
1919+
if (this.hasActiveToolExecutions(serverName, source)) {
1920+
console.log(`Skipping server restart for ${serverName} - tools are currently executing`)
1921+
// Update the config file without triggering a restart
1922+
await this.updateServerToolListWithoutRestart(
1923+
serverName,
1924+
source,
1925+
toolName,
1926+
"disabledTools",
1927+
addToolToDisabledList,
1928+
)
1929+
} else {
1930+
// Normal flow - update and allow restart if needed
1931+
await this.updateServerToolList(serverName, source, toolName, "disabledTools", addToolToDisabledList)
1932+
}
17271933
} catch (error) {
17281934
this.showErrorMessage(`Failed to update settings for tool ${toolName}`, error)
17291935
throw error // Re-throw to ensure the error is properly handled

0 commit comments

Comments
 (0)