fix(mcp): give a time limit for startup (#660)

Carlos Alexandro Becker and Andrey Nering created 2 months ago

* 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 <[e-mail]>

* fix: spinner

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

* Update mcp-tools.go

Co-authored-by: Andrey Nering <[e-mail]>

---------

Signed-off-by: Carlos Alexandro Becker <[e-mail]>
Co-authored-by: Andrey Nering <[e-mail]>

Change summary

internal/llm/agent/agent.go     | 5 ++++-
internal/llm/agent/mcp-tools.go | 2 ++
2 files changed, 6 insertions(+), 1 deletion(-)

Detailed changes

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)
 

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)