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

agent: Load usage eagerly (#29937)

Browse source 
Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <git@maxdeviant.com>
This commit is contained in:
Nate Butler 2025-05-05 14:42:21 -04:00 • committed by GitHub
parent 181cd6294f
commit e64f5ff358
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 89 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

37
crates/agent/src/message_editor.rs
View file

@ -26,7 +26,7 @@ use gpui::{
Task, TextStyle, WeakEntity, linear_color_stop, linear_gradient, point, pulsating_between,
};
use language::{Buffer, Language};
use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent};
use language_model::{ConfiguredModel, LanguageModelRequestMessage, MessageContent, RequestUsage};
use language_model_selector::ToggleModelSelector;
use multi_buffer;
use project::Project;
@ -36,7 +36,7 @@ use settings::Settings;
use std::time::Duration;
use theme::ThemeSettings;
use ui::{Disclosure, KeyBinding, PopoverMenuHandle, Tooltip, prelude::*};
use util::ResultExt as _;
use util::{ResultExt as _, maybe};
use workspace::{CollaboratorId, Workspace};
use zed_llm_client::CompletionMode;
@ -1043,9 +1043,17 @@ impl MessageEditor {
return None;
}
let plan = self
.user_store
.read(cx)
let user_store = self.user_store.read(cx);
let ubb_enable = user_store
.usage_based_billing_enabled()
.map_or(false, |enabled| enabled);
if ubb_enable {
return None;
}
let plan = user_store
.current_plan()
.map(|plan| match plan {
Plan::Free => zed_llm_client::Plan::Free,
@ -1053,7 +1061,24 @@ impl MessageEditor {
Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
})
.unwrap_or(zed_llm_client::Plan::Free);
let usage = self.thread.read(cx).last_usage()?;
let usage = self.thread.read(cx).last_usage().or_else(|| {
maybe!({
let amount = user_store.model_request_usage_amount()?;
let limit = user_store.model_request_usage_limit()?.variant?;
Some(RequestUsage {
amount: amount as i32,
limit: match limit {
proto::usage_limit::Variant::Limited(limited) => {
zed_llm_client::UsageLimit::Limited(limited.limit as i32)
}
proto::usage_limit::Variant::Unlimited(_) => {
zed_llm_client::UsageLimit::Unlimited
}
},
})
})
})?;
Some(
div()