Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update Agent panel to work with CloudUserStore (#35436)

Browse source 
This PR updates the Agent panel to work with the `CloudUserStore`
instead of the `UserStore`, reducing its reliance on being connected to
Collab to function.

Release Notes:

- N/A

---------

Co-authored-by: Richard Feldman <oss@rtfeldman.com>
This commit is contained in:
Marshall Bowers 2025-07-31 21:44:43 -04:00 • committed by GitHub
parent 09b93caa9b
commit 72d354de6c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 212 additions and 108 deletions
Download patch file Download diff file Expand all files Collapse all files

25
crates/language_models/src/language_models.rs
View file

@ -1,7 +1,7 @@
use std::sync::Arc;
use ::settings::{Settings, SettingsStore};
use client::{Client, UserStore};
use client::{Client, CloudUserStore, UserStore};
use collections::HashSet;
use gpui::{App, Context, Entity};
use language_model::{LanguageModelProviderId, LanguageModelRegistry};
@ -26,11 +26,22 @@ use crate::provider::vercel::VercelLanguageModelProvider;
use crate::provider::x_ai::XAiLanguageModelProvider;
pub use crate::settings::*;
pub fn init(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) {
pub fn init(
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
client: Arc<Client>,
cx: &mut App,
) {
crate::settings::init_settings(cx);
let registry = LanguageModelRegistry::global(cx);
registry.update(cx, |registry, cx| {
register_language_model_providers(registry, user_store, client.clone(), cx);
register_language_model_providers(
registry,
user_store,
cloud_user_store,
client.clone(),
cx,
);
});
let mut openai_compatible_providers = AllLanguageModelSettings::get_global(cx)
@ -100,11 +111,17 @@ fn register_openai_compatible_providers(
fn register_language_model_providers(
registry: &mut LanguageModelRegistry,
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
client: Arc<Client>,
cx: &mut Context<LanguageModelRegistry>,
) {
registry.register_provider(
CloudLanguageModelProvider::new(user_store.clone(), client.clone(), cx),
CloudLanguageModelProvider::new(
user_store.clone(),
cloud_user_store.clone(),
client.clone(),
cx,
),
cx,
);

87
crates/language_models/src/provider/cloud.rs
View file

@ -2,11 +2,11 @@ use ai_onboarding::YoungAccountBanner;
use anthropic::AnthropicModelMode;
use anyhow::{Context as _, Result, anyhow};
use chrono::{DateTime, Utc};
use client::{Client, ModelRequestUsage, UserStore, zed_urls};
use client::{Client, CloudUserStore, ModelRequestUsage, UserStore, zed_urls};
use cloud_llm_client::{
CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CURRENT_PLAN_HEADER_NAME, CompletionBody,
CompletionEvent, CompletionRequestStatus, CountTokensBody, CountTokensResponse,
EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, MODEL_REQUESTS_RESOURCE_HEADER_VALUE,
EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, MODEL_REQUESTS_RESOURCE_HEADER_VALUE, Plan,
SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME,
TOOL_USE_LIMIT_REACHED_HEADER_NAME, ZED_VERSION_HEADER_NAME,
};
@ -27,7 +27,6 @@ use language_model::{
LanguageModelToolChoice, LanguageModelToolSchemaFormat, LlmApiToken,
ModelRequestLimitReachedError, PaymentRequiredError, RateLimiter, RefreshLlmTokenListener,
};
use proto::Plan;
use release_channel::AppVersion;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
@ -118,6 +117,7 @@ pub struct State {
client: Arc<Client>,
llm_api_token: LlmApiToken,
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
status: client::Status,
accept_terms_of_service_task: Option<Task<Result<()>>>,
models: Vec<Arc<cloud_llm_client::LanguageModel>>,
@ -133,6 +133,7 @@ impl State {
fn new(
client: Arc<Client>,
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
status: client::Status,
cx: &mut Context<Self>,
) -> Self {
@ -142,6 +143,7 @@ impl State {
client: client.clone(),
llm_api_token: LlmApiToken::default(),
user_store,
cloud_user_store,
status,
accept_terms_of_service_task: None,
models: Vec::new(),
@ -150,12 +152,19 @@ impl State {
recommended_models: Vec::new(),
_fetch_models_task: cx.spawn(async move |this, cx| {
maybe!(async move {
let (client, llm_api_token) = this
.read_with(cx, |this, _cx| (client.clone(), this.llm_api_token.clone()))?;
let (client, cloud_user_store, llm_api_token) =
this.read_with(cx, |this, _cx| {
(
client.clone(),
this.cloud_user_store.clone(),
this.llm_api_token.clone(),
)
})?;
loop {
let status = this.read_with(cx, |this, _cx| this.status)?;
if matches!(status, client::Status::Connected { .. }) {
let is_authenticated =
cloud_user_store.read_with(cx, |this, _cx| this.is_authenticated())?;
if is_authenticated {
break;
}
@ -194,8 +203,8 @@ impl State {
}
}
fn is_signed_out(&self) -> bool {
self.status.is_signed_out()
fn is_signed_out(&self, cx: &App) -> bool {
!self.cloud_user_store.read(cx).is_authenticated()
}
fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
@ -210,10 +219,7 @@ impl State {
}
fn has_accepted_terms_of_service(&self, cx: &App) -> bool {
self.user_store
.read(cx)
.current_user_has_accepted_terms()
.unwrap_or(false)
self.cloud_user_store.read(cx).has_accepted_tos()
}
fn accept_terms_of_service(&mut self, cx: &mut Context<Self>) {
@ -297,11 +303,24 @@ impl State {
}
impl CloudLanguageModelProvider {
pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self {
pub fn new(
user_store: Entity<UserStore>,
cloud_user_store: Entity<CloudUserStore>,
client: Arc<Client>,
cx: &mut App,
) -> Self {
let mut status_rx = client.status();
let status = *status_rx.borrow();
let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx));
let state = cx.new(|cx| {
State::new(
client.clone(),
user_store.clone(),
cloud_user_store.clone(),
status,
cx,
)
});
let state_ref = state.downgrade();
let maintain_client_status = cx.spawn(async move |cx| {
@ -398,7 +417,7 @@ impl LanguageModelProvider for CloudLanguageModelProvider {
fn is_authenticated(&self, cx: &App) -> bool {
let state = self.state.read(cx);
!state.is_signed_out() && state.has_accepted_terms_of_service(cx)
!state.is_signed_out(cx) && state.has_accepted_terms_of_service(cx)
}
fn authenticate(&self, _cx: &mut App) -> Task<Result<(), AuthenticateError>> {
@ -614,9 +633,9 @@ impl CloudLanguageModel {
.and_then(|plan| cloud_llm_client::Plan::from_str(plan).ok())
{
let plan = match plan {
cloud_llm_client::Plan::ZedFree => Plan::Free,
cloud_llm_client::Plan::ZedPro => Plan::ZedPro,
cloud_llm_client::Plan::ZedProTrial => Plan::ZedProTrial,
cloud_llm_client::Plan::ZedFree => proto::Plan::Free,
cloud_llm_client::Plan::ZedPro => proto::Plan::ZedPro,
cloud_llm_client::Plan::ZedProTrial => proto::Plan::ZedProTrial,
};
return Err(anyhow!(ModelRequestLimitReachedError { plan }));
}
@ -1118,7 +1137,7 @@ fn response_lines<T: DeserializeOwned>(
#[derive(IntoElement, RegisterComponent)]
struct ZedAiConfiguration {
is_connected: bool,
plan: Option<proto::Plan>,
plan: Option<Plan>,
subscription_period: Option<(DateTime<Utc>, DateTime<Utc>)>,
eligible_for_trial: bool,
has_accepted_terms_of_service: bool,
@ -1132,15 +1151,15 @@ impl RenderOnce for ZedAiConfiguration {
fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
let young_account_banner = YoungAccountBanner;
let is_pro = self.plan == Some(proto::Plan::ZedPro);
let is_pro = self.plan == Some(Plan::ZedPro);
let subscription_text = match (self.plan, self.subscription_period) {
(Some(proto::Plan::ZedPro), Some(_)) => {
(Some(Plan::ZedPro), Some(_)) => {
"You have access to Zed's hosted models through your Pro subscription."
}
(Some(proto::Plan::ZedProTrial), Some(_)) => {
(Some(Plan::ZedProTrial), Some(_)) => {
"You have access to Zed's hosted models through your Pro trial."
}
(Some(proto::Plan::Free), Some(_)) => {
(Some(Plan::ZedFree), Some(_)) => {
"You have basic access to Zed's hosted models through the Free plan."
}
_ => {
@ -1262,15 +1281,15 @@ impl ConfigurationView {
impl Render for ConfigurationView {
fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
let state = self.state.read(cx);
let user_store = state.user_store.read(cx);
let cloud_user_store = state.cloud_user_store.read(cx);
ZedAiConfiguration {
is_connected: !state.is_signed_out(),
plan: user_store.current_plan(),
subscription_period: user_store.subscription_period(),
eligible_for_trial: user_store.trial_started_at().is_none(),
is_connected: !state.is_signed_out(cx),
plan: cloud_user_store.plan(),
subscription_period: cloud_user_store.subscription_period(),
eligible_for_trial: cloud_user_store.trial_started_at().is_none(),
has_accepted_terms_of_service: state.has_accepted_terms_of_service(cx),
account_too_young: user_store.account_too_young(),
account_too_young: cloud_user_store.account_too_young(),
accept_terms_of_service_in_progress: state.accept_terms_of_service_task.is_some(),
accept_terms_of_service_callback: self.accept_terms_of_service_callback.clone(),
sign_in_callback: self.sign_in_callback.clone(),
@ -1286,7 +1305,7 @@ impl Component for ZedAiConfiguration {
fn preview(_window: &mut Window, _cx: &mut App) -> Option<AnyElement> {
fn configuration(
is_connected: bool,
plan: Option<proto::Plan>,
plan: Option<Plan>,
eligible_for_trial: bool,
account_too_young: bool,
has_accepted_terms_of_service: bool,
@ -1330,15 +1349,15 @@ impl Component for ZedAiConfiguration {
),
single_example(
"Free Plan",
configuration(true, Some(proto::Plan::Free), true, false, true),
configuration(true, Some(Plan::ZedFree), true, false, true),
),
single_example(
"Zed Pro Trial Plan",
configuration(true, Some(proto::Plan::ZedProTrial), true, false, true),
configuration(true, Some(Plan::ZedProTrial), true, false, true),
),
single_example(
"Zed Pro Plan",
configuration(true, Some(proto::Plan::ZedPro), true, false, true),
configuration(true, Some(Plan::ZedPro), true, false, true),
),
])
.into_any_element(),