From 6dd3b09d897536734eee297d108c58b059b0fe08 Mon Sep 17 00:00:00 2001 From: Carlos Alexandro Becker Date: Tue, 15 Jul 2025 16:07:41 -0300 Subject: [PATCH 1/2] refactor: use sync primitives in GetMcpTools Signed-off-by: Carlos Alexandro Becker --- cmd/root.go | 17 ------ internal/llm/agent/agent.go | 2 +- internal/llm/agent/mcp-tools.go | 103 +++++++++++++++++++------------- 3 files changed, 63 insertions(+), 59 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index ad558173c6eb1dd1bf4fdda30524ca7b04793ff5..fe17181c110c68a032495ee806c91651daf3631e 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -92,9 +92,6 @@ to assist developers in writing, debugging, and understanding code directly from } defer app.Shutdown() - // Initialize MCP tools early for both modes - initMCPTools(ctx, app, cfg) - prompt, err = maybePrependStdin(prompt) if err != nil { slog.Error(fmt.Sprintf("Failed to read from stdin: %v", err)) @@ -126,20 +123,6 @@ to assist developers in writing, debugging, and understanding code directly from }, } -func initMCPTools(ctx context.Context, app *app.App, cfg *config.Config) { - go func() { - defer log.RecoverPanic("MCP-goroutine", nil) - - // Create a context with timeout for the initial MCP tools fetch - ctxWithTimeout, cancel := context.WithTimeout(ctx, 30*time.Second) - defer cancel() - - // Set this up once with proper error handling - agent.GetMcpTools(ctxWithTimeout, app.Permissions, cfg) - slog.Info("MCP message handling goroutine exiting") - }() -} - func Execute() { if err := fang.Execute( context.Background(), diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go index 7cd01e91900a6b0a2720a092e6740ab7d989fb6d..0cc8b71029191a7e05fe1e77b437066d3cdd40c5 100644 --- a/internal/llm/agent/agent.go +++ b/internal/llm/agent/agent.go @@ -94,7 +94,7 @@ func NewAgent( ) (Service, error) { ctx := context.Background() cfg := config.Get() - otherTools := GetMcpTools(ctx, permissions, cfg) + otherTools := GetMCPTools(ctx, permissions, cfg) if len(lspClients) > 0 { otherTools = append(otherTools, tools.NewDiagnosticsTool(lspClients)) } diff --git a/internal/llm/agent/mcp-tools.go b/internal/llm/agent/mcp-tools.go index d8610b557896272d94c76c608b9bc00347655be4..4ed4c023719642724c8dd49cbe3293376e3f3c57 100644 --- a/internal/llm/agent/mcp-tools.go +++ b/internal/llm/agent/mcp-tools.go @@ -5,6 +5,7 @@ import ( "encoding/json" "fmt" "log/slog" + "sync" "github.com/charmbracelet/crush/internal/config" "github.com/charmbracelet/crush/internal/llm/tools" @@ -154,8 +155,6 @@ func NewMcpTool(name string, tool mcp.Tool, permissions permission.Service, mcpC } } -var mcpTools []tools.BaseTool - func getTools(ctx context.Context, name string, m config.MCPConfig, permissions permission.Service, c MCPClient, workingDir string) []tools.BaseTool { var stdioTools []tools.BaseTool initRequest := mcp.InitializeRequest{} @@ -183,50 +182,72 @@ func getTools(ctx context.Context, name string, m config.MCPConfig, permissions return stdioTools } -func GetMcpTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool { - if len(mcpTools) > 0 { - return mcpTools - } +var ( + mcpToolsOnce sync.Once + mcpTools []tools.BaseTool +) + +func GetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool { + mcpToolsOnce.Do(func() { + mcpTools = doGetMCPTools(ctx, permissions, cfg) + }) + return mcpTools +} + +func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool { + var mu sync.Mutex + var wg sync.WaitGroup + var result []tools.BaseTool for name, m := range cfg.MCP { if m.Disabled { slog.Debug("skipping disabled mcp", "name", name) continue } - switch m.Type { - case config.MCPStdio: - c, err := client.NewStdioMCPClient( - m.Command, - m.Env, - m.Args..., - ) - if err != nil { - slog.Error("error creating mcp client", "error", err) - continue + wg.Add(1) + go func(name string, m config.MCPConfig) { + defer wg.Done() + switch m.Type { + case config.MCPStdio: + c, err := client.NewStdioMCPClient( + m.Command, + m.Env, + m.Args..., + ) + if err != nil { + slog.Error("error creating mcp client", "error", err) + return + } + + mu.Lock() + result = append(result, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...) + mu.Unlock() + case config.MCPHttp: + c, err := client.NewStreamableHttpClient( + m.URL, + transport.WithHTTPHeaders(m.Headers), + ) + if err != nil { + slog.Error("error creating mcp client", "error", err) + return + } + mu.Lock() + result = append(result, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...) + mu.Unlock() + case config.MCPSse: + c, err := client.NewSSEMCPClient( + m.URL, + client.WithHeaders(m.Headers), + ) + if err != nil { + slog.Error("error creating mcp client", "error", err) + return + } + mu.Lock() + result = append(result, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...) + mu.Unlock() } - - mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...) - case config.MCPHttp: - c, err := client.NewStreamableHttpClient( - m.URL, - transport.WithHTTPHeaders(m.Headers), - ) - if err != nil { - slog.Error("error creating mcp client", "error", err) - continue - } - mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...) - case config.MCPSse: - c, err := client.NewSSEMCPClient( - m.URL, - client.WithHeaders(m.Headers), - ) - if err != nil { - slog.Error("error creating mcp client", "error", err) - continue - } - mcpTools = append(mcpTools, getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...) - } + }(name, m) } - - return mcpTools + wg.Wait() + return result } From 0427603a44382ad2ab8387e3c172697d3afe40e2 Mon Sep 17 00:00:00 2001 From: Carlos Alexandro Becker Date: Tue, 15 Jul 2025 16:41:36 -0300 Subject: [PATCH 2/2] fix: imports Signed-off-by: Carlos Alexandro Becker --- cmd/root.go | 3 --- 1 file changed, 3 deletions(-) diff --git a/cmd/root.go b/cmd/root.go index fe17181c110c68a032495ee806c91651daf3631e..b3ea36c8a976face0bb29c3d77e0d2c82dfb1399 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -6,14 +6,11 @@ import ( "io" "log/slog" "os" - "time" tea "github.com/charmbracelet/bubbletea/v2" "github.com/charmbracelet/crush/internal/app" "github.com/charmbracelet/crush/internal/config" "github.com/charmbracelet/crush/internal/db" - "github.com/charmbracelet/crush/internal/llm/agent" - "github.com/charmbracelet/crush/internal/log" "github.com/charmbracelet/crush/internal/tui" "github.com/charmbracelet/crush/internal/version" "github.com/charmbracelet/fang"