git.iko.soy/crush
on main help sign in
fix: do not init MCP client on every tool request (#459)
cd3ef8dCarlos Alexandro Becker
internal/app/app.go+4
compare with main
unifiedsplit
@@ -268,6 +268,10 @@ func (app *App) InitCoderAgent() error {
268 slog.Error("Failed to create coder agent", "err", err)
268 slog.Error("Failed to create coder agent", "err", err)
269 return err
269 return err
270 }
270 }
271+
272+ // Add MCP client cleanup to shutdown process
273+ app.cleanupFuncs = append(app.cleanupFuncs, agent.CloseMCPClients)
274+
271 setupSubscriber(app.eventsCtx, app.serviceEventsWG, "coderAgent", app.CoderAgent.Subscribe, app.events)
275 setupSubscriber(app.eventsCtx, app.serviceEventsWG, "coderAgent", app.CoderAgent.Subscribe, app.events)
272 return nil
276 return nil
273}
277}
internal/llm/agent/agent.go+3−1
@@ -189,7 +189,9 @@ func NewAgent(
189 tools.NewWriteTool(lspClients, permissions, history, cwd),
189 tools.NewWriteTool(lspClients, permissions, history, cwd),
190 }
190 }
191
191
192+ mcpToolsOnce.Do(func() {
192− mcpTools := GetMCPTools(ctx, permissions, cfg)
193+ mcpTools = doGetMCPTools(ctx, permissions, cfg)
194+ })
193 allTools = append(allTools, mcpTools...)
195 allTools = append(allTools, mcpTools...)
194
196
195 if len(lspClients) > 0 {
197 if len(lspClients) > 0 {
internal/llm/agent/mcp-tools.go+89−106
@@ -11,34 +11,28 @@ import (
11 "github.com/charmbracelet/crush/internal/config"
11 "github.com/charmbracelet/crush/internal/config"
12 "github.com/charmbracelet/crush/internal/csync"
12 "github.com/charmbracelet/crush/internal/csync"
13 "github.com/charmbracelet/crush/internal/llm/tools"
13 "github.com/charmbracelet/crush/internal/llm/tools"
14+ "github.com/charmbracelet/crush/internal/version"
14
15
15 "github.com/charmbracelet/crush/internal/permission"
16 "github.com/charmbracelet/crush/internal/permission"
16− "github.com/charmbracelet/crush/internal/version"
17
17
18 "github.com/mark3labs/mcp-go/client"
18 "github.com/mark3labs/mcp-go/client"
19 "github.com/mark3labs/mcp-go/client/transport"
19 "github.com/mark3labs/mcp-go/client/transport"
20 "github.com/mark3labs/mcp-go/mcp"
20 "github.com/mark3labs/mcp-go/mcp"
21)
21)
22
22
23+var (
24+ mcpToolsOnce sync.Once
25+ mcpTools []tools.BaseTool
26+ mcpClients = csync.NewMap[string, *client.Client]()
27+)
28+
23type McpTool struct {
29type McpTool struct {
24 mcpName string
30 mcpName string
25 tool mcp.Tool
31 tool mcp.Tool
26− client MCPClient
27− mcpConfig config.MCPConfig
28 permissions permission.Service
32 permissions permission.Service
29 workingDir string
33 workingDir string
30}
34}
31
35
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−
42func (b *McpTool) Name() string {
36func (b *McpTool) Name() string {
43 return fmt.Sprintf("mcp_%s_%s", b.mcpName, b.tool.Name)
37 return fmt.Sprintf("mcp_%s_%s", b.mcpName, b.tool.Name)
44}
38}
@@ -56,27 +50,21 @@ func (b *McpTool) Info() tools.ToolInfo {
56 }
50 }
57}
51}
58
52
59−func runTool(ctx context.Context, c MCPClient, toolName string, input string) (tools.ToolResponse, error) {
53+func runTool(ctx context.Context, name, 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
74 var args map[string]any
54 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 {
76 return tools.NewTextErrorResponse(fmt.Sprintf("error parsing parameters: %s", err)), nil
56 return tools.NewTextErrorResponse(fmt.Sprintf("error parsing parameters: %s", err)), nil
77 }
57 }
78− toolRequest.Params.Arguments = args
58+ c, ok := mcpClients.Get(name)
59+ if !ok {
60+ return tools.NewTextErrorResponse("mcp '" + name + "' not available"), nil
61+ }
79− result, err := c.CallTool(ctx, toolRequest)
62+ result, err := c.CallTool(ctx, mcp.CallToolRequest{
63+ Params: mcp.CallToolParams{
64+ Name: toolName,
65+ Arguments: args,
66+ },
67+ })
80 if err != nil {
68 if err != nil {
81 return tools.NewTextErrorResponse(err.Error()), nil
69 return tools.NewTextErrorResponse(err.Error()), nil
82 }
70 }
@@ -114,56 +102,34 @@ func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolRes
114 return tools.ToolResponse{}, permission.ErrorPermissionDenied
102 return tools.ToolResponse{}, permission.ErrorPermissionDenied
115 }
103 }
116
104
117− return runTool(ctx, b.client, b.tool.Name, params.Input)
105+ return runTool(ctx, b.mcpName, b.tool.Name, params.Input)
118}
106}
119
107
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 {
108+func getTools(ctx context.Context, name string, permissions permission.Service, c *client.Client, 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)
109+ result, err := c.ListTools(ctx, mcp.ListToolsRequest{})
147 if err != nil {
110 if err != nil {
148 slog.Error("error listing tools", "error", err)
111 slog.Error("error listing tools", "error", err)
112+ c.Close()
113+ mcpClients.Del(name)
149− return stdioTools
114+ return nil
150 }
115 }
116+ mcpTools := make([]tools.BaseTool, 0, len(result.Tools))
151− for _, t := range tools.Tools {
117+ for _, tool := range result.Tools {
152− stdioTools = append(stdioTools, NewMcpTool(name, c, t, permissions, m, workingDir))
118+ mcpTools = append(mcpTools, &McpTool{
119+ mcpName: name,
120+ tool: tool,
121+ permissions: permissions,
122+ workingDir: workingDir,
123+ })
153 }
124 }
154− return stdioTools
125+ return mcpTools
155}
126}
156
127
157−var (
128+// CloseMCPClients closes all MCP clients. This should be called during application shutdown.
158− mcpToolsOnce sync.Once
129+func CloseMCPClients() {
159− mcpTools []tools.BaseTool
130+ for c := range mcpClients.Seq() {
160−)
131+ _ = c.Close()
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− })
132+ }
166− return mcpTools
167}
133}
168
134
169func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *config.Config) []tools.BaseTool {
135func 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
177 wg.Add(1)
143 wg.Add(1)
178 go func(name string, m config.MCPConfig) {
144 go func(name string, m config.MCPConfig) {
179 defer wg.Done()
145 defer wg.Done()
180− switch m.Type {
146+ c, err := doGetClient(m)
181− case config.MCPStdio:
182− c, err := client.NewStdioMCPClient(
183− m.Command,
184− m.ResolvedEnv(),
185− m.Args...,
186− )
187− if err != nil {
147+ if err != nil {
188− slog.Error("error creating mcp client", "error", err)
148+ slog.Error("error creating mcp client", "error", err)
189− return
149+ return
190− }
150+ }
191−
151+ if err := doInitClient(ctx, name, c); err != nil {
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)
152+ slog.Error("error initializing mcp client", "error", err)
210− return
153+ return
211− }
212− result.Append(getTools(ctx, name, m, permissions, c, cfg.WorkingDir())...)
213 }
154 }
155+ result.Append(getTools(ctx, name, permissions, c, cfg.WorkingDir())...)
214 }(name, m)
156 }(name, m)
215 }
157 }
216 wg.Wait()
158 wg.Wait()
217 return slices.Collect(result.Seq())
159 return slices.Collect(result.Seq())
218}
160}
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+}