@@ -11,34 +11,28 @@ import (
1111 "github.com/charmbracelet/crush/internal/config"
1212 "github.com/charmbracelet/crush/internal/csync"
1313 "github.com/charmbracelet/crush/internal/llm/tools"
14+ "github.com/charmbracelet/crush/internal/version"
1415
1516 "github.com/charmbracelet/crush/internal/permission"
16− "github.com/charmbracelet/crush/internal/version"
1717
1818 "github.com/mark3labs/mcp-go/client"
1919 "github.com/mark3labs/mcp-go/client/transport"
2020 "github.com/mark3labs/mcp-go/mcp"
2121)
2222
23+var (
24+ mcpToolsOnce sync.Once
25+ mcpTools []tools.BaseTool
26+ mcpClients = csync.NewMap[string, *client.Client]()
27+)
28+
2329type McpTool struct {
2430 mcpName string
2531 tool mcp.Tool
26− client MCPClient
27− mcpConfig config.MCPConfig
2832 permissions permission.Service
2933 workingDir string
3034}
3135
32−type MCPClient interface {
33− Initialize(
34− ctx context.Context,
35− request mcp.InitializeRequest,
36− ) (*mcp.InitializeResult, error)
37− ListTools(ctx context.Context, request mcp.ListToolsRequest) (*mcp.ListToolsResult, error)
38− CallTool(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error)
39− Close() error
40−}
41−
4236func (b *McpTool) Name() string {
4337 return fmt.Sprintf("mcp_%s_%s", b.mcpName, b.tool.Name)
4438}
@@ -56,27 +50,21 @@ func (b *McpTool) Info() tools.ToolInfo {
5650 }
5751}
5852
59−func runTool(ctx context.Context, c MCPClient, toolName string, input string) (tools.ToolResponse, error) {
60− initRequest := mcp.InitializeRequest{}
61− initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
62− initRequest.Params.ClientInfo = mcp.Implementation{
63− Name: "Crush",
64− Version: version.Version,
65− }
66−
67− _, err := c.Initialize(ctx, initRequest)
68− if err != nil {
69− return tools.NewTextErrorResponse(err.Error()), nil
70− }
71−
72− toolRequest := mcp.CallToolRequest{}
73− toolRequest.Params.Name = toolName
53+func runTool(ctx context.Context, name, toolName string, input string) (tools.ToolResponse, error) {
7454 var args map[string]any
75− if err = json.Unmarshal([]byte(input), &args); err != nil {
55+ if err := json.Unmarshal([]byte(input), &args); err != nil {
7656 return tools.NewTextErrorResponse(fmt.Sprintf("error parsing parameters: %s", err)), nil
7757 }
78− toolRequest.Params.Arguments = args
79− result, err := c.CallTool(ctx, toolRequest)
58+ c, ok := mcpClients.Get(name)
59+ if !ok {
60+ return tools.NewTextErrorResponse("mcp '" + name + "' not available"), nil
61+ }
62+ result, err := c.CallTool(ctx, mcp.CallToolRequest{
63+ Params: mcp.CallToolParams{
64+ Name: toolName,
65+ Arguments: args,
66+ },
67+ })
8068 if err != nil {
8169 return tools.NewTextErrorResponse(err.Error()), nil
8270 }
@@ -114,56 +102,34 @@ func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolRes
114102 return tools.ToolResponse{}, permission.ErrorPermissionDenied
115103 }
116104
117− return runTool(ctx, b.client, b.tool.Name, params.Input)
105+ return runTool(ctx, b.mcpName, b.tool.Name, params.Input)
118106}
119107
120−func NewMcpTool(name string, c MCPClient, tool mcp.Tool, permissions permission.Service, mcpConfig config.MCPConfig, workingDir string) tools.BaseTool {
121− return &McpTool{
122− mcpName: name,
123− client: c,
124− tool: tool,
125− mcpConfig: mcpConfig,
126− permissions: permissions,
127− workingDir: workingDir,
128− }
129−}
130−
131−func getTools(ctx context.Context, name string, m config.MCPConfig, permissions permission.Service, c MCPClient, workingDir string) []tools.BaseTool {
132− var stdioTools []tools.BaseTool
133− initRequest := mcp.InitializeRequest{}
134− initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
135− initRequest.Params.ClientInfo = mcp.Implementation{
136− Name: "Crush",
137− Version: version.Version,
138− }
139−
140− _, err := c.Initialize(ctx, initRequest)
141− if err != nil {
142− slog.Error("error initializing mcp client", "error", err)
143− return stdioTools
144− }
145− toolsRequest := mcp.ListToolsRequest{}
146− tools, err := c.ListTools(ctx, toolsRequest)
108+func getTools(ctx context.Context, name string, permissions permission.Service, c *client.Client, workingDir string) []tools.BaseTool {
109+ result, err := c.ListTools(ctx, mcp.ListToolsRequest{})
147110 if err != nil {
148111 slog.Error("error listing tools", "error", err)
149− return stdioTools
112+ c.Close()
113+ mcpClients.Del(name)
114+ return nil
150115 }
151− for _, t := range tools.Tools {
152− stdioTools = append(stdioTools, NewMcpTool(name, c, t, permissions, m, workingDir))
116+ mcpTools := make([]tools.BaseTool, 0, len(result.Tools))
117+ for _, tool := range result.Tools {
118+ mcpTools = append(mcpTools, &McpTool{
119+ mcpName: name,
120+ tool: tool,
121+ permissions: permissions,
122+ workingDir: workingDir,
123+ })
153124 }
154− return stdioTools
125+ return mcpTools
155126}
156127
157−var (
158− mcpToolsOnce sync.Once
159− mcpTools []tools.BaseTool
160−)
161−
162−func GetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {
163− mcpToolsOnce.Do(func() {
164− mcpTools = doGetMCPTools(ctx, permissions, cfg)
165− })
166− return mcpTools
128+// CloseMCPClients closes all MCP clients. This should be called during application shutdown.
129+func CloseMCPClients() {
130+ for c := range mcpClients.Seq() {
131+ _ = c.Close()
132+ }
167133}
168134
169135func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {
@@ -177,42 +143,59 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con
177143 wg.Add(1)
178144 go func(name string, m config.MCPConfig) {
179145 defer wg.Done()
180− switch m.Type {
181− case config.MCPStdio:
182− c, err := client.NewStdioMCPClient(
183− m.Command,
184− m.ResolvedEnv(),
185− m.Args...,
186− )
187− if err != nil {
188− slog.Error("error creating mcp client", "error", err)
189− return
190− }
191−
192− result.Append(getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)
193− case config.MCPHttp:
194− c, err := client.NewStreamableHttpClient(
195− m.URL,
196− transport.WithHTTPHeaders(m.ResolvedHeaders()),
197− )
198− if err != nil {
199− slog.Error("error creating mcp client", "error", err)
200− return
201− }
202− result.Append(getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)
203− case config.MCPSse:
204− c, err := client.NewSSEMCPClient(
205− m.URL,
206− client.WithHeaders(m.ResolvedHeaders()),
207− )
208− if err != nil {
209− slog.Error("error creating mcp client", "error", err)
210− return
211− }
212− result.Append(getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)
146+ c, err := doGetClient(m)
147+ if err != nil {
148+ slog.Error("error creating mcp client", "error", err)
149+ return
150+ }
151+ if err := doInitClient(ctx, name, c); err != nil {
152+ slog.Error("error initializing mcp client", "error", err)
153+ return
213154 }
155+ result.Append(getTools(ctx, name, permissions, c, cfg.WorkingDir())...)
214156 }(name, m)
215157 }
216158 wg.Wait()
217159 return slices.Collect(result.Seq())
218160}
161+
162+func doInitClient(ctx context.Context, name string, c *client.Client) error {
163+ initRequest := mcp.InitializeRequest{
164+ Params: mcp.InitializeParams{
165+ ProtocolVersion: mcp.LATEST_PROTOCOL_VERSION,
166+ ClientInfo: mcp.Implementation{
167+ Name: "Crush",
168+ Version: version.Version,
169+ },
170+ },
171+ }
172+ if _, err := c.Initialize(ctx, initRequest); err != nil {
173+ c.Close()
174+ return err
175+ }
176+ mcpClients.Set(name, c)
177+ return nil
178+}
179+
180+func doGetClient(m config.MCPConfig) (*client.Client, error) {
181+ switch m.Type {
182+ case config.MCPStdio:
183+ return client.NewStdioMCPClient(
184+ m.Command,
185+ m.ResolvedEnv(),
186+ m.Args...,
187+ )
188+ case config.MCPHttp:
189+ return client.NewStreamableHttpClient(
190+ m.URL,
191+ transport.WithHTTPHeaders(m.ResolvedHeaders()),
192+ )
193+ case config.MCPSse:
194+ return client.NewSSEMCPClient(
195+ m.URL,
196+ client.WithHeaders(m.ResolvedHeaders()),
197+ )
198+ default:
199+ return nil, fmt.Errorf("unsupported mcp type: %s", m.Type)
200+ }
201+}