From fc3a67f046b2becdd7caad11df80966acb5024ed Mon Sep 17 00:00:00 2001
From: Michael Sloan <michael@zed.dev>
Date: Sun, 4 May 2025 22:32:22 +0200
Subject: [PATCH] Fix mild memory leak for failed cache tasks

---
 crates/language_models/src/provider/google.rs | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/crates/language_models/src/provider/google.rs b/crates/language_models/src/provider/google.rs
index 4d5b1180f2..bf559f1e4c 100644
--- a/crates/language_models/src/provider/google.rs
+++ b/crates/language_models/src/provider/google.rs
@@ -934,14 +934,14 @@ impl GoogleLanguageModel {
                     let cache = self.cache.clone();
                     let executor = cx.background_executor().clone();
                     let task = cx.background_spawn(async move {
-                        match create_cache_future.await.log_err()? {
+                        let result = match create_cache_future.await.log_err()? {
                             CreateCacheResponse::Created(created_cache) => {
                                 log::info!("created cache `{}`", created_cache.name);
                                 CacheEntry::new(
                                     created_cache.name,
                                     created_cache.expire_time.to_utc(),
                                     cache_key,
-                                    cache,
+                                    cache.clone(),
                                     executor,
                                 )
                             }
@@ -949,7 +949,11 @@ impl GoogleLanguageModel {
                                 cache.lock().models_not_supported.insert(model);
                                 None
                             }
+                        };
+                        if result.is_none() {
+                            cache.lock().task_map.remove(&cache_key);
                         }
+                        result
                     });
                     self.cache.lock().task_map.insert(cache_key, task.shared());
                 }