From 07538ff08e04094ceb04e99221474cf5870f3c2c Mon Sep 17 00:00:00 2001
From: Ben Kunkle <ben@zed.dev>
Date: Wed, 17 Dec 2025 18:32:46 -0600
Subject: [PATCH] Make sweep and mercury API tokens use `cx.global` instead of
 `OnceLock` (#45176)

Closes #ISSUE

Release Notes:

- N/A *or* Added/Fixed/Improved ...
---
 crates/edit_prediction/src/mercury.rs  | 19 ++++++++++++-------
 crates/edit_prediction/src/sweep_ai.rs | 19 ++++++++++++-------
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/crates/edit_prediction/src/mercury.rs b/crates/edit_prediction/src/mercury.rs
index b47bd2ad0374eba33e7b8db726c2fa13c0519465..8186fc5d8c609468be04c117eabac11c6c015efd 100644
--- a/crates/edit_prediction/src/mercury.rs
+++ b/crates/edit_prediction/src/mercury.rs
@@ -6,7 +6,7 @@ use crate::{
 use anyhow::{Context as _, Result};
 use futures::AsyncReadExt as _;
 use gpui::{
-    App, AppContext as _, Entity, SharedString, Task,
+    App, AppContext as _, Entity, Global, SharedString, Task,
     http_client::{self, AsyncBody, Method},
 };
 use language::{OffsetRangeExt as _, ToOffset, ToPoint as _};
@@ -300,14 +300,19 @@ pub const MERCURY_CREDENTIALS_URL: SharedString =
     SharedString::new_static("https://api.inceptionlabs.ai/v1/edit/completions");
 pub const MERCURY_CREDENTIALS_USERNAME: &str = "mercury-api-token";
 pub static MERCURY_TOKEN_ENV_VAR: std::sync::LazyLock<EnvVar> = env_var!("MERCURY_AI_TOKEN");
-pub static MERCURY_API_KEY: std::sync::OnceLock<Entity<ApiKeyState>> = std::sync::OnceLock::new();
+
+struct GlobalMercuryApiKey(Entity<ApiKeyState>);
+
+impl Global for GlobalMercuryApiKey {}
 
 pub fn mercury_api_token(cx: &mut App) -> Entity<ApiKeyState> {
-    MERCURY_API_KEY
-        .get_or_init(|| {
-            cx.new(|_| ApiKeyState::new(MERCURY_CREDENTIALS_URL, MERCURY_TOKEN_ENV_VAR.clone()))
-        })
-        .clone()
+    if let Some(global) = cx.try_global::<GlobalMercuryApiKey>() {
+        return global.0.clone();
+    }
+    let entity =
+        cx.new(|_| ApiKeyState::new(MERCURY_CREDENTIALS_URL, MERCURY_TOKEN_ENV_VAR.clone()));
+    cx.set_global(GlobalMercuryApiKey(entity.clone()));
+    entity
 }
 
 pub fn load_mercury_api_token(cx: &mut App) -> Task<Result<(), language_model::AuthenticateError>> {
diff --git a/crates/edit_prediction/src/sweep_ai.rs b/crates/edit_prediction/src/sweep_ai.rs
index 2ed24cd8ef728383ec800acbb2ab7c7b99f07c06..71f28c9213c3440a9267dab7d5a5416dc219f2f3 100644
--- a/crates/edit_prediction/src/sweep_ai.rs
+++ b/crates/edit_prediction/src/sweep_ai.rs
@@ -1,7 +1,7 @@
 use anyhow::Result;
 use futures::AsyncReadExt as _;
 use gpui::{
-    App, AppContext as _, Entity, SharedString, Task,
+    App, AppContext as _, Entity, Global, SharedString, Task,
     http_client::{self, AsyncBody, Method},
 };
 use language::{Point, ToOffset as _};
@@ -272,14 +272,19 @@ pub const SWEEP_CREDENTIALS_URL: SharedString =
     SharedString::new_static("https://autocomplete.sweep.dev");
 pub const SWEEP_CREDENTIALS_USERNAME: &str = "sweep-api-token";
 pub static SWEEP_AI_TOKEN_ENV_VAR: std::sync::LazyLock<EnvVar> = env_var!("SWEEP_AI_TOKEN");
-pub static SWEEP_API_KEY: std::sync::OnceLock<Entity<ApiKeyState>> = std::sync::OnceLock::new();
+
+struct GlobalSweepApiKey(Entity<ApiKeyState>);
+
+impl Global for GlobalSweepApiKey {}
 
 pub fn sweep_api_token(cx: &mut App) -> Entity<ApiKeyState> {
-    SWEEP_API_KEY
-        .get_or_init(|| {
-            cx.new(|_| ApiKeyState::new(SWEEP_CREDENTIALS_URL, SWEEP_AI_TOKEN_ENV_VAR.clone()))
-        })
-        .clone()
+    if let Some(global) = cx.try_global::<GlobalSweepApiKey>() {
+        return global.0.clone();
+    }
+    let entity =
+        cx.new(|_| ApiKeyState::new(SWEEP_CREDENTIALS_URL, SWEEP_AI_TOKEN_ENV_VAR.clone()));
+    cx.set_global(GlobalSweepApiKey(entity.clone()));
+    entity
 }
 
 pub fn load_sweep_api_token(cx: &mut App) -> Task<Result<(), language_model::AuthenticateError>> {