From 42a346ceb034cd1ebe0a10b696ac7c796984521d Mon Sep 17 00:00:00 2001 From: Carlos Alexandro Becker Date: Fri, 8 Aug 2025 16:18:27 -0300 Subject: [PATCH] fix(mcp): give a time limit for startup (#660) * fix(mcp): give a time limit for startup If a MCP gets stuck starting/connecting, chatting will not work. This makes it so it'll timeout after 5s (we can discuss the time), and then everything works accordingly. Signed-off-by: Carlos Alexandro Becker * fix: spinner Signed-off-by: Carlos Alexandro Becker * Update mcp-tools.go Co-authored-by: Andrey Nering --------- Signed-off-by: Carlos Alexandro Becker Co-authored-by: Andrey Nering --- internal/llm/agent/agent.go | 5 ++++- internal/llm/agent/mcp-tools.go | 2 ++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go index 57c30f41078f9a577d2f7bdbca5303da0eb83724..1067a2fd9505d25e0d63988d606c9da5ef2ecabd 100644 --- a/internal/llm/agent/agent.go +++ b/internal/llm/agent/agent.go @@ -449,8 +449,8 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) { ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID) - eventChan := a.provider.StreamResponse(ctx, msgHistory, slices.Collect(a.tools.Seq())) + // Create the assistant message first so the spinner shows immediately assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{ Role: message.Assistant, Parts: []message.ContentPart{}, @@ -461,6 +461,9 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err) } + // Now collect tools (which may block on MCP initialization) + eventChan := a.provider.StreamResponse(ctx, msgHistory, slices.Collect(a.tools.Seq())) + // Add the session and message ID into the context if needed by tools. ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID) diff --git a/internal/llm/agent/mcp-tools.go b/internal/llm/agent/mcp-tools.go index ac4385d2725c600c01a8ac479c8364c99b0cb217..a4d68965c62fa6b9e328e6b0bb9fac790c574f02 100644 --- a/internal/llm/agent/mcp-tools.go +++ b/internal/llm/agent/mcp-tools.go @@ -274,6 +274,8 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con } }() + ctx, cancel := context.WithTimeout(ctx, 10*time.Second) + defer cancel() c, err := createMcpClient(m) if err != nil { updateMCPState(name, MCPStateError, err, nil, 0)