Merge remote-tracking branch 'origin/main' into feat/mcp_notifications

Carlos Alexandro Becker created 1 month ago

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

Change summary

internal/llm/agent/mcp-tools.go                  | 15 ++++++++++-----
internal/tui/components/dialogs/models/keys.go   |  4 ++--
internal/tui/components/dialogs/sessions/keys.go |  4 ++--
3 files changed, 14 insertions(+), 9 deletions(-)

Detailed changes

internal/llm/agent/mcp-tools.go

@@ -377,30 +377,35 @@ func createAndInitializeClient(ctx context.Context, name string, m config.MCPCon
 		}
 	})
 
+	// XXX: ideally we should be able to use context.WithTimeout here, but,
+	// the SSE MCP client will start failing once that context is canceled.
 	timeout := mcpTimeout(m)
-	initCtx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
+	mcpCtx, cancel := context.WithCancel(ctx)
+	cancelTimer := time.AfterFunc(timeout, cancel)
 
-	if err := c.Start(initCtx); err != nil {
+	if err := c.Start(mcpCtx); err != nil {
 		updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
 		slog.Error("error starting mcp client", "error", err, "name", name)
 		_ = c.Close()
+		cancel()
 		return nil, err
 	}
 
-	if _, err := c.Initialize(initCtx, mcpInitRequest); err != nil {
+	if _, err := c.Initialize(mcpCtx, mcpInitRequest); err != nil {
 		updateMCPState(name, MCPStateError, maybeTimeoutErr(err, timeout), nil, 0)
 		slog.Error("error initializing mcp client", "error", err, "name", name)
 		_ = c.Close()
+		cancel()
 		return nil, err
 	}
 
+	cancelTimer.Stop()
 	slog.Info("Initialized mcp client", "name", name)
 	return c, nil
 }
 
 func maybeTimeoutErr(err error, timeout time.Duration) error {
-	if errors.Is(err, context.DeadlineExceeded) {
+	if errors.Is(err, context.Canceled) {
 		return fmt.Errorf("timed out after %s", timeout)
 	}
 	return err

internal/tui/components/dialogs/models/keys.go

@@ -19,7 +19,7 @@ func DefaultKeyMap() KeyMap {
 	return KeyMap{
 		Select: key.NewBinding(
 			key.WithKeys("enter", "ctrl+y"),
-			key.WithHelp("enter", "confirm"),
+			key.WithHelp("enter", "choose"),
 		),
 		Next: key.NewBinding(
 			key.WithKeys("down", "ctrl+n"),
@@ -35,7 +35,7 @@ func DefaultKeyMap() KeyMap {
 		),
 		Close: key.NewBinding(
 			key.WithKeys("esc", "alt+esc"),
-			key.WithHelp("esc", "cancel"),
+			key.WithHelp("esc", "exit"),
 		),
 	}
 }

internal/tui/components/dialogs/sessions/keys.go

@@ -15,7 +15,7 @@ func DefaultKeyMap() KeyMap {
 	return KeyMap{
 		Select: key.NewBinding(
 			key.WithKeys("enter", "tab", "ctrl+y"),
-			key.WithHelp("enter", "confirm"),
+			key.WithHelp("enter", "choose"),
 		),
 		Next: key.NewBinding(
 			key.WithKeys("down", "ctrl+n"),
@@ -27,7 +27,7 @@ func DefaultKeyMap() KeyMap {
 		),
 		Close: key.NewBinding(
 			key.WithKeys("esc", "alt+esc"),
-			key.WithHelp("esc", "cancel"),
+			key.WithHelp("esc", "exit"),
 		),
 	}
 }