git.iko.soy/crush
on main help sign in
feat(mcp): ping and recreate mcp client if needed (#772)
5af4d19Carlos Alexandro Becker
internal/llm/agent/mcp-tools.go+66−30
compare with main
unifiedsplit
@@ -6,6 +6,7 @@ import (
6 "encoding/json"
6 "encoding/json"
7 "fmt"
7 "fmt"
8 "log/slog"
8 "log/slog"
9+ "maps"
9 "slices"
10 "slices"
10 "strings"
11 "strings"
11 "sync"
12 "sync"
@@ -110,9 +111,10 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
110 if err := json.Unmarshal([]byte(input), &args); err != nil {
111 if err := json.Unmarshal([]byte(input), &args); err != nil {
111 return tools.NewTextErrorResponse(fmt.Sprintf("error parsing parameters: %s", err)), nil
112 return tools.NewTextErrorResponse(fmt.Sprintf("error parsing parameters: %s", err)), nil
112 }
113 }
113− c, ok := mcpClients.Get(name)
114+
114− if !ok {
115+ c, err := getOrRenewClient(ctx, name)
116+ if err != nil {
115− return tools.NewTextErrorResponse("mcp '" + name + "' not available"), nil
117+ return tools.NewTextErrorResponse(err.Error()), nil
116 }
118 }
117 result, err := c.CallTool(ctx, mcp.CallToolRequest{
119 result, err := c.CallTool(ctx, mcp.CallToolRequest{
118 Params: mcp.CallToolParams{
120 Params: mcp.CallToolParams{
@@ -135,6 +137,33 @@ func runTool(ctx context.Context, name, toolName string, input string) (tools.To
135 return tools.NewTextResponse(strings.Join(output, "\n")), nil
137 return tools.NewTextResponse(strings.Join(output, "\n")), nil
136}
138}
137
139
140+func getOrRenewClient(ctx context.Context, name string) (*client.Client, error) {
141+ c, ok := mcpClients.Get(name)
142+ if !ok {
143+ return nil, fmt.Errorf("mcp '%s' not available", name)
144+ }
145+
146+ m := config.Get().MCP[name]
147+ state, _ := mcpStates.Get(name)
148+
149+ pingCtx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
150+ defer cancel()
151+ err := c.Ping(pingCtx)
152+ if err == nil {
153+ return c, nil
154+ }
155+ updateMCPState(name, MCPStateError, err, nil, state.ToolCount)
156+
157+ c, err = createAndInitializeClient(ctx, name, m)
158+ if err != nil {
159+ return nil, err
160+ }
161+
162+ updateMCPState(name, MCPStateConnected, nil, c, state.ToolCount)
163+ mcpClients.Set(name, c)
164+ return c, nil
165+}
166+
138func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolResponse, error) {
167func (b *McpTool) Run(ctx context.Context, params tools.ToolCall) (tools.ToolResponse, error) {
139 sessionID, messageID := tools.GetContextValues(ctx)
168 sessionID, messageID := tools.GetContextValues(ctx)
140 if sessionID == "" || messageID == "" {
169 if sessionID == "" || messageID == "" {
@@ -187,11 +216,7 @@ func SubscribeMCPEvents(ctx context.Context) <-chan pubsub.Event[MCPEvent] {
187
216
188// GetMCPStates returns the current state of all MCP clients
217// GetMCPStates returns the current state of all MCP clients
189func GetMCPStates() map[string]MCPClientInfo {
218func GetMCPStates() map[string]MCPClientInfo {
190− states := make(map[string]MCPClientInfo)
191− for name, info := range mcpStates.Seq2() {
219+ return maps.Collect(mcpStates.Seq2())
192− states[name] = info
193− }
194− return states
195}
220}
196
221
197// GetMCPState returns the state of a specific MCP client
222// GetMCPState returns the state of a specific MCP client
@@ -275,32 +300,12 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con
275 }
300 }
276 }()
301 }()
277
302
278− timeout := time.Duration(cmp.Or(m.Timeout, 15)) * time.Second
279− ctx, cancel := context.WithTimeout(ctx, timeout)
303+ ctx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
280 defer cancel()
304 defer cancel()
281− c, err := createMcpClient(m)
305+ c, err := createAndInitializeClient(ctx, name, m)
282 if err != nil {
306 if err != nil {
283− updateMCPState(name, MCPStateError, err, nil, 0)
284− slog.Error("error creating mcp client", "error", err, "name", name)
285 return
307 return
286 }
308 }
287− // Only call Start() for non-stdio clients, as stdio clients auto-start
288− if m.Type != config.MCPStdio {
289− if err := c.Start(ctx); err != nil {
290− updateMCPState(name, MCPStateError, err, nil, 0)
291− slog.Error("error starting mcp client", "error", err, "name", name)
292− _ = c.Close()
293− return
294− }
295− }
296− if _, err := c.Initialize(ctx, mcpInitRequest); err != nil {
297− updateMCPState(name, MCPStateError, err, nil, 0)
298− slog.Error("error initializing mcp client", "error", err, "name", name)
299− _ = c.Close()
300− return
301− }
302−
303− slog.Info("Initialized mcp client", "name", name)
304 mcpClients.Set(name, c)
309 mcpClients.Set(name, c)
305
310
306 tools := getTools(ctx, name, permissions, c, cfg.WorkingDir())
311 tools := getTools(ctx, name, permissions, c, cfg.WorkingDir())
@@ -312,6 +317,33 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con
312 return slices.Collect(result.Seq())
317 return slices.Collect(result.Seq())
313}
318}
314
319
320+func createAndInitializeClient(ctx context.Context, name string, m config.MCPConfig) (*client.Client, error) {
321+ c, err := createMcpClient(m)
322+ if err != nil {
323+ updateMCPState(name, MCPStateError, err, nil, 0)
324+ slog.Error("error creating mcp client", "error", err, "name", name)
325+ return nil, err
326+ }
327+ // Only call Start() for non-stdio clients, as stdio clients auto-start
328+ if m.Type != config.MCPStdio {
329+ if err := c.Start(ctx); err != nil {
330+ updateMCPState(name, MCPStateError, err, nil, 0)
331+ slog.Error("error starting mcp client", "error", err, "name", name)
332+ _ = c.Close()
333+ return nil, err
334+ }
335+ }
336+ if _, err := c.Initialize(ctx, mcpInitRequest); err != nil {
337+ updateMCPState(name, MCPStateError, err, nil, 0)
338+ slog.Error("error initializing mcp client", "error", err, "name", name)
339+ _ = c.Close()
340+ return nil, err
341+ }
342+
343+ slog.Info("Initialized mcp client", "name", name)
344+ return c, nil
345+}
346+
315func createMcpClient(m config.MCPConfig) (*client.Client, error) {
347func createMcpClient(m config.MCPConfig) (*client.Client, error) {
316 switch m.Type {
348 switch m.Type {
317 case config.MCPStdio:
349 case config.MCPStdio:
@@ -343,3 +375,7 @@ type mcpLogger struct{}
343
375
344func (l mcpLogger) Errorf(format string, v ...any) { slog.Error(fmt.Sprintf(format, v...)) }
376func (l mcpLogger) Errorf(format string, v ...any) { slog.Error(fmt.Sprintf(format, v...)) }
345func (l mcpLogger) Infof(format string, v ...any) { slog.Info(fmt.Sprintf(format, v...)) }
377func (l mcpLogger) Infof(format string, v ...any) { slog.Info(fmt.Sprintf(format, v...)) }
378+
379+func mcpTimeout(m config.MCPConfig) time.Duration {
380+ return time.Duration(cmp.Or(m.Timeout, 15)) * time.Second
381+}