From e2bb9700c4e3b9b0cc283e97bcad1e9fd44d8249 Mon Sep 17 00:00:00 2001
From: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
Date: Mon, 14 Jul 2025 22:39:03 -0300
Subject: [PATCH] fix: timeout and context cancel (#180)

* fix: timeout

* fix: handle context cancel

* chore: smaller diff

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>

* fix: once

---------

Signed-off-by: Carlos Alexandro Becker <caarlos0@users.noreply.github.com>
---
 cmd/root.go                 | 20 ++------------------
 internal/llm/agent/agent.go | 11 +++++++----
 2 files changed, 9 insertions(+), 22 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 24cde93e21c6c5f96abd0dc6259270b5987d8a25..8e1f0839323c29489031ef48567e31f105612ef3 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -6,7 +6,6 @@ import (
 	"io"
 	"log/slog"
 	"os"
-	"sync"
 	"time"
 
 	tea "github.com/charmbracelet/bubbletea/v2"
@@ -86,23 +85,7 @@ to assist developers in writing, debugging, and understanding code directly from
 			slog.Error(fmt.Sprintf("Failed to create app instance: %v", err))
 			return err
 		}
-
-		// Set up shutdown handling that works for both normal exit and signal interruption
-		var shutdownOnce sync.Once
-		shutdown := func() {
-			shutdownOnce.Do(func() {
-				slog.Info("Shutting down application")
-				app.Shutdown()
-			})
-		}
-		defer shutdown()
-
-		// Handle context cancellation (from signals) in a goroutine
-		go func() {
-			<-ctx.Done()
-			slog.Info("Context cancelled, initiating shutdown")
-			shutdown()
-		}()
+		defer app.Shutdown()
 
 		// Initialize MCP tools early for both modes
 		initMCPTools(ctx, app, cfg)
@@ -125,6 +108,7 @@ to assist developers in writing, debugging, and understanding code directly from
 			tea.WithAltScreen(),
 			tea.WithKeyReleases(),
 			tea.WithUniformKeyLayout(),
+			tea.WithContext(ctx),
 			tea.WithMouseCellMotion(),            // Use cell motion instead of all motion to reduce event flooding
 			tea.WithFilter(tui.MouseEventFilter), // Filter mouse events based on focus state
 		)
diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 6c7844eaa2811570824fba489a8c7a7581fa201f..56fb431b3b705a656cdfbf9df426b8ce8c7298c4 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -800,11 +800,14 @@ func (a *agent) CancelAll() {
 		a.Cancel(key.(string)) // key is sessionID
 		return true
 	})
-	for {
-		if a.IsBusy() {
+
+	timeout := time.After(5 * time.Second)
+	for a.IsBusy() {
+		select {
+		case <-timeout:
+			return
+		default:
 			time.Sleep(200 * time.Millisecond)
-		} else {
-			break
 		}
 	}
 }