chore: fix cancellation

Kujtim Hoxha created 3 months ago

Change summary

internal/app/app.go         | 6 +++---
internal/llm/agent/agent.go | 7 +++++++
2 files changed, 10 insertions(+), 3 deletions(-)

Detailed changes

internal/app/app.go

@@ -279,6 +279,9 @@ func (app *App) Subscribe(program *tea.Program) {
 
 // Shutdown performs a clean shutdown of the application
 func (app *App) Shutdown() {
+	if app.CoderAgent != nil {
+		app.CoderAgent.CancelAll()
+	}
 	app.cancelFuncsMutex.Lock()
 	for _, cancel := range app.watcherCancelFuncs {
 		cancel()
@@ -298,9 +301,6 @@ func (app *App) Shutdown() {
 		}
 		cancel()
 	}
-	if app.CoderAgent != nil {
-		app.CoderAgent.CancelAll()
-	}
 
 	for _, cleanup := range app.cleanupFuncs {
 		if cleanup != nil {

internal/llm/agent/agent.go

@@ -801,6 +801,13 @@ func (a *agent) CancelAll() {
 		a.Cancel(key.(string)) // key is sessionID
 		return true
 	})
+	for {
+		if a.IsBusy() {
+			time.Sleep(200 * time.Millisecond)
+		} else {
+			break
+		}
+	}
 }
 
 func (a *agent) UpdateModel() error {