Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions go.mod
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+github.com/gdamore/tcell/v2 v2.13.8
+github.com/rivo/tview v0.42.0

These two dependencies have changed from indirect to direct dependencies, but there is no code referencing them in the PR. This could be a side effect of go mod tidy, or there may be other uncommitted local code that introduced these dependencies.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's go mod tidy effect

Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ require (
github.com/bwmarrin/discordgo v0.29.0
github.com/caarlos0/env/v11 v11.3.1
github.com/chzyer/readline v1.5.1
github.com/gdamore/tcell/v2 v2.13.8
github.com/google/uuid v1.6.0
github.com/gorilla/websocket v1.5.3
github.com/larksuite/oapi-sdk-go/v3 v3.5.3
Expand All @@ -16,6 +17,7 @@ require (
github.com/mymmrac/telego v1.6.0
github.com/open-dingtalk/dingtalk-stream-sdk-go v0.9.1
github.com/openai/openai-go/v3 v3.22.0
github.com/rivo/tview v0.42.0
github.com/slack-go/slack v0.17.3
github.com/spf13/cobra v1.10.2
github.com/stretchr/testify v1.11.1
Expand All @@ -35,7 +37,6 @@ require (
github.com/dustin/go-humanize v1.0.1 // indirect
github.com/elliotchance/orderedmap/v3 v3.1.0 // indirect
github.com/gdamore/encoding v1.0.1 // indirect
github.com/gdamore/tcell/v2 v2.13.8 // indirect
github.com/inconshreveable/mousetrap v1.1.0 // indirect
github.com/lucasb-eyer/go-colorful v1.3.0 // indirect
github.com/mattn/go-colorable v0.1.14 // indirect
Expand All @@ -44,7 +45,6 @@ require (
github.com/petermattis/goid v0.0.0-20260113132338-7c7de50cc741 // indirect
github.com/pmezard/go-difflib v1.0.0 // indirect
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/rivo/tview v0.42.0 // indirect
github.com/rivo/uniseg v0.4.7 // indirect
github.com/rs/zerolog v1.34.0 // indirect
github.com/spf13/pflag v1.0.10 // indirect
Expand Down
21 changes: 11 additions & 10 deletions pkg/agent/loop.go
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,17 @@ func (al *AgentLoop) Run(ctx context.Context) error {
// Initialize MCP servers for all agents
if al.cfg.Tools.MCP.Enabled {
mcpManager := mcp.NewManager()
// Ensure MCP connections are cleaned up on exit, regardless of initialization success
// This fixes resource leak when LoadFromMCPConfig partially succeeds then fails
defer func() {
if err := mcpManager.Close(); err != nil {
logger.ErrorCF("agent", "Failed to close MCP manager",
map[string]any{
"error": err.Error(),
})
}
}()

defaultAgent := al.registry.GetDefaultAgent()
var workspacePath string
if defaultAgent != nil && defaultAgent.Workspace != "" {
Expand All @@ -192,16 +203,6 @@ func (al *AgentLoop) Run(ctx context.Context) error {
"error": err.Error(),
})
} else {
// Ensure MCP connections are cleaned up on exit, only if initialization succeeded
defer func() {
if err := mcpManager.Close(); err != nil {
logger.ErrorCF("agent", "Failed to close MCP manager",
map[string]any{
"error": err.Error(),
})
}
}()

// Register MCP tools for all agents
servers := mcpManager.GetServers()
uniqueTools := 0
Expand Down
24 changes: 15 additions & 9 deletions pkg/mcp/manager.go
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (
"path/filepath"
"strings"
"sync"
"sync/atomic"

"github.com/modelcontextprotocol/go-sdk/mcp"

Expand Down Expand Up @@ -108,7 +109,7 @@ type ServerConnection struct {
type Manager struct {
servers map[string]*ServerConnection
mu sync.RWMutex
closed bool
closed atomic.Bool // changed from bool to atomic.Bool to avoid TOCTOU race
wg sync.WaitGroup // tracks in-flight CallTool calls
}

Expand Down Expand Up @@ -440,14 +441,20 @@ func (m *Manager) CallTool(
serverName, toolName string,
arguments map[string]any,
) (*mcp.CallToolResult, error) {
// Check if closed before acquiring lock (fast path)
if m.closed.Load() {
return nil, fmt.Errorf("manager is closed")
}

m.mu.RLock()
if m.closed {
// Double-check after acquiring lock to prevent TOCTOU race
if m.closed.Load() {
m.mu.RUnlock()
return nil, fmt.Errorf("manager is closed")
}
conn, ok := m.servers[serverName]
if ok {
m.wg.Add(1)
m.wg.Add(1) // Add to WaitGroup while holding the lock
}
m.mu.RUnlock()

Expand All @@ -471,15 +478,14 @@ func (m *Manager) CallTool(

// Close closes all server connections
func (m *Manager) Close() error {
m.mu.Lock()
if m.closed {
m.mu.Unlock()
return nil
// Use Swap to atomically set closed=true and get the previous value
// This prevents TOCTOU race with CallTool's closed check
if m.closed.Swap(true) {
return nil // already closed
}
m.closed = true
m.mu.Unlock()

// Wait for all in-flight CallTool calls to finish before closing sessions
// After closed=true is set, no new CallTool can start (they check closed first)
m.wg.Wait()

m.mu.Lock()
Expand Down
2 changes: 1 addition & 1 deletion pkg/mcp/manager_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -268,7 +268,7 @@ func TestGetAllTools_FiltersEmptyTools(t *testing.T) {
func TestCallTool_ErrorsForClosedOrMissingServer(t *testing.T) {
t.Run("manager closed", func(t *testing.T) {
mgr := NewManager()
mgr.closed = true
mgr.closed.Store(true)

_, err := mgr.CallTool(context.Background(), "s1", "tool", nil)
if err == nil || !strings.Contains(err.Error(), "manager is closed") {
Expand Down