Fix Codestral API key credentials URL mismatch (#48513)

Closes #46506

Release Notes:

- Fixed an issue where the codestral URL used for credentials would be
different than the one used for requests causing authentication errors

Cargo.lock

@@ -3300,9 +3300,8 @@ dependencies = [
  "http_client",
  "icons",
  "language",
- "language_models",
+ "language_model",
  "log",
- "mistral",
  "serde",
  "serde_json",
  "text",
@@ -15215,6 +15214,7 @@ dependencies = [
  "assets",
  "bm25",
  "client",
+ "codestral",
  "component",
  "copilot",
  "copilot_ui",
@@ -15228,7 +15228,6 @@ dependencies = [
  "heck 0.5.0",
  "itertools 0.14.0",
  "language",
- "language_models",
  "log",
  "menu",
  "node_runtime",

crates/codestral/Cargo.toml

@@ -17,9 +17,8 @@ gpui.workspace = true
 http_client.workspace = true
 icons.workspace = true
 language.workspace = true
-language_models.workspace = true
+language_model.workspace = true
 log.workspace = true
-mistral.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 text.workspace = true

crates/codestral/src/codestral.rs

@@ -4,15 +4,15 @@ use edit_prediction_types::{
     EditPrediction, EditPredictionDelegate, EditPredictionDiscardReason, EditPredictionIconSet,
 };
 use futures::AsyncReadExt;
-use gpui::{App, Context, Entity, Task};
+use gpui::{App, AppContext as _, Context, Entity, Global, SharedString, Task};
 use http_client::HttpClient;
 use icons::IconName;
 use language::{
     Anchor, Buffer, BufferSnapshot, EditPreview, ToPoint, language_settings::all_language_settings,
 };
-use language_models::MistralLanguageModelProvider;
-use mistral::CODESTRAL_API_URL;
+use language_model::{ApiKeyState, AuthenticateError, EnvVar, env_var};
 use serde::{Deserialize, Serialize};
+
 use std::{
     ops::Range,
     sync::Arc,
@@ -20,8 +20,50 @@ use std::{
 };
 use text::{OffsetRangeExt as _, ToOffset};
 
+pub const CODESTRAL_API_URL: &str = "https://codestral.mistral.ai";
 pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(150);
 
+static CODESTRAL_API_KEY_ENV_VAR: std::sync::LazyLock<EnvVar> = env_var!("CODESTRAL_API_KEY");
+
+struct GlobalCodestralApiKey(Entity<ApiKeyState>);
+
+impl Global for GlobalCodestralApiKey {}
+
+pub fn codestral_api_key_state(cx: &mut App) -> Entity<ApiKeyState> {
+    if let Some(global) = cx.try_global::<GlobalCodestralApiKey>() {
+        return global.0.clone();
+    }
+    let entity =
+        cx.new(|cx| ApiKeyState::new(codestral_api_url(cx), CODESTRAL_API_KEY_ENV_VAR.clone()));
+    cx.set_global(GlobalCodestralApiKey(entity.clone()));
+    entity
+}
+
+pub fn codestral_api_key(cx: &App) -> Option<Arc<str>> {
+    let url = codestral_api_url(cx);
+    cx.try_global::<GlobalCodestralApiKey>()?
+        .0
+        .read(cx)
+        .key(&url)
+}
+
+pub fn load_codestral_api_key(cx: &mut App) -> Task<Result<(), AuthenticateError>> {
+    let api_url = codestral_api_url(cx);
+    codestral_api_key_state(cx).update(cx, |key_state, cx| {
+        key_state.load_if_needed(api_url, |s| s, cx)
+    })
+}
+
+pub fn codestral_api_url(cx: &App) -> SharedString {
+    all_language_settings(None, cx)
+        .edit_predictions
+        .codestral
+        .api_url
+        .clone()
+        .unwrap_or_else(|| CODESTRAL_API_URL.to_string())
+        .into()
+}
+
 /// Represents a completion that has been received and processed from Codestral.
 /// This struct maintains the state needed to interpolate the completion as the user types.
 #[derive(Clone)]
@@ -59,21 +101,8 @@ impl CodestralEditPredictionDelegate {
         }
     }
 
-    pub fn has_api_key(cx: &App) -> bool {
-        Self::api_key(cx).is_some()
-    }
-
-    /// This is so we can immediately show Codestral as a provider users can
-    /// switch to in the edit prediction menu, if the API has been added
-    pub fn ensure_api_key_loaded(http_client: Arc<dyn HttpClient>, cx: &mut App) {
-        MistralLanguageModelProvider::global(http_client, cx)
-            .load_codestral_api_key(cx)
-            .detach();
-    }
-
-    fn api_key(cx: &App) -> Option<Arc<str>> {
-        MistralLanguageModelProvider::try_global(cx)
-            .and_then(|provider| provider.codestral_api_key(CODESTRAL_API_URL, cx))
+    pub fn ensure_api_key_loaded(cx: &mut App) {
+        load_codestral_api_key(cx).detach();
     }
 
     /// Uses Codestral's Fill-in-the-Middle API for code completion.
@@ -180,7 +209,7 @@ impl EditPredictionDelegate for CodestralEditPredictionDelegate {
     }
 
     fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
-        Self::api_key(cx).is_some()
+        codestral_api_key(cx).is_some()
     }
 
     fn is_refreshing(&self, _cx: &App) -> bool {
@@ -196,7 +225,7 @@ impl EditPredictionDelegate for CodestralEditPredictionDelegate {
     ) {
         log::debug!("Codestral: Refresh called (debounce: {})", debounce);
 
-        let Some(api_key) = Self::api_key(cx) else {
+        let Some(api_key) = codestral_api_key(cx) else {
             log::warn!("Codestral: No API key configured, skipping refresh");
             return;
         };
@@ -221,12 +250,7 @@ impl EditPredictionDelegate for CodestralEditPredictionDelegate {
             .clone()
             .unwrap_or_else(|| "codestral-latest".to_string());
         let max_tokens = settings.edit_predictions.codestral.max_tokens;
-        let api_url = settings
-            .edit_predictions
-            .codestral
-            .api_url
-            .clone()
-            .unwrap_or_else(|| CODESTRAL_API_URL.to_string());
+        let api_url = codestral_api_url(cx).to_string();
 
         self.pending_request = Some(cx.spawn(async move |this, cx| {
             if debounce {

crates/edit_prediction_ui/src/edit_prediction_button.rs

@@ -1,7 +1,7 @@
 use anyhow::Result;
 use client::{Client, UserStore, zed_urls};
 use cloud_llm_client::UsageLimit;
-use codestral::CodestralEditPredictionDelegate;
+use codestral::{self, CodestralEditPredictionDelegate};
 use copilot::Status;
 use edit_prediction::{EditPredictionStore, Zeta2FeatureFlag};
 use edit_prediction_types::EditPredictionDelegateHandle;
@@ -287,7 +287,7 @@ impl Render for EditPredictionButton {
 
             EditPredictionProvider::Codestral => {
                 let enabled = self.editor_enabled.unwrap_or(true);
-                let has_api_key = CodestralEditPredictionDelegate::has_api_key(cx);
+                let has_api_key = codestral::codestral_api_key(cx).is_some();
                 let this = cx.weak_entity();
                 let file = self.file.clone();
                 let language = self.language.clone();
@@ -600,7 +600,6 @@ impl EditPredictionButton {
         fs: Arc<dyn Fs>,
         user_store: Entity<UserStore>,
         popover_menu_handle: PopoverMenuHandle<ContextMenu>,
-        client: Arc<Client>,
         project: Entity<Project>,
         cx: &mut Context<Self>,
     ) -> Self {
@@ -630,7 +629,7 @@ impl EditPredictionButton {
         })
         .detach();
 
-        CodestralEditPredictionDelegate::ensure_api_key_loaded(client.http_client(), cx);
+        CodestralEditPredictionDelegate::ensure_api_key_loaded(cx);
 
         Self {
             editor_subscription: None,
@@ -1493,7 +1492,7 @@ pub fn get_available_providers(cx: &mut App) -> Vec<EditPredictionProvider> {
         }
     }
 
-    if CodestralEditPredictionDelegate::has_api_key(cx) {
+    if codestral::codestral_api_key(cx).is_some() {
         providers.push(EditPredictionProvider::Codestral);
     }
 

crates/language_models/src/provider/mistral.rs

@@ -11,7 +11,7 @@ use language_model::{
     LanguageModelRequest, LanguageModelToolChoice, LanguageModelToolResultContent,
     LanguageModelToolUse, MessageContent, RateLimiter, Role, StopReason, TokenUsage, env_var,
 };
-pub use mistral::{CODESTRAL_API_URL, MISTRAL_API_URL, StreamResponse};
+pub use mistral::{MISTRAL_API_URL, StreamResponse};
 pub use settings::MistralAvailableModel as AvailableModel;
 use settings::{Settings, SettingsStore};
 use std::collections::HashMap;
@@ -29,9 +29,6 @@ const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new(
 const API_KEY_ENV_VAR_NAME: &str = "MISTRAL_API_KEY";
 static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
 
-const CODESTRAL_API_KEY_ENV_VAR_NAME: &str = "CODESTRAL_API_KEY";
-static CODESTRAL_API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(CODESTRAL_API_KEY_ENV_VAR_NAME);
-
 #[derive(Default, Clone, Debug, PartialEq)]
 pub struct MistralSettings {
     pub api_url: String,
@@ -45,20 +42,6 @@ pub struct MistralLanguageModelProvider {
 
 pub struct State {
     api_key_state: ApiKeyState,
-    codestral_api_key_state: Entity<ApiKeyState>,
-}
-
-pub fn codestral_api_key(cx: &mut App) -> Entity<ApiKeyState> {
-    // IMPORTANT:
-    // Do not store `Entity<T>` handles in process-wide statics (e.g. `OnceLock`).
-    //
-    // `Entity<T>` is tied to a particular `App`/entity-map context. Caching it globally can
-    // cause panics like "used a entity with the wrong context" when tests (or multiple apps)
-    // create distinct `App` instances in the same process.
-    //
-    // If we want a per-process singleton, store plain data (e.g. env var names) and create
-    // the entity per-App instead.
-    cx.new(|_| ApiKeyState::new(CODESTRAL_API_URL.into(), CODESTRAL_API_KEY_ENV_VAR.clone()))
 }
 
 impl State {
@@ -77,15 +60,6 @@ impl State {
         self.api_key_state
             .load_if_needed(api_url, |this| &mut this.api_key_state, cx)
     }
-
-    fn authenticate_codestral(
-        &mut self,
-        cx: &mut Context<Self>,
-    ) -> Task<Result<(), AuthenticateError>> {
-        self.codestral_api_key_state.update(cx, |state, cx| {
-            state.load_if_needed(CODESTRAL_API_URL.into(), |state| state, cx)
-        })
-    }
 }
 
 struct GlobalMistralLanguageModelProvider(Arc<MistralLanguageModelProvider>);
@@ -112,7 +86,6 @@ impl MistralLanguageModelProvider {
             .detach();
             State {
                 api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
-                codestral_api_key_state: codestral_api_key(cx),
             }
         });
 
@@ -121,19 +94,6 @@ impl MistralLanguageModelProvider {
         cx.global::<GlobalMistralLanguageModelProvider>().0.clone()
     }
 
-    pub fn load_codestral_api_key(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
-        self.state
-            .update(cx, |state, cx| state.authenticate_codestral(cx))
-    }
-
-    pub fn codestral_api_key(&self, url: &str, cx: &App) -> Option<Arc<str>> {
-        self.state
-            .read(cx)
-            .codestral_api_key_state
-            .read(cx)
-            .key(url)
-    }
-
     fn create_language_model(&self, model: mistral::Model) -> Arc<dyn LanguageModel> {
         Arc::new(MistralLanguageModel {
             id: LanguageModelId::from(model.id().to_string()),

crates/mistral/src/mistral.rs

@@ -7,7 +7,6 @@ use std::convert::TryFrom;
 use strum::EnumIter;
 
 pub const MISTRAL_API_URL: &str = "https://api.mistral.ai/v1";
-pub const CODESTRAL_API_URL: &str = "https://codestral.mistral.ai";
 
 #[derive(Clone, Copy, Serialize, Deserialize, Debug, Eq, PartialEq)]
 #[serde(rename_all = "lowercase")]

crates/settings_ui/Cargo.toml

@@ -21,6 +21,7 @@ agent_settings.workspace = true
 anyhow.workspace = true
 bm25 = "2.3.2"
 component.workspace = true
+codestral.workspace = true
 copilot.workspace = true
 copilot_ui.workspace = true
 edit_prediction.workspace = true
@@ -32,7 +33,6 @@ fuzzy.workspace = true
 gpui.workspace = true
 heck.workspace = true
 itertools.workspace = true
-language_models.workspace = true
 language.workspace = true
 log.workspace = true
 menu.workspace = true

crates/settings_ui/src/pages/edit_prediction_provider_setup.rs

@@ -1,3 +1,4 @@
+use codestral::{CODESTRAL_API_URL, codestral_api_key_state, codestral_api_url};
 use edit_prediction::{
     ApiKeyState,
     mercury::{MERCURY_CREDENTIALS_URL, mercury_api_token},
@@ -6,7 +7,7 @@ use edit_prediction::{
 use edit_prediction_ui::{get_available_providers, set_completion_provider};
 use gpui::{Entity, ScrollHandle, prelude::*};
 use language::language_settings::AllLanguageSettings;
-use language_models::provider::mistral::{CODESTRAL_API_URL, codestral_api_key};
+
 use settings::Settings as _;
 use ui::{ButtonLink, ConfiguredApiCard, ContextMenu, DropdownMenu, DropdownStyle, prelude::*};
 use workspace::AppState;
@@ -69,8 +70,8 @@ pub(crate) fn render_edit_prediction_setup_page(
                 IconName::AiMistral,
                 "Codestral",
                 "https://console.mistral.ai/codestral".into(),
-                codestral_api_key(cx),
-                |cx| language_models::MistralLanguageModelProvider::api_url(cx),
+                codestral_api_key_state(cx),
+                |cx| codestral_api_url(cx),
                 Some(
                     settings_window
                         .render_sub_page_items_section(
@@ -175,7 +176,7 @@ fn render_api_key_provider(
     cx: &mut Context<SettingsWindow>,
 ) -> impl IntoElement {
     let weak_page = cx.weak_entity();
-    _ = window.use_keyed_state(title, cx, |_, cx| {
+    _ = window.use_keyed_state(current_url(cx), cx, |_, cx| {
         let task = api_key_state.update(cx, |key_state, cx| {
             key_state.load_if_needed(current_url(cx), |state| state, cx)
         });

crates/zed/src/zed.rs

@@ -421,7 +421,6 @@ pub fn initialize_workspace(
                 app_state.fs.clone(),
                 app_state.user_store.clone(),
                 edit_prediction_menu_handle.clone(),
-                app_state.client.clone(),
                 workspace.project().clone(),
                 cx,
             )

crates/zed/src/zed/edit_prediction_registry.rs

@@ -1,5 +1,5 @@
 use client::{Client, UserStore};
-use codestral::CodestralEditPredictionDelegate;
+use codestral::{CodestralEditPredictionDelegate, load_codestral_api_key};
 use collections::HashMap;
 use copilot::CopilotEditPredictionDelegate;
 use edit_prediction::{ZedEditPredictionDelegate, Zeta2FeatureFlag};
@@ -7,7 +7,7 @@ use editor::Editor;
 use feature_flags::FeatureFlagAppExt;
 use gpui::{AnyWindowHandle, App, AppContext as _, Context, Entity, WeakEntity};
 use language::language_settings::{EditPredictionProvider, all_language_settings};
-use language_models::MistralLanguageModelProvider;
+
 use settings::{EXPERIMENTAL_ZETA2_EDIT_PREDICTION_PROVIDER_NAME, SettingsStore};
 use std::{cell::RefCell, rc::Rc, sync::Arc};
 use supermaven::{Supermaven, SupermavenEditPredictionDelegate};
@@ -111,8 +111,7 @@ fn assign_edit_prediction_providers(
     cx: &mut App,
 ) {
     if provider == EditPredictionProvider::Codestral {
-        let mistral = MistralLanguageModelProvider::global(client.http_client(), cx);
-        mistral.load_codestral_api_key(cx).detach();
+        load_codestral_api_key(cx).detach();
     }
     for (editor, window) in editors.borrow().iter() {
         _ = window.update(cx, |_window, window, cx| {