Add per-server MCP network timeout configuration

This commit is contained in:
Matt Rubens
2025-01-26 15:32:17 -05:00
parent 3deeb0c28b
commit 57518e10b3
7 changed files with 341 additions and 91 deletions

View File

@@ -1177,6 +1177,16 @@ export class ClineProvider implements vscode.WebviewViewProvider {
}
await this.postStateToWebview()
break
case "updateMcpTimeout":
if (message.serverName && typeof message.timeout === "number") {
try {
await this.mcpHub?.updateServerTimeout(message.serverName, message.timeout)
} catch (error) {
console.error(`Failed to update timeout for ${message.serverName}:`, error)
vscode.window.showErrorMessage(`Failed to update server timeout`)
}
}
break
case "updateCustomMode":
if (message.modeConfig) {
await this.customModesManager.updateCustomMode(message.modeConfig.slug, message.modeConfig)

View File

@@ -35,12 +35,13 @@ export type McpConnection = {
// StdioServerParameters
const AlwaysAllowSchema = z.array(z.string()).default([])
const StdioConfigSchema = z.object({
export const StdioConfigSchema = z.object({
command: z.string(),
args: z.array(z.string()).optional(),
env: z.record(z.string()).optional(),
alwaysAllow: AlwaysAllowSchema.optional(),
disabled: z.boolean().optional(),
timeout: z.number().min(1).max(3600).optional().default(60),
})
const McpSettingsSchema = z.object({
@@ -238,28 +239,6 @@ export class McpHub {
}
transport.start = async () => {} // No-op now, .connect() won't fail
// // Set up notification handlers
// client.setNotificationHandler(
// // @ts-ignore-next-line
// { method: "notifications/tools/list_changed" },
// async () => {
// console.log(`Tools changed for server: ${name}`)
// connection.server.tools = await this.fetchTools(name)
// await this.notifyWebviewOfServerChanges()
// },
// )
// client.setNotificationHandler(
// // @ts-ignore-next-line
// { method: "notifications/resources/list_changed" },
// async () => {
// console.log(`Resources changed for server: ${name}`)
// connection.server.resources = await this.fetchResources(name)
// connection.server.resourceTemplates = await this.fetchResourceTemplates(name)
// await this.notifyWebviewOfServerChanges()
// },
// )
// Connect
await client.connect(transport)
connection.server.status = "connected"
@@ -339,10 +318,6 @@ export class McpHub {
const connection = this.connections.find((conn) => conn.server.name === name)
if (connection) {
try {
// connection.client.removeNotificationHandler("notifications/tools/list_changed")
// connection.client.removeNotificationHandler("notifications/resources/list_changed")
// connection.client.removeNotificationHandler("notifications/stderr")
// connection.client.removeNotificationHandler("notifications/stderr")
await connection.transport.close()
await connection.client.close()
} catch (error) {
@@ -468,8 +443,6 @@ export class McpHub {
})
}
// Public methods for server management
public async toggleServerDisabled(serverName: string, disabled: boolean): Promise<void> {
let settingsPath: string
try {
@@ -545,6 +518,59 @@ export class McpHub {
}
}
public async updateServerTimeout(serverName: string, timeout: number): Promise<void> {
let settingsPath: string
try {
settingsPath = await this.getMcpSettingsFilePath()
// Ensure the settings file exists and is accessible
try {
await fs.access(settingsPath)
} catch (error) {
console.error("Settings file not accessible:", error)
throw new Error("Settings file not accessible")
}
const content = await fs.readFile(settingsPath, "utf-8")
const config = JSON.parse(content)
// Validate the config structure
if (!config || typeof config !== "object") {
throw new Error("Invalid config structure")
}
if (!config.mcpServers || typeof config.mcpServers !== "object") {
config.mcpServers = {}
}
if (config.mcpServers[serverName]) {
// Create a new server config object to ensure clean structure
const serverConfig = {
...config.mcpServers[serverName],
timeout,
}
config.mcpServers[serverName] = serverConfig
// Write the entire config back
const updatedConfig = {
mcpServers: config.mcpServers,
}
await fs.writeFile(settingsPath, JSON.stringify(updatedConfig, null, 2))
await this.notifyWebviewOfServerChanges()
}
} catch (error) {
console.error("Failed to update server timeout:", error)
if (error instanceof Error) {
console.error("Error details:", error.message, error.stack)
}
vscode.window.showErrorMessage(
`Failed to update server timeout: ${error instanceof Error ? error.message : String(error)}`,
)
throw error
}
}
async readResource(serverName: string, uri: string): Promise<McpResourceResponse> {
const connection = this.connections.find((conn) => conn.server.name === serverName)
if (!connection) {
@@ -579,6 +605,16 @@ export class McpHub {
throw new Error(`Server "${serverName}" is disabled and cannot be used`)
}
let timeout: number
try {
const parsedConfig = StdioConfigSchema.parse(JSON.parse(connection.server.config))
timeout = (parsedConfig.timeout ?? 60) * 1000
} catch (error) {
console.error("Failed to parse server config for timeout:", error)
// Default to 60 seconds if parsing fails
timeout = 60 * 1000
}
return await connection.client.request(
{
method: "tools/call",
@@ -588,6 +624,9 @@ export class McpHub {
},
},
CallToolResultSchema,
{
timeout,
},
)
}

View File

@@ -2,8 +2,8 @@ import type { McpHub as McpHubType } from "../McpHub"
import type { ClineProvider } from "../../../core/webview/ClineProvider"
import type { ExtensionContext, Uri } from "vscode"
import type { McpConnection } from "../McpHub"
import { StdioConfigSchema } from "../McpHub"
const vscode = require("vscode")
const fs = require("fs/promises")
const { McpHub } = require("../McpHub")
@@ -280,6 +280,7 @@ describe("McpHub", () => {
},
},
expect.any(Object),
expect.objectContaining({ timeout: 60000 }), // Default 60 second timeout
)
})
@@ -288,5 +289,192 @@ describe("McpHub", () => {
"No connection found for server: non-existent-server",
)
})
describe("timeout configuration", () => {
it("should validate timeout values", () => {
// Test valid timeout values
const validConfig = {
command: "test",
timeout: 60,
}
expect(() => StdioConfigSchema.parse(validConfig)).not.toThrow()
// Test invalid timeout values
const invalidConfigs = [
{ command: "test", timeout: 0 }, // Too low
{ command: "test", timeout: 3601 }, // Too high
{ command: "test", timeout: -1 }, // Negative
]
invalidConfigs.forEach((config) => {
expect(() => StdioConfigSchema.parse(config)).toThrow()
})
})
it("should use default timeout of 60 seconds if not specified", async () => {
const mockConnection: McpConnection = {
server: {
name: "test-server",
config: JSON.stringify({ command: "test" }), // No timeout specified
status: "connected",
},
client: {
request: jest.fn().mockResolvedValue({ content: [] }),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]
await mcpHub.callTool("test-server", "test-tool")
expect(mockConnection.client.request).toHaveBeenCalledWith(
expect.anything(),
expect.anything(),
expect.objectContaining({ timeout: 60000 }), // 60 seconds in milliseconds
)
})
it("should apply configured timeout to tool calls", async () => {
const mockConnection: McpConnection = {
server: {
name: "test-server",
config: JSON.stringify({ command: "test", timeout: 120 }), // 2 minutes
status: "connected",
},
client: {
request: jest.fn().mockResolvedValue({ content: [] }),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]
await mcpHub.callTool("test-server", "test-tool")
expect(mockConnection.client.request).toHaveBeenCalledWith(
expect.anything(),
expect.anything(),
expect.objectContaining({ timeout: 120000 }), // 120 seconds in milliseconds
)
})
})
describe("updateServerTimeout", () => {
it("should update server timeout in settings file", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
command: "node",
args: ["test.js"],
timeout: 60,
},
},
}
// Mock reading initial config
;(fs.readFile as jest.Mock).mockResolvedValueOnce(JSON.stringify(mockConfig))
await mcpHub.updateServerTimeout("test-server", 120)
// Verify the config was updated correctly
const writeCall = (fs.writeFile as jest.Mock).mock.calls[0]
const writtenConfig = JSON.parse(writeCall[1])
expect(writtenConfig.mcpServers["test-server"].timeout).toBe(120)
})
it("should fallback to default timeout when config has invalid timeout", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
command: "node",
args: ["test.js"],
timeout: 60,
},
},
}
// Mock initial read
;(fs.readFile as jest.Mock).mockResolvedValueOnce(JSON.stringify(mockConfig))
// Update with invalid timeout
await mcpHub.updateServerTimeout("test-server", 3601)
// Config is written
expect(fs.writeFile).toHaveBeenCalled()
// Setup connection with invalid timeout
const mockConnection: McpConnection = {
server: {
name: "test-server",
config: JSON.stringify({
command: "node",
args: ["test.js"],
timeout: 3601, // Invalid timeout
}),
status: "connected",
},
client: {
request: jest.fn().mockResolvedValue({ content: [] }),
} as any,
transport: {} as any,
}
mcpHub.connections = [mockConnection]
// Call tool - should use default timeout
await mcpHub.callTool("test-server", "test-tool")
// Verify default timeout was used
expect(mockConnection.client.request).toHaveBeenCalledWith(
expect.anything(),
expect.anything(),
expect.objectContaining({ timeout: 60000 }), // Default 60 seconds
)
})
it("should accept valid timeout values", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
command: "node",
args: ["test.js"],
timeout: 60,
},
},
}
;(fs.readFile as jest.Mock).mockResolvedValueOnce(JSON.stringify(mockConfig))
// Test valid timeout values
const validTimeouts = [1, 60, 3600]
for (const timeout of validTimeouts) {
await mcpHub.updateServerTimeout("test-server", timeout)
expect(fs.writeFile).toHaveBeenCalled()
jest.clearAllMocks() // Reset for next iteration
;(fs.readFile as jest.Mock).mockResolvedValueOnce(JSON.stringify(mockConfig))
}
})
it("should notify webview after updating timeout", async () => {
const mockConfig = {
mcpServers: {
"test-server": {
command: "node",
args: ["test.js"],
timeout: 60,
},
},
}
;(fs.readFile as jest.Mock).mockResolvedValueOnce(JSON.stringify(mockConfig))
await mcpHub.updateServerTimeout("test-server", 120)
expect(mockProvider.postMessageToWebview).toHaveBeenCalledWith(
expect.objectContaining({
type: "mcpServers",
}),
)
})
})
})
})

View File

@@ -51,6 +51,7 @@ export interface WebviewMessage {
| "restartMcpServer"
| "toggleToolAlwaysAllow"
| "toggleMcpServer"
| "updateMcpTimeout"
| "fuzzyMatchThreshold"
| "preferredLanguage"
| "writeDelayMs"
@@ -99,6 +100,7 @@ export interface WebviewMessage {
query?: string
slug?: string
modeConfig?: ModeConfig
timeout?: number
}
export type ClineAskResponse = "yesButtonClicked" | "noButtonClicked" | "messageResponse"

View File

@@ -7,6 +7,7 @@ export type McpServer = {
resources?: McpResource[]
resourceTemplates?: McpResourceTemplate[]
disabled?: boolean
timeout?: number
}
export type McpTool = {