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