cloud.rs

   1use ai_onboarding::YoungAccountBanner;
   2use anthropic::AnthropicModelMode;
   3use anyhow::{Context as _, Result, anyhow};
   4use chrono::{DateTime, Utc};
   5use client::{Client, ModelRequestUsage, UserStore, zed_urls};
   6use cloud_llm_client::{
   7    CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CLIENT_SUPPORTS_X_AI_HEADER_NAME,
   8    CURRENT_PLAN_HEADER_NAME, CompletionBody, CompletionEvent, CompletionRequestStatus,
   9    CountTokensBody, CountTokensResponse, EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse,
  10    MODEL_REQUESTS_RESOURCE_HEADER_VALUE, Plan, PlanV1, PlanV2,
  11    SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME,
  12    TOOL_USE_LIMIT_REACHED_HEADER_NAME, ZED_VERSION_HEADER_NAME,
  13};
  14use futures::{
  15    AsyncBufReadExt, FutureExt, Stream, StreamExt, future::BoxFuture, stream::BoxStream,
  16};
  17use google_ai::GoogleModelMode;
  18use gpui::{AnyElement, AnyView, App, AsyncApp, Context, Entity, Subscription, Task};
  19use http_client::http::{HeaderMap, HeaderValue};
  20use http_client::{AsyncBody, HttpClient, HttpRequestExt, Method, Response, StatusCode};
  21use language_model::{
  22    AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,
  23    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
  24    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
  25    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
  26    LanguageModelToolSchemaFormat, LlmApiToken, ModelRequestLimitReachedError,
  27    PaymentRequiredError, RateLimiter, RefreshLlmTokenListener,
  28};
  29use release_channel::AppVersion;
  30use schemars::JsonSchema;
  31use semver::Version;
  32use serde::{Deserialize, Serialize, de::DeserializeOwned};
  33use settings::SettingsStore;
  34pub use settings::ZedDotDevAvailableModel as AvailableModel;
  35pub use settings::ZedDotDevAvailableProvider as AvailableProvider;
  36use smol::io::{AsyncReadExt, BufReader};
  37use std::pin::Pin;
  38use std::str::FromStr as _;
  39use std::sync::Arc;
  40use std::time::Duration;
  41use thiserror::Error;
  42use ui::{TintColor, prelude::*};
  43use util::{ResultExt as _, maybe};
  44
  45use crate::provider::anthropic::{
  46    AnthropicEventMapper, count_anthropic_tokens_with_tiktoken, into_anthropic,
  47};
  48use crate::provider::google::{GoogleEventMapper, into_google};
  49use crate::provider::open_ai::{OpenAiEventMapper, count_open_ai_tokens, into_open_ai};
  50use crate::provider::x_ai::count_xai_tokens;
  51
  52const PROVIDER_ID: LanguageModelProviderId = language_model::ZED_CLOUD_PROVIDER_ID;
  53const PROVIDER_NAME: LanguageModelProviderName = language_model::ZED_CLOUD_PROVIDER_NAME;
  54
  55#[derive(Default, Clone, Debug, PartialEq)]
  56pub struct ZedDotDevSettings {
  57    pub available_models: Vec<AvailableModel>,
  58}
  59#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
  60#[serde(tag = "type", rename_all = "lowercase")]
  61pub enum ModelMode {
  62    #[default]
  63    Default,
  64    Thinking {
  65        /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
  66        budget_tokens: Option<u32>,
  67    },
  68}
  69
  70impl From<ModelMode> for AnthropicModelMode {
  71    fn from(value: ModelMode) -> Self {
  72        match value {
  73            ModelMode::Default => AnthropicModelMode::Default,
  74            ModelMode::Thinking { budget_tokens } => AnthropicModelMode::Thinking { budget_tokens },
  75        }
  76    }
  77}
  78
  79pub struct CloudLanguageModelProvider {
  80    client: Arc<Client>,
  81    state: Entity<State>,
  82    _maintain_client_status: Task<()>,
  83}
  84
  85pub struct State {
  86    client: Arc<Client>,
  87    llm_api_token: LlmApiToken,
  88    user_store: Entity<UserStore>,
  89    status: client::Status,
  90    models: Vec<Arc<cloud_llm_client::LanguageModel>>,
  91    default_model: Option<Arc<cloud_llm_client::LanguageModel>>,
  92    default_fast_model: Option<Arc<cloud_llm_client::LanguageModel>>,
  93    recommended_models: Vec<Arc<cloud_llm_client::LanguageModel>>,
  94    _fetch_models_task: Task<()>,
  95    _settings_subscription: Subscription,
  96    _llm_token_subscription: Subscription,
  97}
  98
  99impl State {
 100    fn new(
 101        client: Arc<Client>,
 102        user_store: Entity<UserStore>,
 103        status: client::Status,
 104        cx: &mut Context<Self>,
 105    ) -> Self {
 106        let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
 107        let mut current_user = user_store.read(cx).watch_current_user();
 108        Self {
 109            client: client.clone(),
 110            llm_api_token: LlmApiToken::default(),
 111            user_store,
 112            status,
 113            models: Vec::new(),
 114            default_model: None,
 115            default_fast_model: None,
 116            recommended_models: Vec::new(),
 117            _fetch_models_task: cx.spawn(async move |this, cx| {
 118                maybe!(async move {
 119                    let (client, llm_api_token) = this
 120                        .read_with(cx, |this, _cx| (client.clone(), this.llm_api_token.clone()))?;
 121
 122                    while current_user.borrow().is_none() {
 123                        current_user.next().await;
 124                    }
 125
 126                    let response =
 127                        Self::fetch_models(client.clone(), llm_api_token.clone()).await?;
 128                    this.update(cx, |this, cx| this.update_models(response, cx))?;
 129                    anyhow::Ok(())
 130                })
 131                .await
 132                .context("failed to fetch Zed models")
 133                .log_err();
 134            }),
 135            _settings_subscription: cx.observe_global::<SettingsStore>(|_, cx| {
 136                cx.notify();
 137            }),
 138            _llm_token_subscription: cx.subscribe(
 139                &refresh_llm_token_listener,
 140                move |this, _listener, _event, cx| {
 141                    let client = this.client.clone();
 142                    let llm_api_token = this.llm_api_token.clone();
 143                    cx.spawn(async move |this, cx| {
 144                        llm_api_token.refresh(&client).await?;
 145                        let response = Self::fetch_models(client, llm_api_token).await?;
 146                        this.update(cx, |this, cx| {
 147                            this.update_models(response, cx);
 148                        })
 149                    })
 150                    .detach_and_log_err(cx);
 151                },
 152            ),
 153        }
 154    }
 155
 156    fn is_signed_out(&self, cx: &App) -> bool {
 157        self.user_store.read(cx).current_user().is_none()
 158    }
 159
 160    fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
 161        let client = self.client.clone();
 162        cx.spawn(async move |state, cx| {
 163            client.sign_in_with_optional_connect(true, cx).await?;
 164            state.update(cx, |_, cx| cx.notify())
 165        })
 166    }
 167    fn update_models(&mut self, response: ListModelsResponse, cx: &mut Context<Self>) {
 168        let mut models = Vec::new();
 169
 170        for model in response.models {
 171            models.push(Arc::new(model.clone()));
 172
 173            // Right now we represent thinking variants of models as separate models on the client,
 174            // so we need to insert variants for any model that supports thinking.
 175            if model.supports_thinking {
 176                models.push(Arc::new(cloud_llm_client::LanguageModel {
 177                    id: cloud_llm_client::LanguageModelId(format!("{}-thinking", model.id).into()),
 178                    display_name: format!("{} Thinking", model.display_name),
 179                    ..model
 180                }));
 181            }
 182        }
 183
 184        self.default_model = models
 185            .iter()
 186            .find(|model| {
 187                response
 188                    .default_model
 189                    .as_ref()
 190                    .is_some_and(|default_model_id| &model.id == default_model_id)
 191            })
 192            .cloned();
 193        self.default_fast_model = models
 194            .iter()
 195            .find(|model| {
 196                response
 197                    .default_fast_model
 198                    .as_ref()
 199                    .is_some_and(|default_fast_model_id| &model.id == default_fast_model_id)
 200            })
 201            .cloned();
 202        self.recommended_models = response
 203            .recommended_models
 204            .iter()
 205            .filter_map(|id| models.iter().find(|model| &model.id == id))
 206            .cloned()
 207            .collect();
 208        self.models = models;
 209        cx.notify();
 210    }
 211
 212    async fn fetch_models(
 213        client: Arc<Client>,
 214        llm_api_token: LlmApiToken,
 215    ) -> Result<ListModelsResponse> {
 216        let http_client = &client.http_client();
 217        let token = llm_api_token.acquire(&client).await?;
 218
 219        let request = http_client::Request::builder()
 220            .method(Method::GET)
 221            .header(CLIENT_SUPPORTS_X_AI_HEADER_NAME, "true")
 222            .uri(http_client.build_zed_llm_url("/models", &[])?.as_ref())
 223            .header("Authorization", format!("Bearer {token}"))
 224            .body(AsyncBody::empty())?;
 225        let mut response = http_client
 226            .send(request)
 227            .await
 228            .context("failed to send list models request")?;
 229
 230        if response.status().is_success() {
 231            let mut body = String::new();
 232            response.body_mut().read_to_string(&mut body).await?;
 233            Ok(serde_json::from_str(&body)?)
 234        } else {
 235            let mut body = String::new();
 236            response.body_mut().read_to_string(&mut body).await?;
 237            anyhow::bail!(
 238                "error listing models.\nStatus: {:?}\nBody: {body}",
 239                response.status(),
 240            );
 241        }
 242    }
 243}
 244
 245impl CloudLanguageModelProvider {
 246    pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self {
 247        let mut status_rx = client.status();
 248        let status = *status_rx.borrow();
 249
 250        let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx));
 251
 252        let state_ref = state.downgrade();
 253        let maintain_client_status = cx.spawn(async move |cx| {
 254            while let Some(status) = status_rx.next().await {
 255                if let Some(this) = state_ref.upgrade() {
 256                    _ = this.update(cx, |this, cx| {
 257                        if this.status != status {
 258                            this.status = status;
 259                            cx.notify();
 260                        }
 261                    });
 262                } else {
 263                    break;
 264                }
 265            }
 266        });
 267
 268        Self {
 269            client,
 270            state,
 271            _maintain_client_status: maintain_client_status,
 272        }
 273    }
 274
 275    fn create_language_model(
 276        &self,
 277        model: Arc<cloud_llm_client::LanguageModel>,
 278        llm_api_token: LlmApiToken,
 279    ) -> Arc<dyn LanguageModel> {
 280        Arc::new(CloudLanguageModel {
 281            id: LanguageModelId(SharedString::from(model.id.0.clone())),
 282            model,
 283            llm_api_token,
 284            client: self.client.clone(),
 285            request_limiter: RateLimiter::new(4),
 286        })
 287    }
 288}
 289
 290impl LanguageModelProviderState for CloudLanguageModelProvider {
 291    type ObservableEntity = State;
 292
 293    fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
 294        Some(self.state.clone())
 295    }
 296}
 297
 298impl LanguageModelProvider for CloudLanguageModelProvider {
 299    fn id(&self) -> LanguageModelProviderId {
 300        PROVIDER_ID
 301    }
 302
 303    fn name(&self) -> LanguageModelProviderName {
 304        PROVIDER_NAME
 305    }
 306
 307    fn icon(&self) -> IconName {
 308        IconName::AiZed
 309    }
 310
 311    fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
 312        let default_model = self.state.read(cx).default_model.clone()?;
 313        let llm_api_token = self.state.read(cx).llm_api_token.clone();
 314        Some(self.create_language_model(default_model, llm_api_token))
 315    }
 316
 317    fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
 318        let default_fast_model = self.state.read(cx).default_fast_model.clone()?;
 319        let llm_api_token = self.state.read(cx).llm_api_token.clone();
 320        Some(self.create_language_model(default_fast_model, llm_api_token))
 321    }
 322
 323    fn recommended_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
 324        let llm_api_token = self.state.read(cx).llm_api_token.clone();
 325        self.state
 326            .read(cx)
 327            .recommended_models
 328            .iter()
 329            .cloned()
 330            .map(|model| self.create_language_model(model, llm_api_token.clone()))
 331            .collect()
 332    }
 333
 334    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
 335        let llm_api_token = self.state.read(cx).llm_api_token.clone();
 336        self.state
 337            .read(cx)
 338            .models
 339            .iter()
 340            .cloned()
 341            .map(|model| self.create_language_model(model, llm_api_token.clone()))
 342            .collect()
 343    }
 344
 345    fn is_authenticated(&self, cx: &App) -> bool {
 346        let state = self.state.read(cx);
 347        !state.is_signed_out(cx)
 348    }
 349
 350    fn authenticate(&self, _cx: &mut App) -> Task<Result<(), AuthenticateError>> {
 351        Task::ready(Ok(()))
 352    }
 353
 354    fn configuration_view(
 355        &self,
 356        _target_agent: language_model::ConfigurationViewTargetAgent,
 357        _: &mut Window,
 358        cx: &mut App,
 359    ) -> AnyView {
 360        cx.new(|_| ConfigurationView::new(self.state.clone()))
 361            .into()
 362    }
 363
 364    fn reset_credentials(&self, _cx: &mut App) -> Task<Result<()>> {
 365        Task::ready(Ok(()))
 366    }
 367}
 368
 369pub struct CloudLanguageModel {
 370    id: LanguageModelId,
 371    model: Arc<cloud_llm_client::LanguageModel>,
 372    llm_api_token: LlmApiToken,
 373    client: Arc<Client>,
 374    request_limiter: RateLimiter,
 375}
 376
 377struct PerformLlmCompletionResponse {
 378    response: Response<AsyncBody>,
 379    usage: Option<ModelRequestUsage>,
 380    tool_use_limit_reached: bool,
 381    includes_status_messages: bool,
 382}
 383
 384impl CloudLanguageModel {
 385    async fn perform_llm_completion(
 386        client: Arc<Client>,
 387        llm_api_token: LlmApiToken,
 388        app_version: Option<Version>,
 389        body: CompletionBody,
 390    ) -> Result<PerformLlmCompletionResponse> {
 391        let http_client = &client.http_client();
 392
 393        let mut token = llm_api_token.acquire(&client).await?;
 394        let mut refreshed_token = false;
 395
 396        loop {
 397            let request = http_client::Request::builder()
 398                .method(Method::POST)
 399                .uri(http_client.build_zed_llm_url("/completions", &[])?.as_ref())
 400                .when_some(app_version.as_ref(), |builder, app_version| {
 401                    builder.header(ZED_VERSION_HEADER_NAME, app_version.to_string())
 402                })
 403                .header("Content-Type", "application/json")
 404                .header("Authorization", format!("Bearer {token}"))
 405                .header(CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, "true")
 406                .body(serde_json::to_string(&body)?.into())?;
 407
 408            let mut response = http_client.send(request).await?;
 409            let status = response.status();
 410            if status.is_success() {
 411                let includes_status_messages = response
 412                    .headers()
 413                    .get(SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME)
 414                    .is_some();
 415
 416                let tool_use_limit_reached = response
 417                    .headers()
 418                    .get(TOOL_USE_LIMIT_REACHED_HEADER_NAME)
 419                    .is_some();
 420
 421                let usage = if includes_status_messages {
 422                    None
 423                } else {
 424                    ModelRequestUsage::from_headers(response.headers()).ok()
 425                };
 426
 427                return Ok(PerformLlmCompletionResponse {
 428                    response,
 429                    usage,
 430                    includes_status_messages,
 431                    tool_use_limit_reached,
 432                });
 433            }
 434
 435            if !refreshed_token
 436                && response
 437                    .headers()
 438                    .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
 439                    .is_some()
 440            {
 441                token = llm_api_token.refresh(&client).await?;
 442                refreshed_token = true;
 443                continue;
 444            }
 445
 446            if status == StatusCode::FORBIDDEN
 447                && response
 448                    .headers()
 449                    .get(SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME)
 450                    .is_some()
 451            {
 452                if let Some(MODEL_REQUESTS_RESOURCE_HEADER_VALUE) = response
 453                    .headers()
 454                    .get(SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME)
 455                    .and_then(|resource| resource.to_str().ok())
 456                    && let Some(plan) = response
 457                        .headers()
 458                        .get(CURRENT_PLAN_HEADER_NAME)
 459                        .and_then(|plan| plan.to_str().ok())
 460                        .and_then(|plan| cloud_llm_client::PlanV1::from_str(plan).ok())
 461                        .map(Plan::V1)
 462                {
 463                    return Err(anyhow!(ModelRequestLimitReachedError { plan }));
 464                }
 465            } else if status == StatusCode::PAYMENT_REQUIRED {
 466                return Err(anyhow!(PaymentRequiredError));
 467            }
 468
 469            let mut body = String::new();
 470            let headers = response.headers().clone();
 471            response.body_mut().read_to_string(&mut body).await?;
 472            return Err(anyhow!(ApiError {
 473                status,
 474                body,
 475                headers
 476            }));
 477        }
 478    }
 479}
 480
 481#[derive(Debug, Error)]
 482#[error("cloud language model request failed with status {status}: {body}")]
 483struct ApiError {
 484    status: StatusCode,
 485    body: String,
 486    headers: HeaderMap<HeaderValue>,
 487}
 488
 489/// Represents error responses from Zed's cloud API.
 490///
 491/// Example JSON for an upstream HTTP error:
 492/// ```json
 493/// {
 494///   "code": "upstream_http_error",
 495///   "message": "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout",
 496///   "upstream_status": 503
 497/// }
 498/// ```
 499#[derive(Debug, serde::Deserialize)]
 500struct CloudApiError {
 501    code: String,
 502    message: String,
 503    #[serde(default)]
 504    #[serde(deserialize_with = "deserialize_optional_status_code")]
 505    upstream_status: Option<StatusCode>,
 506    #[serde(default)]
 507    retry_after: Option<f64>,
 508}
 509
 510fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result<Option<StatusCode>, D::Error>
 511where
 512    D: serde::Deserializer<'de>,
 513{
 514    let opt: Option<u16> = Option::deserialize(deserializer)?;
 515    Ok(opt.and_then(|code| StatusCode::from_u16(code).ok()))
 516}
 517
 518impl From<ApiError> for LanguageModelCompletionError {
 519    fn from(error: ApiError) -> Self {
 520        if let Ok(cloud_error) = serde_json::from_str::<CloudApiError>(&error.body) {
 521            if cloud_error.code.starts_with("upstream_http_") {
 522                let status = if let Some(status) = cloud_error.upstream_status {
 523                    status
 524                } else if cloud_error.code.ends_with("_error") {
 525                    error.status
 526                } else {
 527                    // If there's a status code in the code string (e.g. "upstream_http_429")
 528                    // then use that; otherwise, see if the JSON contains a status code.
 529                    cloud_error
 530                        .code
 531                        .strip_prefix("upstream_http_")
 532                        .and_then(|code_str| code_str.parse::<u16>().ok())
 533                        .and_then(|code| StatusCode::from_u16(code).ok())
 534                        .unwrap_or(error.status)
 535                };
 536
 537                return LanguageModelCompletionError::UpstreamProviderError {
 538                    message: cloud_error.message,
 539                    status,
 540                    retry_after: cloud_error.retry_after.map(Duration::from_secs_f64),
 541                };
 542            }
 543
 544            return LanguageModelCompletionError::from_http_status(
 545                PROVIDER_NAME,
 546                error.status,
 547                cloud_error.message,
 548                None,
 549            );
 550        }
 551
 552        let retry_after = None;
 553        LanguageModelCompletionError::from_http_status(
 554            PROVIDER_NAME,
 555            error.status,
 556            error.body,
 557            retry_after,
 558        )
 559    }
 560}
 561
 562impl LanguageModel for CloudLanguageModel {
 563    fn id(&self) -> LanguageModelId {
 564        self.id.clone()
 565    }
 566
 567    fn name(&self) -> LanguageModelName {
 568        LanguageModelName::from(self.model.display_name.clone())
 569    }
 570
 571    fn provider_id(&self) -> LanguageModelProviderId {
 572        PROVIDER_ID
 573    }
 574
 575    fn provider_name(&self) -> LanguageModelProviderName {
 576        PROVIDER_NAME
 577    }
 578
 579    fn upstream_provider_id(&self) -> LanguageModelProviderId {
 580        use cloud_llm_client::LanguageModelProvider::*;
 581        match self.model.provider {
 582            Anthropic => language_model::ANTHROPIC_PROVIDER_ID,
 583            OpenAi => language_model::OPEN_AI_PROVIDER_ID,
 584            Google => language_model::GOOGLE_PROVIDER_ID,
 585            XAi => language_model::X_AI_PROVIDER_ID,
 586        }
 587    }
 588
 589    fn upstream_provider_name(&self) -> LanguageModelProviderName {
 590        use cloud_llm_client::LanguageModelProvider::*;
 591        match self.model.provider {
 592            Anthropic => language_model::ANTHROPIC_PROVIDER_NAME,
 593            OpenAi => language_model::OPEN_AI_PROVIDER_NAME,
 594            Google => language_model::GOOGLE_PROVIDER_NAME,
 595            XAi => language_model::X_AI_PROVIDER_NAME,
 596        }
 597    }
 598
 599    fn supports_tools(&self) -> bool {
 600        self.model.supports_tools
 601    }
 602
 603    fn supports_images(&self) -> bool {
 604        self.model.supports_images
 605    }
 606
 607    fn supports_streaming_tools(&self) -> bool {
 608        self.model.supports_streaming_tools
 609    }
 610
 611    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
 612        match choice {
 613            LanguageModelToolChoice::Auto
 614            | LanguageModelToolChoice::Any
 615            | LanguageModelToolChoice::None => true,
 616        }
 617    }
 618
 619    fn supports_burn_mode(&self) -> bool {
 620        self.model.supports_max_mode
 621    }
 622
 623    fn telemetry_id(&self) -> String {
 624        format!("zed.dev/{}", self.model.id)
 625    }
 626
 627    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
 628        match self.model.provider {
 629            cloud_llm_client::LanguageModelProvider::Anthropic
 630            | cloud_llm_client::LanguageModelProvider::OpenAi
 631            | cloud_llm_client::LanguageModelProvider::XAi => {
 632                LanguageModelToolSchemaFormat::JsonSchema
 633            }
 634            cloud_llm_client::LanguageModelProvider::Google => {
 635                LanguageModelToolSchemaFormat::JsonSchemaSubset
 636            }
 637        }
 638    }
 639
 640    fn max_token_count(&self) -> u64 {
 641        self.model.max_token_count as u64
 642    }
 643
 644    fn max_token_count_in_burn_mode(&self) -> Option<u64> {
 645        self.model
 646            .max_token_count_in_max_mode
 647            .filter(|_| self.model.supports_max_mode)
 648            .map(|max_token_count| max_token_count as u64)
 649    }
 650
 651    fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
 652        match &self.model.provider {
 653            cloud_llm_client::LanguageModelProvider::Anthropic => {
 654                Some(LanguageModelCacheConfiguration {
 655                    min_total_token: 2_048,
 656                    should_speculate: true,
 657                    max_cache_anchors: 4,
 658                })
 659            }
 660            cloud_llm_client::LanguageModelProvider::OpenAi
 661            | cloud_llm_client::LanguageModelProvider::XAi
 662            | cloud_llm_client::LanguageModelProvider::Google => None,
 663        }
 664    }
 665
 666    fn count_tokens(
 667        &self,
 668        request: LanguageModelRequest,
 669        cx: &App,
 670    ) -> BoxFuture<'static, Result<u64>> {
 671        match self.model.provider {
 672            cloud_llm_client::LanguageModelProvider::Anthropic => cx
 673                .background_spawn(async move { count_anthropic_tokens_with_tiktoken(request) })
 674                .boxed(),
 675            cloud_llm_client::LanguageModelProvider::OpenAi => {
 676                let model = match open_ai::Model::from_id(&self.model.id.0) {
 677                    Ok(model) => model,
 678                    Err(err) => return async move { Err(anyhow!(err)) }.boxed(),
 679                };
 680                count_open_ai_tokens(request, model, cx)
 681            }
 682            cloud_llm_client::LanguageModelProvider::XAi => {
 683                let model = match x_ai::Model::from_id(&self.model.id.0) {
 684                    Ok(model) => model,
 685                    Err(err) => return async move { Err(anyhow!(err)) }.boxed(),
 686                };
 687                count_xai_tokens(request, model, cx)
 688            }
 689            cloud_llm_client::LanguageModelProvider::Google => {
 690                let client = self.client.clone();
 691                let llm_api_token = self.llm_api_token.clone();
 692                let model_id = self.model.id.to_string();
 693                let generate_content_request =
 694                    into_google(request, model_id.clone(), GoogleModelMode::Default);
 695                async move {
 696                    let http_client = &client.http_client();
 697                    let token = llm_api_token.acquire(&client).await?;
 698
 699                    let request_body = CountTokensBody {
 700                        provider: cloud_llm_client::LanguageModelProvider::Google,
 701                        model: model_id,
 702                        provider_request: serde_json::to_value(&google_ai::CountTokensRequest {
 703                            generate_content_request,
 704                        })?,
 705                    };
 706                    let request = http_client::Request::builder()
 707                        .method(Method::POST)
 708                        .uri(
 709                            http_client
 710                                .build_zed_llm_url("/count_tokens", &[])?
 711                                .as_ref(),
 712                        )
 713                        .header("Content-Type", "application/json")
 714                        .header("Authorization", format!("Bearer {token}"))
 715                        .body(serde_json::to_string(&request_body)?.into())?;
 716                    let mut response = http_client.send(request).await?;
 717                    let status = response.status();
 718                    let headers = response.headers().clone();
 719                    let mut response_body = String::new();
 720                    response
 721                        .body_mut()
 722                        .read_to_string(&mut response_body)
 723                        .await?;
 724
 725                    if status.is_success() {
 726                        let response_body: CountTokensResponse =
 727                            serde_json::from_str(&response_body)?;
 728
 729                        Ok(response_body.tokens as u64)
 730                    } else {
 731                        Err(anyhow!(ApiError {
 732                            status,
 733                            body: response_body,
 734                            headers
 735                        }))
 736                    }
 737                }
 738                .boxed()
 739            }
 740        }
 741    }
 742
 743    fn stream_completion(
 744        &self,
 745        request: LanguageModelRequest,
 746        cx: &AsyncApp,
 747    ) -> BoxFuture<
 748        'static,
 749        Result<
 750            BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>,
 751            LanguageModelCompletionError,
 752        >,
 753    > {
 754        let thread_id = request.thread_id.clone();
 755        let prompt_id = request.prompt_id.clone();
 756        let intent = request.intent;
 757        let mode = request.mode;
 758        let app_version = cx.update(|cx| AppVersion::global(cx)).ok();
 759        let thinking_allowed = request.thinking_allowed;
 760        let provider_name = provider_name(&self.model.provider);
 761        match self.model.provider {
 762            cloud_llm_client::LanguageModelProvider::Anthropic => {
 763                let request = into_anthropic(
 764                    request,
 765                    self.model.id.to_string(),
 766                    1.0,
 767                    self.model.max_output_tokens as u64,
 768                    if thinking_allowed && self.model.id.0.ends_with("-thinking") {
 769                        AnthropicModelMode::Thinking {
 770                            budget_tokens: Some(4_096),
 771                        }
 772                    } else {
 773                        AnthropicModelMode::Default
 774                    },
 775                );
 776                let client = self.client.clone();
 777                let llm_api_token = self.llm_api_token.clone();
 778                let future = self.request_limiter.stream(async move {
 779                    let PerformLlmCompletionResponse {
 780                        response,
 781                        usage,
 782                        includes_status_messages,
 783                        tool_use_limit_reached,
 784                    } = Self::perform_llm_completion(
 785                        client.clone(),
 786                        llm_api_token,
 787                        app_version,
 788                        CompletionBody {
 789                            thread_id,
 790                            prompt_id,
 791                            intent,
 792                            mode,
 793                            provider: cloud_llm_client::LanguageModelProvider::Anthropic,
 794                            model: request.model.clone(),
 795                            provider_request: serde_json::to_value(&request)
 796                                .map_err(|e| anyhow!(e))?,
 797                        },
 798                    )
 799                    .await
 800                    .map_err(|err| match err.downcast::<ApiError>() {
 801                        Ok(api_err) => anyhow!(LanguageModelCompletionError::from(api_err)),
 802                        Err(err) => anyhow!(err),
 803                    })?;
 804
 805                    let mut mapper = AnthropicEventMapper::new();
 806                    Ok(map_cloud_completion_events(
 807                        Box::pin(
 808                            response_lines(response, includes_status_messages)
 809                                .chain(usage_updated_event(usage))
 810                                .chain(tool_use_limit_reached_event(tool_use_limit_reached)), // .map(|_| {}),
 811                        ),
 812                        &provider_name,
 813                        move |event| mapper.map_event(event),
 814                    ))
 815                });
 816                async move { Ok(future.await?.boxed()) }.boxed()
 817            }
 818            cloud_llm_client::LanguageModelProvider::OpenAi => {
 819                let client = self.client.clone();
 820                let request = into_open_ai(
 821                    request,
 822                    &self.model.id.0,
 823                    self.model.supports_parallel_tool_calls,
 824                    true,
 825                    None,
 826                    None,
 827                );
 828                let llm_api_token = self.llm_api_token.clone();
 829                let future = self.request_limiter.stream(async move {
 830                    let PerformLlmCompletionResponse {
 831                        response,
 832                        usage,
 833                        includes_status_messages,
 834                        tool_use_limit_reached,
 835                    } = Self::perform_llm_completion(
 836                        client.clone(),
 837                        llm_api_token,
 838                        app_version,
 839                        CompletionBody {
 840                            thread_id,
 841                            prompt_id,
 842                            intent,
 843                            mode,
 844                            provider: cloud_llm_client::LanguageModelProvider::OpenAi,
 845                            model: request.model.clone(),
 846                            provider_request: serde_json::to_value(&request)
 847                                .map_err(|e| anyhow!(e))?,
 848                        },
 849                    )
 850                    .await?;
 851
 852                    let mut mapper = OpenAiEventMapper::new();
 853                    Ok(map_cloud_completion_events(
 854                        Box::pin(
 855                            response_lines(response, includes_status_messages)
 856                                .chain(usage_updated_event(usage))
 857                                .chain(tool_use_limit_reached_event(tool_use_limit_reached)),
 858                        ),
 859                        &provider_name,
 860                        move |event| mapper.map_event(event),
 861                    ))
 862                });
 863                async move { Ok(future.await?.boxed()) }.boxed()
 864            }
 865            cloud_llm_client::LanguageModelProvider::XAi => {
 866                let client = self.client.clone();
 867                let request = into_open_ai(
 868                    request,
 869                    &self.model.id.0,
 870                    self.model.supports_parallel_tool_calls,
 871                    false,
 872                    None,
 873                    None,
 874                );
 875                let llm_api_token = self.llm_api_token.clone();
 876                let future = self.request_limiter.stream(async move {
 877                    let PerformLlmCompletionResponse {
 878                        response,
 879                        usage,
 880                        includes_status_messages,
 881                        tool_use_limit_reached,
 882                    } = Self::perform_llm_completion(
 883                        client.clone(),
 884                        llm_api_token,
 885                        app_version,
 886                        CompletionBody {
 887                            thread_id,
 888                            prompt_id,
 889                            intent,
 890                            mode,
 891                            provider: cloud_llm_client::LanguageModelProvider::XAi,
 892                            model: request.model.clone(),
 893                            provider_request: serde_json::to_value(&request)
 894                                .map_err(|e| anyhow!(e))?,
 895                        },
 896                    )
 897                    .await?;
 898
 899                    let mut mapper = OpenAiEventMapper::new();
 900                    Ok(map_cloud_completion_events(
 901                        Box::pin(
 902                            response_lines(response, includes_status_messages)
 903                                .chain(usage_updated_event(usage))
 904                                .chain(tool_use_limit_reached_event(tool_use_limit_reached)),
 905                        ),
 906                        &provider_name,
 907                        move |event| mapper.map_event(event),
 908                    ))
 909                });
 910                async move { Ok(future.await?.boxed()) }.boxed()
 911            }
 912            cloud_llm_client::LanguageModelProvider::Google => {
 913                let client = self.client.clone();
 914                let request =
 915                    into_google(request, self.model.id.to_string(), GoogleModelMode::Default);
 916                let llm_api_token = self.llm_api_token.clone();
 917                let future = self.request_limiter.stream(async move {
 918                    let PerformLlmCompletionResponse {
 919                        response,
 920                        usage,
 921                        includes_status_messages,
 922                        tool_use_limit_reached,
 923                    } = Self::perform_llm_completion(
 924                        client.clone(),
 925                        llm_api_token,
 926                        app_version,
 927                        CompletionBody {
 928                            thread_id,
 929                            prompt_id,
 930                            intent,
 931                            mode,
 932                            provider: cloud_llm_client::LanguageModelProvider::Google,
 933                            model: request.model.model_id.clone(),
 934                            provider_request: serde_json::to_value(&request)
 935                                .map_err(|e| anyhow!(e))?,
 936                        },
 937                    )
 938                    .await?;
 939
 940                    let mut mapper = GoogleEventMapper::new();
 941                    Ok(map_cloud_completion_events(
 942                        Box::pin(
 943                            response_lines(response, includes_status_messages)
 944                                .chain(usage_updated_event(usage))
 945                                .chain(tool_use_limit_reached_event(tool_use_limit_reached)),
 946                        ),
 947                        &provider_name,
 948                        move |event| mapper.map_event(event),
 949                    ))
 950                });
 951                async move { Ok(future.await?.boxed()) }.boxed()
 952            }
 953        }
 954    }
 955}
 956
 957fn map_cloud_completion_events<T, F>(
 958    stream: Pin<Box<dyn Stream<Item = Result<CompletionEvent<T>>> + Send>>,
 959    provider: &LanguageModelProviderName,
 960    mut map_callback: F,
 961) -> BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
 962where
 963    T: DeserializeOwned + 'static,
 964    F: FnMut(T) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
 965        + Send
 966        + 'static,
 967{
 968    let provider = provider.clone();
 969    stream
 970        .flat_map(move |event| {
 971            futures::stream::iter(match event {
 972                Err(error) => {
 973                    vec![Err(LanguageModelCompletionError::from(error))]
 974                }
 975                Ok(CompletionEvent::Status(event)) => {
 976                    vec![
 977                        LanguageModelCompletionEvent::from_completion_request_status(
 978                            event,
 979                            provider.clone(),
 980                        ),
 981                    ]
 982                }
 983                Ok(CompletionEvent::Event(event)) => map_callback(event),
 984            })
 985        })
 986        .boxed()
 987}
 988
 989fn provider_name(provider: &cloud_llm_client::LanguageModelProvider) -> LanguageModelProviderName {
 990    match provider {
 991        cloud_llm_client::LanguageModelProvider::Anthropic => {
 992            language_model::ANTHROPIC_PROVIDER_NAME
 993        }
 994        cloud_llm_client::LanguageModelProvider::OpenAi => language_model::OPEN_AI_PROVIDER_NAME,
 995        cloud_llm_client::LanguageModelProvider::Google => language_model::GOOGLE_PROVIDER_NAME,
 996        cloud_llm_client::LanguageModelProvider::XAi => language_model::X_AI_PROVIDER_NAME,
 997    }
 998}
 999
