From 57f4acf00f0232aa1d1b8bc167cd7bce372d799f Mon Sep 17 00:00:00 2001
From: dj-oyu <68707227+dj-oyu@users.noreply.github.com>
Date: Fri, 20 Mar 2026 06:34:39 +0900
Subject: [PATCH] feat: add per-server timeout for MCP tool calls

MCP CallTool had no timeout, so a hanging remote API could block the
agent indefinitely. Add a configurable timeout (default 60s) per MCP
server via `"timeout": <seconds>` in config, applied as
context.WithTimeout in Manager.CallTool.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 pkg/config/config.go |  2 ++
 pkg/mcp/manager.go   | 16 ++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 776e472b7..98662a4ff 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -846,6 +846,8 @@ type MCPServerConfig struct {
 	URL string `json:"url,omitempty"`
 	// Headers are HTTP headers to send with requests (sse/http only)
 	Headers map[string]string `json:"headers,omitempty"`
+	// Timeout is the maximum duration in seconds for tool calls to this server (default: 60)
+	Timeout int `json:"timeout,omitempty"`
 }
 
 // MCPConfig defines configuration for all MCP servers
diff --git a/pkg/mcp/manager.go b/pkg/mcp/manager.go
index 7b63cc979..861999981 100644
--- a/pkg/mcp/manager.go
+++ b/pkg/mcp/manager.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 	"sync"
 	"sync/atomic"
+	"time"
 
 	"github.com/modelcontextprotocol/go-sdk/mcp"
 
@@ -103,6 +104,7 @@ type ServerConnection struct {
 	Client  *mcp.Client
 	Session *mcp.ClientSession
 	Tools   []*mcp.Tool
+	Timeout time.Duration
 }
 
 // Manager manages multiple MCP server connections
@@ -401,6 +403,12 @@ func (m *Manager) ConnectServer(
 			})
 	}
 
+	// Resolve timeout (default 60s)
+	timeout := 60 * time.Second
+	if cfg.Timeout > 0 {
+		timeout = time.Duration(cfg.Timeout) * time.Second
+	}
+
 	// Store connection
 	m.mu.Lock()
 	m.servers[name] = &ServerConnection{
@@ -408,6 +416,7 @@ func (m *Manager) ConnectServer(
 		Client:  client,
 		Session: session,
 		Tools:   tools,
+		Timeout: timeout,
 	}
 	m.mu.Unlock()
 
@@ -463,6 +472,13 @@ func (m *Manager) CallTool(
 	}
 	defer m.wg.Done()
 
+	// Apply per-server timeout
+	if conn.Timeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, conn.Timeout)
+		defer cancel()
+	}
+
 	params := &mcp.CallToolParams{
 		Name:      toolName,
 		Arguments: arguments,