language_models: Use weak entities in subscribe/observes around the language model registry (#52312)

Ben Brandt and Lukas Wirth created 2 weeks ago

## Context

I was getting some leak detection failures in evals and tracked it down
to these entities getting passed into observe/subscribe callbacks and
causing cycles.

Release Notes:

- N/A

Co-authored-by: Lukas Wirth <[e-mail]>

Change summary

crates/language_model/src/model/cloud_model.rs |  6 +
crates/language_models/src/language_models.rs  | 62 +++++++++++--------
2 files changed, 40 insertions(+), 28 deletions(-)

Detailed changes

crates/language_model/src/model/cloud_model.rs

@@ -154,9 +154,11 @@ impl RefreshLlmTokenListener {
 
     fn new(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut Context<Self>) -> Self {
         client.add_message_to_client_handler({
-            let this = cx.entity();
+            let this = cx.weak_entity();
             move |message, cx| {
-                Self::handle_refresh_llm_token(this.clone(), message, cx);
+                if let Some(this) = this.upgrade() {
+                    Self::handle_refresh_llm_token(this, message, cx);
+                }
             }
         });
 

crates/language_models/src/language_models.rs

@@ -39,37 +39,43 @@ pub fn init(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) {
     // Subscribe to extension store events to track LLM extension installations
     if let Some(extension_store) = extension_host::ExtensionStore::try_global(cx) {
         cx.subscribe(&extension_store, {
-            let registry = registry.clone();
-            move |extension_store, event, cx| match event {
-                extension_host::Event::ExtensionInstalled(extension_id) => {
-                    if let Some(manifest) = extension_store
-                        .read(cx)
-                        .extension_manifest_for_id(extension_id)
-                    {
-                        if !manifest.language_model_providers.is_empty() {
-                            registry.update(cx, |registry, cx| {
-                                registry.extension_installed(extension_id.clone(), cx);
-                            });
+            let registry = registry.downgrade();
+            move |extension_store, event, cx| {
+                let Some(registry) = registry.upgrade() else {
+                    return;
+                };
+                match event {
+                    extension_host::Event::ExtensionInstalled(extension_id) => {
+                        if let Some(manifest) = extension_store
+                            .read(cx)
+                            .extension_manifest_for_id(extension_id)
+                        {
+                            if !manifest.language_model_providers.is_empty() {
+                                registry.update(cx, |registry, cx| {
+                                    registry.extension_installed(extension_id.clone(), cx);
+                                });
+                            }
                         }
                     }
-                }
-                extension_host::Event::ExtensionUninstalled(extension_id) => {
-                    registry.update(cx, |registry, cx| {
-                        registry.extension_uninstalled(extension_id, cx);
-                    });
-                }
-                extension_host::Event::ExtensionsUpdated => {
-                    let mut new_ids = HashSet::default();
-                    for (extension_id, entry) in extension_store.read(cx).installed_extensions() {
-                        if !entry.manifest.language_model_providers.is_empty() {
-                            new_ids.insert(extension_id.clone());
+                    extension_host::Event::ExtensionUninstalled(extension_id) => {
+                        registry.update(cx, |registry, cx| {
+                            registry.extension_uninstalled(extension_id, cx);
+                        });
+                    }
+                    extension_host::Event::ExtensionsUpdated => {
+                        let mut new_ids = HashSet::default();
+                        for (extension_id, entry) in extension_store.read(cx).installed_extensions()
+                        {
+                            if !entry.manifest.language_model_providers.is_empty() {
+                                new_ids.insert(extension_id.clone());
+                            }
                         }
+                        registry.update(cx, |registry, cx| {
+                            registry.sync_installed_llm_extensions(new_ids, cx);
+                        });
                     }
-                    registry.update(cx, |registry, cx| {
-                        registry.sync_installed_llm_extensions(new_ids, cx);
-                    });
+                    _ => {}
                 }
-                _ => {}
             }
         })
         .detach();
@@ -101,7 +107,11 @@ pub fn init(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) {
             cx,
         );
     });
+    let registry = registry.downgrade();
     cx.observe_global::<SettingsStore>(move |cx| {
+        let Some(registry) = registry.upgrade() else {
+            return;
+        };
         let openai_compatible_providers_new = AllLanguageModelSettings::get_global(cx)
             .openai_compatible
             .keys()