diff --git a/internal/llm/agent/agent.go b/internal/llm/agent/agent.go
index 0541ec7d42c35cc05c6dc2129f85225613a51daa..fb25d365634f6c0c544e9183b4e34d5f86907c03 100644
--- a/internal/llm/agent/agent.go
+++ b/internal/llm/agent/agent.go
@@ -355,10 +355,7 @@ func (a *agent) Run(ctx context.Context, sessionID string, content string, attac
 	}
 
 	genCtx, cancel := context.WithCancel(ctx)
-	defer cancel() // Ensure cancel is always called
-
 	a.activeRequests.Set(sessionID, cancel)
-	defer a.activeRequests.Del(sessionID) // Clean up on exit
 
 	go func() {
 		slog.Debug("Request started", "sessionID", sessionID)
@@ -374,6 +371,8 @@ func (a *agent) Run(ctx context.Context, sessionID string, content string, attac
 			slog.Error(result.Error.Error())
 		}
 		slog.Debug("Request completed", "sessionID", sessionID)
+		a.activeRequests.Del(sessionID)
+		cancel()
 		a.Publish(pubsub.CreatedEvent, result)
 		events <- result
 		close(events)