42a346c · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

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

42a346cCarlos Alexandro Becker1 year ago

internal/llm/agent/agent.go+4−1

compare with main

unifiedsplit

@@ -449,8 +449,8 @@ func (a *agent) createUserMessage(ctx context.Context, sessionID, content string

449

450func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msgHistory []message.Message) (message.Message, *message.Message, error) {

451 ctx = context.WithValue(ctx, tools.SessionIDContextKey, sessionID)

452− eventChan := a.provider.StreamResponse(ctx, msgHistory, slices.Collect(a.tools.Seq()))

453

452

453+ // Create the assistant message first so the spinner shows immediately

454 assistantMsg, err := a.messages.Create(ctx, sessionID, message.CreateMessageParams{

455 Role: message.Assistant,

456 Parts: []message.ContentPart{},

@@ -461,6 +461,9 @@ func (a *agent) streamAndHandleEvents(ctx context.Context, sessionID string, msg

461 return assistantMsg, nil, fmt.Errorf("failed to create assistant message: %w", err)

462 }

463

464+ // Now collect tools (which may block on MCP initialization)

465+ eventChan := a.provider.StreamResponse(ctx, msgHistory, slices.Collect(a.tools.Seq()))

466+

464 // Add the session and message ID into the context if needed by tools.

467 // Add the session and message ID into the context if needed by tools.

465 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

468 ctx = context.WithValue(ctx, tools.MessageIDContextKey, assistantMsg.ID)

466

469

internal/llm/agent/mcp-tools.go+2

@@ -274,6 +274,8 @@ func doGetMCPTools(ctx context.Context, permissions permission.Service, cfg *con

274 }

275 }()

276

277+ ctx, cancel := context.WithTimeout(ctx, 10*time.Second)

278+ defer cancel()

277 c, err := createMcpClient(m)

279 c, err := createMcpClient(m)

278 if err != nil {

280 if err != nil {

279 updateMCPState(name, MCPStateError, err, nil, 0)

281 updateMCPState(name, MCPStateError, err, nil, 0)