From b7dcd4e4d64daaa8d9eaa9210666a15591eaff76 Mon Sep 17 00:00:00 2001
From: Thorsten Ball <mrnugget@gmail.com>
Date: Wed, 14 Aug 2024 12:13:09 +0200
Subject: [PATCH] assistant panel: Fix pending completions not being cleaned up
 (#16201)

Turns out that you could always cancel a completion, even if it was
already done and completed, because it was never cleaned up.

Release Notes:

- N/A
---
 crates/assistant/src/context.rs | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/crates/assistant/src/context.rs b/crates/assistant/src/context.rs
index ca57bd86dd..d0891faa66 100644
--- a/crates/assistant/src/context.rs
+++ b/crates/assistant/src/context.rs
@@ -1713,6 +1713,8 @@ impl Context {
             .insert_message_after(assistant_message.id, Role::User, MessageStatus::Done, cx)
             .unwrap();
 
+        let pending_completion_id = post_inc(&mut self.completion_count);
+
         let task = cx.spawn({
             |this, mut cx| async move {
                 let stream = model.stream_completion(request, &cx);
@@ -1761,10 +1763,9 @@ impl Context {
                         })?;
                         smol::future::yield_now().await;
                     }
-
                     this.update(&mut cx, |this, cx| {
                         this.pending_completions
-                            .retain(|completion| completion.id != this.completion_count);
+                            .retain(|completion| completion.id != pending_completion_id);
                         this.summarize(false, cx);
                     })?;
 
@@ -1806,7 +1807,7 @@ impl Context {
         });
 
         self.pending_completions.push(PendingCompletion {
-            id: post_inc(&mut self.completion_count),
+            id: pending_completion_id,
             assistant_message_id: assistant_message.id,
             _task: task,
         });