1000fn usage_updated_event<T>(
1001    usage: Option<ModelRequestUsage>,
1002) -> impl Stream<Item = Result<CompletionEvent<T>>> {
1003    futures::stream::iter(usage.map(|usage| {
1004        Ok(CompletionEvent::Status(
1005            CompletionRequestStatus::UsageUpdated {
1006                amount: usage.amount as usize,
1007                limit: usage.limit,
1008            },
1009        ))
1010    }))
1011}
1012
1013fn tool_use_limit_reached_event<T>(
1014    tool_use_limit_reached: bool,
1015) -> impl Stream<Item = Result<CompletionEvent<T>>> {
1016    futures::stream::iter(tool_use_limit_reached.then(|| {
1017        Ok(CompletionEvent::Status(
1018            CompletionRequestStatus::ToolUseLimitReached,
1019        ))
1020    }))
1021}
1022
1023fn response_lines<T: DeserializeOwned>(
1024    response: Response<AsyncBody>,
1025    includes_status_messages: bool,
1026) -> impl Stream<Item = Result<CompletionEvent<T>>> {
1027    futures::stream::try_unfold(
1028        (String::new(), BufReader::new(response.into_body())),
1029        move |(mut line, mut body)| async move {
1030            match body.read_line(&mut line).await {
1031                Ok(0) => Ok(None),
1032                Ok(_) => {
1033                    let event = if includes_status_messages {
1034                        serde_json::from_str::<CompletionEvent<T>>(&line)?
1035                    } else {
1036                        CompletionEvent::Event(serde_json::from_str::<T>(&line)?)
1037                    };
1038
1039                    line.clear();
1040                    Ok(Some((event, (line, body))))
1041                }
1042                Err(e) => Err(e.into()),
1043            }
1044        },
1045    )
1046}
1047
1048#[derive(IntoElement, RegisterComponent)]
1049struct ZedAiConfiguration {
1050    is_connected: bool,
1051    plan: Option<Plan>,
1052    subscription_period: Option<(DateTime<Utc>, DateTime<Utc>)>,
1053    eligible_for_trial: bool,
1054    account_too_young: bool,
1055    sign_in_callback: Arc<dyn Fn(&mut Window, &mut App) + Send + Sync>,
1056}
1057
1058impl RenderOnce for ZedAiConfiguration {
1059    fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
1060        let is_pro = self.plan.is_some_and(|plan| {
1061            matches!(plan, Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro))
1062        });
1063        let subscription_text = match (self.plan, self.subscription_period) {
1064            (Some(Plan::V1(PlanV1::ZedPro) | Plan::V2(PlanV2::ZedPro)), Some(_)) => {
1065                "You have access to Zed's hosted models through your Pro subscription."
1066            }
1067            (Some(Plan::V1(PlanV1::ZedProTrial) | Plan::V2(PlanV2::ZedProTrial)), Some(_)) => {
1068                "You have access to Zed's hosted models through your Pro trial."
1069            }
1070            (Some(Plan::V1(PlanV1::ZedFree)), Some(_)) => {
1071                "You have basic access to Zed's hosted models through the Free plan."
1072            }
1073            (Some(Plan::V2(PlanV2::ZedFree)), Some(_)) => {
1074                if self.eligible_for_trial {
1075                    "Subscribe for access to Zed's hosted models. Start with a 14 day free trial."
1076                } else {
1077                    "Subscribe for access to Zed's hosted models."
1078                }
1079            }
1080            _ => {
1081                if self.eligible_for_trial {
1082                    "Subscribe for access to Zed's hosted models. Start with a 14 day free trial."
1083                } else {
1084                    "Subscribe for access to Zed's hosted models."
1085                }
1086            }
1087        };
1088
1089        let manage_subscription_buttons = if is_pro {
1090            Button::new("manage_settings", "Manage Subscription")
1091                .full_width()
1092                .style(ButtonStyle::Tinted(TintColor::Accent))
1093                .on_click(|_, _, cx| cx.open_url(&zed_urls::account_url(cx)))
1094                .into_any_element()
1095        } else if self.plan.is_none() || self.eligible_for_trial {
1096            Button::new("start_trial", "Start 14-day Free Pro Trial")
1097                .full_width()
1098                .style(ui::ButtonStyle::Tinted(ui::TintColor::Accent))
1099                .on_click(|_, _, cx| cx.open_url(&zed_urls::start_trial_url(cx)))
1100                .into_any_element()
1101        } else {
1102            Button::new("upgrade", "Upgrade to Pro")
1103                .full_width()
1104                .style(ui::ButtonStyle::Tinted(ui::TintColor::Accent))
1105                .on_click(|_, _, cx| cx.open_url(&zed_urls::upgrade_to_zed_pro_url(cx)))
1106                .into_any_element()
1107        };
1108
1109        if !self.is_connected {
1110            return v_flex()
1111                .gap_2()
1112                .child(Label::new("Sign in to have access to Zed's complete agentic experience with hosted models."))
1113                .child(
1114                    Button::new("sign_in", "Sign In to use Zed AI")
1115                        .icon_color(Color::Muted)
1116                        .icon(IconName::Github)
1117                        .icon_size(IconSize::Small)
1118                        .icon_position(IconPosition::Start)
1119                        .full_width()
1120                        .on_click({
1121                            let callback = self.sign_in_callback.clone();
1122                            move |_, window, cx| (callback)(window, cx)
1123                        }),
1124                );
1125        }
1126
1127        v_flex().gap_2().w_full().map(|this| {
1128            if self.account_too_young {
1129                this.child(YoungAccountBanner).child(
1130                    Button::new("upgrade", "Upgrade to Pro")
1131                        .style(ui::ButtonStyle::Tinted(ui::TintColor::Accent))
1132                        .full_width()
1133                        .on_click(|_, _, cx| cx.open_url(&zed_urls::upgrade_to_zed_pro_url(cx))),
1134                )
1135            } else {
1136                this.text_sm()
1137                    .child(subscription_text)
1138                    .child(manage_subscription_buttons)
1139            }
1140        })
1141    }
1142}
1143
1144struct ConfigurationView {
1145    state: Entity<State>,
1146    sign_in_callback: Arc<dyn Fn(&mut Window, &mut App) + Send + Sync>,
1147}
1148
1149impl ConfigurationView {
1150    fn new(state: Entity<State>) -> Self {
1151        let sign_in_callback = Arc::new({
1152            let state = state.clone();
1153            move |_window: &mut Window, cx: &mut App| {
1154                state.update(cx, |state, cx| {
1155                    state.authenticate(cx).detach_and_log_err(cx);
1156                });
1157            }
1158        });
1159
1160        Self {
1161            state,
1162            sign_in_callback,
1163        }
1164    }
1165}
1166
1167impl Render for ConfigurationView {
1168    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1169        let state = self.state.read(cx);
1170        let user_store = state.user_store.read(cx);
1171
1172        ZedAiConfiguration {
1173            is_connected: !state.is_signed_out(cx),
1174            plan: user_store.plan(),
1175            subscription_period: user_store.subscription_period(),
1176            eligible_for_trial: user_store.trial_started_at().is_none(),
1177            account_too_young: user_store.account_too_young(),
1178            sign_in_callback: self.sign_in_callback.clone(),
1179        }
1180    }
1181}
1182
1183impl Component for ZedAiConfiguration {
1184    fn name() -> &'static str {
1185        "AI Configuration Content"
1186    }
1187
1188    fn sort_name() -> &'static str {
1189        "AI Configuration Content"
1190    }
1191
1192    fn scope() -> ComponentScope {
1193        ComponentScope::Onboarding
1194    }
1195
1196    fn preview(_window: &mut Window, _cx: &mut App) -> Option<AnyElement> {
1197        fn configuration(
1198            is_connected: bool,
1199            plan: Option<Plan>,
1200            eligible_for_trial: bool,
1201            account_too_young: bool,
1202        ) -> AnyElement {
1203            ZedAiConfiguration {
1204                is_connected,
1205                plan,
1206                subscription_period: plan
1207                    .is_some()
1208                    .then(|| (Utc::now(), Utc::now() + chrono::Duration::days(7))),
1209                eligible_for_trial,
1210                account_too_young,
1211                sign_in_callback: Arc::new(|_, _| {}),
1212            }
1213            .into_any_element()
1214        }
1215
1216        Some(
1217            v_flex()
1218                .p_4()
1219                .gap_4()
1220                .children(vec![
1221                    single_example("Not connected", configuration(false, None, false, false)),
1222                    single_example(
1223                        "Accept Terms of Service",
1224                        configuration(true, None, true, false),
1225                    ),
1226                    single_example(
1227                        "No Plan - Not eligible for trial",
1228                        configuration(true, None, false, false),
1229                    ),
1230                    single_example(
1231                        "No Plan - Eligible for trial",
1232                        configuration(true, None, true, false),
1233                    ),
1234                    single_example(
1235                        "Free Plan",
1236                        configuration(true, Some(Plan::V1(PlanV1::ZedFree)), true, false),
1237                    ),
1238                    single_example(
1239                        "Zed Pro Trial Plan",
1240                        configuration(true, Some(Plan::V1(PlanV1::ZedProTrial)), true, false),
1241                    ),
1242                    single_example(
1243                        "Zed Pro Plan",
1244                        configuration(true, Some(Plan::V1(PlanV1::ZedPro)), true, false),
1245                    ),
1246                ])
1247                .into_any_element(),
1248        )
1249    }
1250}
1251
1252#[cfg(test)]
1253mod tests {
1254    use super::*;
1255    use http_client::http::{HeaderMap, StatusCode};
1256    use language_model::LanguageModelCompletionError;
1257
1258    #[test]
1259    fn test_api_error_conversion_with_upstream_http_error() {
1260        // upstream_http_error with 503 status should become ServerOverloaded
1261        let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout","upstream_status":503}"#;
1262
1263        let api_error = ApiError {
1264            status: StatusCode::INTERNAL_SERVER_ERROR,
1265            body: error_body.to_string(),
1266            headers: HeaderMap::new(),
1267        };
1268
1269        let completion_error: LanguageModelCompletionError = api_error.into();
1270
1271        match completion_error {
1272            LanguageModelCompletionError::UpstreamProviderError { message, .. } => {
1273                assert_eq!(
1274                    message,
1275                    "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout"
1276                );
1277            }
1278            _ => panic!(
1279                "Expected UpstreamProviderError for upstream 503, got: {:?}",
1280                completion_error
1281            ),
1282        }
1283
1284        // upstream_http_error with 500 status should become ApiInternalServerError
1285        let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the OpenAI API: internal server error","upstream_status":500}"#;
1286
1287        let api_error = ApiError {
1288            status: StatusCode::INTERNAL_SERVER_ERROR,
1289            body: error_body.to_string(),
1290            headers: HeaderMap::new(),
1291        };
1292
1293        let completion_error: LanguageModelCompletionError = api_error.into();
1294
1295        match completion_error {
1296            LanguageModelCompletionError::UpstreamProviderError { message, .. } => {
1297                assert_eq!(
1298                    message,
1299                    "Received an error from the OpenAI API: internal server error"
1300                );
1301            }
1302            _ => panic!(
1303                "Expected UpstreamProviderError for upstream 500, got: {:?}",
1304                completion_error
1305            ),
1306        }
1307
1308        // upstream_http_error with 429 status should become RateLimitExceeded
1309        let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Google API: rate limit exceeded","upstream_status":429}"#;
1310
1311        let api_error = ApiError {
1312            status: StatusCode::INTERNAL_SERVER_ERROR,
1313            body: error_body.to_string(),
1314            headers: HeaderMap::new(),
1315        };
1316
1317        let completion_error: LanguageModelCompletionError = api_error.into();
1318
1319        match completion_error {
1320            LanguageModelCompletionError::UpstreamProviderError { message, .. } => {
1321                assert_eq!(
1322                    message,
1323                    "Received an error from the Google API: rate limit exceeded"
1324                );
1325            }
1326            _ => panic!(
1327                "Expected UpstreamProviderError for upstream 429, got: {:?}",
1328                completion_error
1329            ),
1330        }
1331
1332        // Regular 500 error without upstream_http_error should remain ApiInternalServerError for Zed
1333        let error_body = "Regular internal server error";
1334
1335        let api_error = ApiError {
1336            status: StatusCode::INTERNAL_SERVER_ERROR,
1337            body: error_body.to_string(),
1338            headers: HeaderMap::new(),
1339        };
1340
1341        let completion_error: LanguageModelCompletionError = api_error.into();
1342
1343        match completion_error {
1344            LanguageModelCompletionError::ApiInternalServerError { provider, message } => {
1345                assert_eq!(provider, PROVIDER_NAME);
1346                assert_eq!(message, "Regular internal server error");
1347            }
1348            _ => panic!(
1349                "Expected ApiInternalServerError for regular 500, got: {:?}",
1350                completion_error
1351            ),
1352        }
1353
1354        // upstream_http_429 format should be converted to UpstreamProviderError
1355        let error_body = r#"{"code":"upstream_http_429","message":"Upstream Anthropic rate limit exceeded.","retry_after":30.5}"#;
1356
1357        let api_error = ApiError {
1358            status: StatusCode::INTERNAL_SERVER_ERROR,
1359            body: error_body.to_string(),
1360            headers: HeaderMap::new(),
1361        };
1362
1363        let completion_error: LanguageModelCompletionError = api_error.into();
1364
1365        match completion_error {
1366            LanguageModelCompletionError::UpstreamProviderError {
1367                message,
1368                status,
1369                retry_after,
1370            } => {
1371                assert_eq!(message, "Upstream Anthropic rate limit exceeded.");
1372                assert_eq!(status, StatusCode::TOO_MANY_REQUESTS);
1373                assert_eq!(retry_after, Some(Duration::from_secs_f64(30.5)));
1374            }
1375            _ => panic!(
1376                "Expected UpstreamProviderError for upstream_http_429, got: {:?}",
1377                completion_error
1378            ),
1379        }
1380
1381        // Invalid JSON in error body should fall back to regular error handling
1382        let error_body = "Not JSON at all";
1383
1384        let api_error = ApiError {
1385            status: StatusCode::INTERNAL_SERVER_ERROR,
1386            body: error_body.to_string(),
1387            headers: HeaderMap::new(),
1388        };
1389
1390        let completion_error: LanguageModelCompletionError = api_error.into();
1391
1392        match completion_error {
1393            LanguageModelCompletionError::ApiInternalServerError { provider, .. } => {
1394                assert_eq!(provider, PROVIDER_NAME);
1395            }
1396            _ => panic!(
1397                "Expected ApiInternalServerError for invalid JSON, got: {:?}",
1398                completion_error
1399            ),
1400        }
1401    }
1402}