Skip to repository content1041 lines · 39.4 KB · rust
tenant.openagents/omega
No repository description is available.
OpenAgents Git authority 2026-07-28T04:33:24.668Z Public web read
NIP-34 coordinate
30617:7649603503856e5148d571eac2766b288a8ff1e9e35d380337a1d2b0015b4f92:omegaMaintainersHidden in public view
References2 branches · 1 tag
Read-only clone
git clone https://openagents.com/git/tenant.openagents/omega.gitBrowse files
cloud.rs
1use anyhow::{Result, anyhow};
2use client::{
3 Client, RefreshLlmTokenListener, TelemetrySettings, UserStore, global_llm_token, zed_urls,
4};
5use cloud_api_client::LlmApiToken;
6use cloud_api_types::OrganizationId;
7use cloud_api_types::Plan;
8use futures::FutureExt;
9use futures::StreamExt;
10use futures::future::BoxFuture;
11use gpui::{AnyElement, App, AppContext, Context, Entity, Subscription, Task, TaskExt};
12use language_model::{
13 AuthenticateError, FastModeConfirmation, IconOrSvg, InlineDescription, LanguageModel,
14 LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
15 LanguageModelProviderState, ProviderSettingsView, ZED_CLOUD_PROVIDER_ID,
16 ZED_CLOUD_PROVIDER_NAME,
17};
18use language_models_cloud::{CloudLlmTokenProvider, CloudModelProvider};
19use rand::{Rng as _, SeedableRng as _, rngs::StdRng};
20use release_channel::AppVersion;
21
22use settings::SettingsStore;
23pub use settings::ZedDotDevAvailableModel as AvailableModel;
24pub use settings::ZedDotDevAvailableProvider as AvailableProvider;
25use std::sync::Arc;
26use std::time::Duration;
27use ui::{TintColor, prelude::*};
28
29const PROVIDER_ID: LanguageModelProviderId = ZED_CLOUD_PROVIDER_ID;
30const PROVIDER_NAME: LanguageModelProviderName = ZED_CLOUD_PROVIDER_NAME;
31const MODELS_REFRESH_DEBOUNCE: Duration = Duration::from_secs(5 * 60);
32
33struct ClientTokenProvider {
34 client: Arc<Client>,
35 llm_api_token: LlmApiToken,
36 user_store: Entity<UserStore>,
37}
38
39impl CloudLlmTokenProvider for ClientTokenProvider {
40 type AuthContext = Option<OrganizationId>;
41
42 fn auth_context(&self, cx: &impl AppContext) -> Self::AuthContext {
43 self.user_store.read_with(cx, |user_store, _| {
44 user_store
45 .current_organization()
46 .map(|organization| organization.id.clone())
47 })
48 }
49
50 fn cached_token(
51 &self,
52 organization_id: Self::AuthContext,
53 ) -> BoxFuture<'static, Result<String>> {
54 let client = self.client.clone();
55 let llm_api_token = self.llm_api_token.clone();
56 Box::pin(async move {
57 let organization_id =
58 organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
59 client
60 .cached_llm_token(&llm_api_token, organization_id)
61 .await
62 })
63 }
64
65 fn refresh_token(
66 &self,
67 organization_id: Self::AuthContext,
68 ) -> BoxFuture<'static, Result<String>> {
69 let client = self.client.clone();
70 let llm_api_token = self.llm_api_token.clone();
71 Box::pin(async move {
72 let organization_id =
73 organization_id.ok_or_else(|| anyhow!("No organization selected."))?;
74 client
75 .refresh_llm_token(&llm_api_token, organization_id)
76 .await
77 })
78 }
79
80 fn has_data_retention_consent(&self, cx: &impl AppContext) -> bool {
81 cx.read_global(|settings_store: &SettingsStore, _| {
82 settings_store
83 .get::<TelemetrySettings>(None)
84 .anthropic_retention
85 })
86 }
87}
88
89#[derive(Default, Clone, Debug, PartialEq)]
90pub struct ZedDotDevSettings {
91 pub available_models: Vec<AvailableModel>,
92}
93
94pub struct CloudLanguageModelProvider {
95 state: Entity<State>,
96 _maintain_client_status: Task<()>,
97}
98
99pub struct State {
100 client: Arc<Client>,
101 user_store: Entity<UserStore>,
102 status: client::Status,
103 provider: Entity<CloudModelProvider<ClientTokenProvider>>,
104 pending_models_refresh: Option<Task<()>>,
105 _user_store_subscription: Subscription,
106 _settings_subscription: Subscription,
107 _llm_token_subscription: Subscription,
108 _provider_subscription: Subscription,
109 _cloud_reconnect_task: Task<()>,
110}
111
112impl State {
113 fn new(
114 client: Arc<Client>,
115 user_store: Entity<UserStore>,
116 status: client::Status,
117 cx: &mut Context<Self>,
118 ) -> Self {
119 let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
120 let token_provider = Arc::new(ClientTokenProvider {
121 client: client.clone(),
122 llm_api_token: global_llm_token(cx),
123 user_store: user_store.clone(),
124 });
125
126 let provider = cx.new(|cx| {
127 CloudModelProvider::new(
128 token_provider.clone(),
129 client.http_client(),
130 Some(AppVersion::global(cx)),
131 )
132 });
133
134 let cloud_reconnect_task = cx.spawn({
135 let client = client.clone();
136 async move |this, cx| {
137 let mut connection_id_rx = client.cloud_connection_id();
138 while let Some(connection_id) = connection_id_rx.next().await {
139 // The initial value `0` means no connection has been
140 // established since this `Client` was created; only real
141 // reconnects trigger a refresh.
142 if connection_id == 0 {
143 continue;
144 }
145 if this
146 .update(cx, |this, cx| this.schedule_debounced_models_refresh(cx))
147 .is_err()
148 {
149 break;
150 }
151 }
152 }
153 });
154
155 Self {
156 client: client.clone(),
157 user_store: user_store.clone(),
158 status,
159 pending_models_refresh: None,
160 _provider_subscription: cx.observe(&provider, |_, _, cx| cx.notify()),
161 provider,
162 _user_store_subscription: cx.subscribe(
163 &user_store,
164 move |this, _user_store, event, cx| match event {
165 client::user::Event::PrivateUserInfoUpdated => {
166 let status = *client.status().borrow();
167 if status.is_signed_out() {
168 return;
169 }
170
171 this.refresh_models(cx);
172 }
173 _ => {}
174 },
175 ),
176 _settings_subscription: cx.observe_global::<SettingsStore>(|_, cx| {
177 cx.notify();
178 }),
179 _llm_token_subscription: cx.subscribe(
180 &refresh_llm_token_listener,
181 move |this, _listener, _event, cx| {
182 this.refresh_models(cx);
183 },
184 ),
185 _cloud_reconnect_task: cloud_reconnect_task,
186 }
187 }
188
189 fn is_signed_out(&self, cx: &App) -> bool {
190 self.status.is_signed_out() || self.user_store.read(cx).current_user().is_none()
191 }
192
193 fn sign_in(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
194 let client = self.client.clone();
195 let mut current_user = self.user_store.read(cx).watch_current_user();
196 cx.spawn(async move |state, cx| {
197 client.sign_in_with_optional_connect(true, cx).await?;
198 while current_user.borrow().is_none() {
199 current_user.next().await;
200 }
201 state.update(cx, |_, cx| {
202 cx.notify();
203 })
204 })
205 }
206
207 fn refresh_models(&mut self, cx: &mut Context<Self>) {
208 self.provider.update(cx, |provider, cx| {
209 provider.refresh_models(cx).detach_and_log_err(cx);
210 });
211 }
212
213 /// Schedules a model list refresh, replacing any previously scheduled
214 /// refresh.
215 fn schedule_debounced_models_refresh(&mut self, cx: &mut Context<Self>) {
216 self.pending_models_refresh = Some(cx.spawn(async move |this, cx| {
217 #[cfg(any(test, feature = "test-support"))]
218 let mut rng = StdRng::seed_from_u64(0);
219 #[cfg(not(any(test, feature = "test-support")))]
220 let mut rng = StdRng::from_os_rng();
221 let jitter = Duration::from_millis(
222 rng.random_range(0..MODELS_REFRESH_DEBOUNCE.as_millis() as u64),
223 );
224 cx.background_executor()
225 .timer(MODELS_REFRESH_DEBOUNCE + jitter)
226 .await;
227 this.update(cx, |this, cx| this.refresh_models(cx)).ok();
228 }));
229 }
230}
231
232impl CloudLanguageModelProvider {
233 pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self {
234 let mut status_rx = client.status();
235 let status = *status_rx.borrow();
236
237 let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx));
238
239 let state_ref = state.downgrade();
240 let maintain_client_status = cx.spawn(async move |cx| {
241 while let Some(status) = status_rx.next().await {
242 if let Some(this) = state_ref.upgrade() {
243 _ = this.update(cx, |this, cx| {
244 if this.status != status {
245 this.status = status;
246 if status.is_signed_out() {
247 this.provider.update(cx, |provider, cx| {
248 provider.clear_models();
249 cx.notify();
250 });
251 }
252 cx.notify();
253 }
254 });
255 } else {
256 break;
257 }
258 }
259 });
260
261 Self {
262 state,
263 _maintain_client_status: maintain_client_status,
264 }
265 }
266}
267
268impl LanguageModelProviderState for CloudLanguageModelProvider {
269 type ObservableEntity = State;
270
271 fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
272 Some(self.state.clone())
273 }
274}
275
276impl LanguageModelProvider for CloudLanguageModelProvider {
277 fn id(&self) -> LanguageModelProviderId {
278 PROVIDER_ID
279 }
280
281 fn name(&self) -> LanguageModelProviderName {
282 PROVIDER_NAME
283 }
284
285 fn icon(&self) -> IconOrSvg {
286 IconOrSvg::Icon(IconName::AiZed)
287 }
288
289 fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
290 let state = self.state.read(cx);
291 let provider = state.provider.read(cx);
292 let model = provider.default_model()?;
293 Some(provider.create_model(model))
294 }
295
296 fn default_fast_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
297 let state = self.state.read(cx);
298 let provider = state.provider.read(cx);
299 let model = provider.default_fast_model()?;
300 Some(provider.create_model(model))
301 }
302
303 fn recommended_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
304 let state = self.state.read(cx);
305 let provider = state.provider.read(cx);
306 provider
307 .recommended_models()
308 .iter()
309 .map(|model| provider.create_model(model))
310 .collect()
311 }
312
313 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
314 let state = self.state.read(cx);
315 let provider = state.provider.read(cx);
316 provider
317 .models()
318 .iter()
319 .map(|model| provider.create_model(model))
320 .collect()
321 }
322
323 fn is_authenticated(&self, cx: &App) -> bool {
324 let state = self.state.read(cx);
325 !state.is_signed_out(cx)
326 }
327
328 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
329 if self.is_authenticated(cx) {
330 return Task::ready(Ok(()));
331 }
332 let mut status = self.state.read(cx).client.status();
333 let mut current_user = self.state.read(cx).user_store.read(cx).watch_current_user();
334 if !status.borrow().is_signing_in() {
335 return Task::ready(Ok(()));
336 }
337 cx.background_spawn(async move {
338 while status.borrow().is_signing_in() {
339 status.next().await;
340 }
341 while current_user.borrow().is_none() {
342 let current_status = *status.borrow();
343 if !matches!(
344 current_status,
345 client::Status::Authenticated
346 | client::Status::Reauthenticated
347 | client::Status::Connected { .. }
348 ) {
349 return Err(AuthenticateError::Other(anyhow!(
350 "sign-in did not complete: {current_status:?}"
351 )));
352 }
353 futures::select_biased! {
354 _ = current_user.next().fuse() => {},
355 _ = status.next().fuse() => {},
356 }
357 }
358 Ok(())
359 })
360 }
361
362 fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView> {
363 let state = self.state.read(cx);
364 let user_store = state.user_store.read(cx);
365 let is_zed_model_provider_enabled = user_store
366 .current_organization_configuration()
367 .map_or(true, |config| config.is_zed_model_provider_enabled);
368 let description = InlineDescription::Text(
369 zed_ai_description(
370 !state.is_signed_out(cx),
371 user_store.plan(),
372 is_zed_model_provider_enabled,
373 user_store.trial_started_at().is_none(),
374 )
375 .into(),
376 );
377
378 let title = if state.is_signed_out(cx) {
379 None
380 } else {
381 match state.user_store.read(cx).plan() {
382 Some(Plan::ZedPro) => Some("Subscribed to Pro".into()),
383 Some(Plan::ZedProTrial) => Some("Subscribed to Pro Trial".into()),
384 Some(Plan::ZedStudent) => Some("Subscribed to Student".into()),
385 Some(Plan::ZedBusiness) => Some("Subscribed to Business".into()),
386 Some(Plan::ZedVip) => Some("Subscribed to VIP".into()),
387 Some(Plan::ZedFree) | None => None,
388 }
389 };
390
391 Some(ProviderSettingsView::Inline(
392 language_model::InlineProviderSettings {
393 title,
394 description: Some(description),
395 create_view: Arc::new({
396 let state = self.state.clone();
397 move |_window, cx| {
398 cx.new(|_| ConfigurationView::new(state.clone(), true))
399 .into()
400 }
401 }),
402 },
403 ))
404 }
405
406 fn authentication_error_message(&self) -> SharedString {
407 "Failed to sign in with your Zed account (401).".into()
408 }
409
410 fn missing_credentials_error_message(&self) -> SharedString {
411 "You are not signed in to your Zed account. \
412 Sign in to continue."
413 .into()
414 }
415
416 fn fast_mode_confirmation(&self, _cx: &App) -> Option<FastModeConfirmation> {
417 Some(FastModeConfirmation {
418 title: "Enable Fast Mode for Zed AI?".into(),
419 message: "Fast mode routes requests through the upstream provider's fast mode or priority tier. The \
420 upstream provider's premium per-token pricing applies and is passed through to \
421 your Zed billing."
422 .into(),
423 })
424 }
425}
426
427#[derive(IntoElement, RegisterComponent)]
428struct ZedAiConfiguration {
429 is_connected: bool,
430 plan: Option<Plan>,
431 is_zed_model_provider_enabled: bool,
432 eligible_for_trial: bool,
433 account_too_young: bool,
434 compact: bool,
435 sign_in_callback: Arc<dyn Fn(&mut Window, &mut App) + Send + Sync>,
436}
437
438fn zed_ai_description(
439 is_connected: bool,
440 _plan: Option<Plan>,
441 is_zed_model_provider_enabled: bool,
442 _eligible_for_trial: bool,
443) -> &'static str {
444 // Omega does not sell a hosted AI service, so this used to render a
445 // subscription pitch for a product that does not exist here, naming Zed as
446 // the vendor: Pro, Student, Business, VIP, and "Start with a 14 day free
447 // trial". omega#16 forbids presenting Zed as the product, and the copy also
448 // told the operator to buy something they cannot buy. See OMEGA-DELTA-0008.
449 //
450 // The plan and trial-eligibility inputs are kept in the signature because
451 // the upstream call sites still pass them; the distinction they encoded is
452 // one Omega deliberately does not make.
453 if !is_connected {
454 return "Sign in to use hosted models from a configured provider.";
455 }
456 if is_zed_model_provider_enabled {
457 "Hosted models are available from your configured provider."
458 } else {
459 "Hosted models are disabled by your current configuration."
460 }
461}
462
463impl RenderOnce for ZedAiConfiguration {
464 fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
465 let has_paid_plan = matches!(
466 self.plan,
467 Some(Plan::ZedPro | Plan::ZedStudent | Plan::ZedBusiness | Plan::ZedVip)
468 );
469
470 let description = zed_ai_description(
471 self.is_connected,
472 self.plan,
473 self.is_zed_model_provider_enabled,
474 self.eligible_for_trial,
475 );
476
477 let manage_subscription_buttons = if has_paid_plan {
478 Button::new("manage_settings", "Manage Subscription")
479 .when(!self.compact, |this| {
480 this.full_width().label_size(LabelSize::Small)
481 })
482 .when(self.compact, |this| this.size(ButtonSize::Medium))
483 .style(ButtonStyle::Tinted(TintColor::Accent))
484 .on_click(|_, _, cx| cx.open_url(&zed_urls::account_url(cx)))
485 .into_any_element()
486 } else if self.plan.is_none() || self.eligible_for_trial {
487 Button::new("start_trial", "Start 14-day Free Pro Trial")
488 .when(!self.compact, |this| {
489 this.full_width().label_size(LabelSize::Small)
490 })
491 .when(self.compact, |this| this.size(ButtonSize::Medium))
492 .style(ui::ButtonStyle::Tinted(ui::TintColor::Accent))
493 .on_click(|_, _, cx| cx.open_url(&zed_urls::start_trial_url(cx)))
494 .into_any_element()
495 } else {
496 Button::new("upgrade", "Upgrade to Pro")
497 .when(!self.compact, |this| {
498 this.full_width().label_size(LabelSize::Small)
499 })
500 .when(self.compact, |this| this.size(ButtonSize::Medium))
501 .style(ui::ButtonStyle::Tinted(ui::TintColor::Accent))
502 .on_click(|_, _, cx| cx.open_url(&zed_urls::upgrade_to_zed_pro_url(cx)))
503 .into_any_element()
504 };
505
506 if !self.is_connected {
507 return v_flex()
508 .gap_2()
509 .when(!self.compact, |this| this.child(Label::new(description)))
510 .child(
511 Button::new("sign_in", "Sign In to use Zed AI")
512 .start_icon(
513 Icon::new(IconName::Github)
514 .size(IconSize::Small)
515 .color(Color::Muted),
516 )
517 .when(!self.compact, |this| this.full_width())
518 .on_click({
519 let callback = self.sign_in_callback.clone();
520 move |_, window, cx| (callback)(window, cx)
521 }),
522 );
523 }
524
525 v_flex()
526 .gap_2()
527 .when(!self.compact, |this| this.w_full())
528 .map(|this| {
529 if self.account_too_young {
530 this.child(
531 Button::new("upgrade", "Upgrade to Pro")
532 .style(ui::ButtonStyle::Tinted(ui::TintColor::Accent))
533 .when(!self.compact, |this| this.full_width())
534 .on_click(|_, _, cx| {
535 cx.open_url(&zed_urls::upgrade_to_zed_pro_url(cx))
536 }),
537 )
538 } else {
539 this.when(!self.compact, |this| this.text_sm().child(description))
540 .child(manage_subscription_buttons)
541 }
542 })
543 }
544}
545
546struct ConfigurationView {
547 state: Entity<State>,
548 compact: bool,
549 sign_in_callback: Arc<dyn Fn(&mut Window, &mut App) + Send + Sync>,
550}
551
552impl ConfigurationView {
553 fn new(state: Entity<State>, compact: bool) -> Self {
554 let sign_in_callback = Arc::new({
555 let state = state.clone();
556 move |_window: &mut Window, cx: &mut App| {
557 state.update(cx, |state, cx| {
558 state.sign_in(cx).detach_and_log_err(cx);
559 });
560 }
561 });
562
563 Self {
564 state,
565 compact,
566 sign_in_callback,
567 }
568 }
569}
570
571impl Render for ConfigurationView {
572 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
573 let state = self.state.read(cx);
574 let user_store = state.user_store.read(cx);
575
576 let is_zed_model_provider_enabled = user_store
577 .current_organization_configuration()
578 .map_or(true, |config| config.is_zed_model_provider_enabled);
579
580 ZedAiConfiguration {
581 is_connected: !state.is_signed_out(cx),
582 plan: user_store.plan(),
583 is_zed_model_provider_enabled,
584 eligible_for_trial: user_store.trial_started_at().is_none(),
585 account_too_young: user_store.account_too_young(),
586 compact: self.compact,
587 sign_in_callback: self.sign_in_callback.clone(),
588 }
589 }
590}
591
592#[cfg(test)]
593mod tests {
594 use super::*;
595 use client::{Credentials, test::make_get_authenticated_user_response};
596 use clock::FakeSystemClock;
597 use feature_flags::FeatureFlagAppExt as _;
598 use gpui::TestAppContext;
599 use http_client::{FakeHttpClient, Method, Response};
600 use std::sync::{
601 Arc, Mutex,
602 atomic::{AtomicUsize, Ordering},
603 };
604
605 const TEST_USER_ID: u64 = 42;
606
607 fn init_test(cx: &mut App) -> (Arc<Client>, Entity<UserStore>, CloudLanguageModelProvider) {
608 let settings_store = SettingsStore::test(cx);
609 cx.set_global(settings_store);
610 cx.set_global(db::AppDatabase::test_new());
611 let app_version = AppVersion::global(cx);
612 release_channel::init_test(app_version, release_channel::ReleaseChannel::Dev, cx);
613 gpui_tokio::init(cx);
614 cx.update_flags(false, Vec::new());
615
616 let client = Client::new(
617 Arc::new(FakeSystemClock::new()),
618 FakeHttpClient::with_404_response(),
619 cx,
620 );
621 let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
622 RefreshLlmTokenListener::register(client.clone(), user_store.clone(), cx);
623 let provider = CloudLanguageModelProvider::new(user_store.clone(), client.clone(), cx);
624
625 (client, user_store, provider)
626 }
627
628 fn override_authenticate(
629 client: &Arc<Client>,
630 authenticate_rx: futures::channel::oneshot::Receiver<anyhow::Result<Credentials>>,
631 ) {
632 let authenticate_rx = Arc::new(Mutex::new(Some(authenticate_rx)));
633 client.override_authenticate(move |cx| {
634 let authenticate_rx = authenticate_rx.clone();
635 cx.background_spawn(async move {
636 let authenticate_rx = authenticate_rx
637 .lock()
638 .expect("authenticate receiver lock poisoned")
639 .take()
640 .expect("authenticate receiver already used");
641 authenticate_rx.await?
642 })
643 });
644 }
645
646 fn respond_to_authenticated_user_after(
647 client: &Arc<Client>,
648 authenticated_user_rx: futures::channel::oneshot::Receiver<()>,
649 ) {
650 let authenticated_user_rx = Arc::new(Mutex::new(Some(authenticated_user_rx)));
651 client
652 .http_client()
653 .as_fake()
654 .replace_handler(move |old_handler, request| {
655 let authenticated_user_rx = authenticated_user_rx.clone();
656 async move {
657 if request.method() == Method::GET && request.uri().path() == "/client/users/me"
658 {
659 let authenticated_user_rx = authenticated_user_rx
660 .lock()
661 .expect("authenticated user receiver lock poisoned")
662 .take();
663 if let Some(authenticated_user_rx) = authenticated_user_rx {
664 authenticated_user_rx.await.ok();
665 }
666
667 return Ok(Response::builder()
668 .status(200)
669 .body(
670 serde_json::to_string(&make_get_authenticated_user_response(
671 TEST_USER_ID as i32,
672 format!("user-{TEST_USER_ID}"),
673 ))
674 .expect("failed to serialize authenticated user response")
675 .into(),
676 )
677 .expect("failed to build authenticated user response"));
678 }
679
680 old_handler(request).await
681 }
682 });
683 }
684
685 async fn sign_in_until_authenticating(
686 client: Arc<Client>,
687 cx: &mut TestAppContext,
688 ) -> Task<anyhow::Result<Credentials>> {
689 let mut status = client.status();
690 let sign_in_task = cx.update(|cx| {
691 cx.spawn({
692 let client = client.clone();
693 async move |cx| client.sign_in(false, cx).await
694 })
695 });
696
697 while !status.borrow().is_signing_in() {
698 status.next().await;
699 }
700
701 sign_in_task
702 }
703
704 fn test_cloud_model(
705 model_id: cloud_llm_client::LanguageModelId,
706 ) -> cloud_llm_client::LanguageModel {
707 cloud_llm_client::LanguageModel {
708 provider: cloud_llm_client::LanguageModelProvider::Anthropic,
709 id: model_id,
710 display_name: "Test Model".to_string(),
711 is_latest: true,
712 max_token_count: 200_000,
713 max_token_count_in_max_mode: None,
714 max_output_tokens: 8_192,
715 supports_tools: true,
716 supports_images: false,
717 supports_thinking: false,
718 supports_disabling_thinking: false,
719 supports_fast_mode: false,
720 supports_server_side_compaction: false,
721 supported_effort_levels: Vec::new(),
722 supports_streaming_tools: false,
723 supports_parallel_tool_calls: false,
724 is_disabled: false,
725 disabled_reason: None,
726 }
727 }
728
729 #[gpui::test]
730 async fn provider_authenticate_does_not_start_sign_in_when_signed_out(cx: &mut TestAppContext) {
731 let (client, _user_store, provider) = cx.update(init_test);
732 let authenticate_calls = Arc::new(AtomicUsize::new(0));
733 client.override_authenticate({
734 let authenticate_calls = authenticate_calls.clone();
735 move |_| {
736 authenticate_calls.fetch_add(1, Ordering::SeqCst);
737 Task::ready(Err(anyhow::anyhow!(
738 "provider authenticate should not start sign-in"
739 )))
740 }
741 });
742
743 assert!(!cx.read(|cx| provider.is_authenticated(cx)));
744 assert!(matches!(
745 *client.status().borrow(),
746 client::Status::SignedOut
747 ));
748
749 cx.update(|cx| provider.authenticate(cx))
750 .now_or_never()
751 .expect("authenticate should return immediately when signed out")
752 .expect("authenticate should not fail when no sign-in is in progress");
753 cx.executor().run_until_parked();
754
755 assert_eq!(authenticate_calls.load(Ordering::SeqCst), 0);
756 assert!(matches!(
757 *client.status().borrow(),
758 client::Status::SignedOut
759 ));
760 assert!(!cx.read(|cx| provider.is_authenticated(cx)));
761 }
762
763 #[gpui::test]
764 async fn provider_authenticate_waits_for_current_user(cx: &mut TestAppContext) {
765 let (client, _user_store, provider) = cx.update(init_test);
766 let (authenticate_tx, authenticate_rx) = futures::channel::oneshot::channel();
767 let (authenticated_user_tx, authenticated_user_rx) = futures::channel::oneshot::channel();
768 override_authenticate(&client, authenticate_rx);
769 respond_to_authenticated_user_after(&client, authenticated_user_rx);
770
771 let sign_in_task = sign_in_until_authenticating(client.clone(), cx).await;
772 let authenticate_task = cx.update(|cx| provider.authenticate(cx));
773 authenticate_tx
774 .send(Ok(Credentials {
775 user_id: TEST_USER_ID,
776 access_token: "token".to_string(),
777 }))
778 .expect("authenticate receiver dropped");
779
780 cx.executor().run_until_parked();
781 assert!(!cx.read(|cx| provider.is_authenticated(cx)));
782
783 authenticated_user_tx
784 .send(())
785 .expect("authenticated user receiver dropped");
786 sign_in_task
787 .await
788 .expect("sign-in should complete after user response");
789 authenticate_task
790 .await
791 .expect("provider authentication should complete after current user is populated");
792 assert!(cx.read(|cx| provider.is_authenticated(cx)));
793
794 cx.update(|cx| provider.authenticate(cx))
795 .now_or_never()
796 .expect("already-authenticated provider should authenticate immediately")
797 .unwrap();
798 }
799
800 #[gpui::test]
801 async fn provider_authenticate_returns_error_when_sign_in_fails(cx: &mut TestAppContext) {
802 let (client, _user_store, provider) = cx.update(init_test);
803 let (authenticate_tx, authenticate_rx) = futures::channel::oneshot::channel();
804 override_authenticate(&client, authenticate_rx);
805
806 let sign_in_task = sign_in_until_authenticating(client.clone(), cx).await;
807 let authenticate_task = cx.update(|cx| provider.authenticate(cx));
808 authenticate_tx
809 .send(Err(anyhow::anyhow!("test authentication failed")))
810 .expect("authenticate receiver dropped");
811
812 sign_in_task
813 .await
814 .expect_err("sign-in should report authentication failure");
815 let error = authenticate_task
816 .await
817 .expect_err("provider authentication should fail when sign-in fails");
818 assert!(error.to_string().contains("AuthenticationError"));
819 }
820
821 #[gpui::test]
822 async fn provided_models_surface_disabled_reason(cx: &mut TestAppContext) {
823 let (_client, _user_store, provider) = cx.update(init_test);
824 let model_id = cloud_llm_client::LanguageModelId(Arc::from("disabled-model"));
825 let disabled_reason = "This model is temporarily unavailable.";
826
827 cx.update(|cx| {
828 let cloud_model_provider = provider.state.read(cx).provider.clone();
829 cloud_model_provider.update(cx, |cloud_model_provider, cx| {
830 let mut model = test_cloud_model(model_id.clone());
831 model.is_disabled = true;
832 model.disabled_reason = Some(disabled_reason.to_string());
833 cloud_model_provider.update_models(cloud_llm_client::ListModelsResponse {
834 models: vec![model],
835 default_model: Some(model_id.clone()),
836 default_fast_model: None,
837 recommended_models: vec![model_id],
838 });
839 cx.notify();
840 });
841 });
842
843 let model = cx.read(|cx| {
844 provider
845 .provided_models(cx)
846 .into_iter()
847 .next()
848 .expect("disabled model should be provided")
849 });
850 assert_eq!(
851 model.is_disabled(),
852 Some(language_model::DisabledReason::new(disabled_reason))
853 );
854 }
855
856 #[gpui::test]
857 async fn sign_out_hides_cached_cloud_models(cx: &mut TestAppContext) {
858 let (client, _user_store, provider) = cx.update(init_test);
859 let (authenticate_tx, authenticate_rx) = futures::channel::oneshot::channel();
860 let (authenticated_user_tx, authenticated_user_rx) = futures::channel::oneshot::channel();
861 override_authenticate(&client, authenticate_rx);
862 respond_to_authenticated_user_after(&client, authenticated_user_rx);
863
864 let sign_in_task = sign_in_until_authenticating(client.clone(), cx).await;
865 authenticate_tx
866 .send(Ok(Credentials {
867 user_id: TEST_USER_ID,
868 access_token: "token".to_string(),
869 }))
870 .expect("authenticate receiver dropped");
871 authenticated_user_tx
872 .send(())
873 .expect("authenticated user receiver dropped");
874 sign_in_task.await.expect("sign-in should complete");
875 cx.executor().run_until_parked();
876
877 let model_id = cloud_llm_client::LanguageModelId(Arc::from("test-model"));
878 cx.update(|cx| {
879 let cloud_model_provider = provider.state.read(cx).provider.clone();
880 cloud_model_provider.update(cx, |cloud_model_provider, cx| {
881 cloud_model_provider.update_models(cloud_llm_client::ListModelsResponse {
882 models: vec![test_cloud_model(model_id.clone())],
883 default_model: Some(model_id.clone()),
884 default_fast_model: None,
885 recommended_models: vec![model_id],
886 });
887 cx.notify();
888 });
889 });
890
891 assert!(cx.read(|cx| provider.is_authenticated(cx)));
892 assert_eq!(cx.read(|cx| provider.provided_models(cx).len()), 1);
893 assert!(cx.read(|cx| provider.default_model(cx).is_some()));
894 assert_eq!(cx.read(|cx| provider.recommended_models(cx).len()), 1);
895
896 cx.update(|cx| {
897 cx.spawn({
898 let client = client.clone();
899 async move |cx| client.sign_out(cx).await
900 })
901 })
902 .await;
903 cx.executor().run_until_parked();
904
905 assert!(!cx.read(|cx| provider.is_authenticated(cx)));
906 assert!(cx.read(|cx| provider.provided_models(cx).is_empty()));
907 assert!(cx.read(|cx| provider.default_model(cx).is_none()));
908 assert!(cx.read(|cx| provider.recommended_models(cx).is_empty()));
909 }
910}
911
912impl Component for ZedAiConfiguration {
913 fn name() -> &'static str {
914 "AI Configuration Content"
915 }
916
917 fn sort_name() -> &'static str {
918 "AI Configuration Content"
919 }
920
921 fn scope() -> ComponentScope {
922 ComponentScope::Onboarding
923 }
924
925 fn description() -> &'static str {
926 "The configuration surface for Zed's hosted AI models, \
927 showing the user's connection status, current plan, trial eligibility, \
928 and entry points for enabling the Zed model provider."
929 }
930
931 fn preview(_window: &mut Window, _cx: &mut App) -> AnyElement {
932 struct PreviewConfiguration {
933 plan: Option<Plan>,
934 is_connected: bool,
935 is_zed_model_provider_enabled: bool,
936 eligible_for_trial: bool,
937 }
938
939 let configuration = |config: PreviewConfiguration| -> AnyElement {
940 ZedAiConfiguration {
941 is_connected: config.is_connected,
942 plan: config.plan,
943 is_zed_model_provider_enabled: config.is_zed_model_provider_enabled,
944 eligible_for_trial: config.eligible_for_trial,
945 account_too_young: false,
946 compact: false,
947 sign_in_callback: Arc::new(|_, _| {}),
948 }
949 .into_any_element()
950 };
951
952 v_flex()
953 .p_4()
954 .gap_4()
955 .children(vec![
956 single_example(
957 "Not connected",
958 configuration(PreviewConfiguration {
959 plan: None,
960 is_connected: false,
961 is_zed_model_provider_enabled: true,
962 eligible_for_trial: false,
963 }),
964 ),
965 single_example(
966 "Accept Terms of Service",
967 configuration(PreviewConfiguration {
968 plan: None,
969 is_connected: true,
970 is_zed_model_provider_enabled: true,
971 eligible_for_trial: true,
972 }),
973 ),
974 single_example(
975 "No Plan - Not eligible for trial",
976 configuration(PreviewConfiguration {
977 plan: None,
978 is_connected: true,
979 is_zed_model_provider_enabled: true,
980 eligible_for_trial: false,
981 }),
982 ),
983 single_example(
984 "No Plan - Eligible for trial",
985 configuration(PreviewConfiguration {
986 plan: None,
987 is_connected: true,
988 is_zed_model_provider_enabled: true,
989 eligible_for_trial: true,
990 }),
991 ),
992 single_example(
993 "Free Plan",
994 configuration(PreviewConfiguration {
995 plan: Some(Plan::ZedFree),
996 is_connected: true,
997 is_zed_model_provider_enabled: true,
998 eligible_for_trial: true,
999 }),
1000 ),
1001 single_example(
1002 "Zed Pro Trial Plan",
1003 configuration(PreviewConfiguration {
1004 plan: Some(Plan::ZedProTrial),
1005 is_connected: true,
1006 is_zed_model_provider_enabled: true,
1007 eligible_for_trial: true,
1008 }),
1009 ),
1010 single_example(
1011 "Zed Pro Plan",
1012 configuration(PreviewConfiguration {
1013 plan: Some(Plan::ZedPro),
1014 is_connected: true,
1015 is_zed_model_provider_enabled: true,
1016 eligible_for_trial: true,
1017 }),
1018 ),
1019 single_example(
1020 "Business Plan - Zed models enabled",
1021 configuration(PreviewConfiguration {
1022 plan: Some(Plan::ZedBusiness),
1023 is_connected: true,
1024 is_zed_model_provider_enabled: true,
1025 eligible_for_trial: false,
1026 }),
1027 ),
1028 single_example(
1029 "Business Plan - Zed models disabled",
1030 configuration(PreviewConfiguration {
1031 plan: Some(Plan::ZedBusiness),
1032 is_connected: true,
1033 is_zed_model_provider_enabled: false,
1034 eligible_for_trial: false,
1035 }),
1036 ),
1037 ])
1038 .into_any_element()
1039 }
1040}
1041