git.iko.soy/crush
on main help sign in
fix(mcp): improve timeout errors (#1108)
2fd3b43Carlos Alexandro Becker
internal/llm/agent/mcp-tools.go+19−6
compare with main
unifiedsplit
@@ -152,13 +152,14 @@ func getOrRenewClient(ctx context.Context, name string) (*client.Client, error)
152 m := config.Get().MCP[name]
152 m := config.Get().MCP[name]
153 state, _ := mcpStates.Get(name)
153 state, _ := mcpStates.Get(name)
154
154
155+ timeout := mcpTimeout(m)
155− pingCtx, cancel := context.WithTimeout(ctx, mcpTimeout(m))
156+ pingCtx, cancel := context.WithTimeout(ctx, timeout)
156 defer cancel()
157 defer cancel()
157 err := c.Ping(pingCtx)
158 err := c.Ping(pingCtx)
158 if err == nil {
159 if err == nil {
159 return c, nil
160 return c, nil
160 }
161 }
161− updateMCPState(name, MCPStateError, err, nil, state.ToolCount)
162+ updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, state.ToolCount)
162
163
163 c, err = createAndInitializeClient(ctx, name, m)
164 c, err = createAndInitializeClient(ctx, name, m)
164 if err != nil {
165 if err != nil {
@@ -334,17 +335,22 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon
334 slog.Error("error creating mcp client", "error", err, "name", name)
335 slog.Error("error creating mcp client", "error", err, "name", name)
335 return nil, err
336 return nil, err
336 }
337 }
338+
339+ timeout := mcpTimeout(m)
340+ initCtx, cancel := context.WithTimeout(ctx, timeout)
341+ defer cancel()
342+
337 // Only call Start() for non-stdio clients, as stdio clients auto-start
343 // Only call Start() for non-stdio clients, as stdio clients auto-start
338 if m.Type != config.MCPStdio {
344 if m.Type != config.MCPStdio {
339− if err := c.Start(ctx); err != nil {
345+ if err := c.Start(initCtx); err != nil {
340− updateMCPState(name, MCPStateError, err, nil, 0)
346+ updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
341 slog.Error("error starting mcp client", "error", err, "name", name)
347 slog.Error("error starting mcp client", "error", err, "name", name)
342 _ = c.Close()
348 _ = c.Close()
343 return nil, err
349 return nil, err
344 }
350 }
345 }
351 }
346− if _, err := c.Initialize(ctx, mcpInitRequest); err != nil {
352+ if _, err := c.Initialize(initCtx, mcpInitRequest); err != nil {
347− updateMCPState(name, MCPStateError, err, nil, 0)
353+ updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
348 slog.Error("error initializing mcp client", "error", err, "name", name)
354 slog.Error("error initializing mcp client", "error", err, "name", name)
349 _ = c.Close()
355 _ = c.Close()
350 return nil, err
356 return nil, err
@@ -354,6 +360,13 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon
354 return c, nil
360 return c, nil
355}
361}
356
362
363+func maybeTimeoutErr(err error, timeout time.Duration) error {
364+ if errors.Is(err, context.DeadlineExceeded) {
365+ return fmt.Errorf("timed out after %s", timeout)
366+ }
367+ return err
368+}
369+
357func createMcpClient(name string, m config.MCPConfig) (*client.Client, error) {
370func createMcpClient(name string, m config.MCPConfig) (*client.Client, error) {
358 switch m.Type {
371 switch m.Type {
359 case config.MCPStdio:
372 case config.MCPStdio: