agent: Fix issues with usage display sometimes showing initially fetched usage (#33125)
Having `Thread::last_usage` as an override of the initially fetched usage could cause the initial usage to be displayed when the current thread is empty or in text threads. Fix is to just store last usage info in `UserStore` and not have these overrides Release Notes: - Agent: Fixed request usage display to always include the most recently known usage - there were some cases where it would show the initially requested usage.
This commit is contained in:
parent
e0c0b6f95d
commit
7e801dccb0
13 changed files with 188 additions and 211 deletions
|
@ -7,6 +7,7 @@ use agent_settings::{AgentProfileId, AgentSettings, CompletionMode};
|
|||
use anyhow::{Result, anyhow};
|
||||
use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
|
||||
use chrono::{DateTime, Utc};
|
||||
use client::{ModelRequestUsage, RequestUsage};
|
||||
use collections::HashMap;
|
||||
use editor::display_map::CreaseMetadata;
|
||||
use feature_flags::{self, FeatureFlagAppExt};
|
||||
|
@ -22,8 +23,8 @@ use language_model::{
|
|||
LanguageModelId, LanguageModelKnownError, LanguageModelRegistry, LanguageModelRequest,
|
||||
LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
|
||||
LanguageModelToolResultContent, LanguageModelToolUseId, MessageContent,
|
||||
ModelRequestLimitReachedError, PaymentRequiredError, RequestUsage, Role, SelectedModel,
|
||||
StopReason, TokenUsage,
|
||||
ModelRequestLimitReachedError, PaymentRequiredError, Role, SelectedModel, StopReason,
|
||||
TokenUsage,
|
||||
};
|
||||
use postage::stream::Stream as _;
|
||||
use project::Project;
|
||||
|
@ -38,7 +39,7 @@ use ui::Window;
|
|||
use util::{ResultExt as _, post_inc};
|
||||
|
||||
use uuid::Uuid;
|
||||
use zed_llm_client::{CompletionIntent, CompletionRequestStatus};
|
||||
use zed_llm_client::{CompletionIntent, CompletionRequestStatus, UsageLimit};
|
||||
|
||||
use crate::ThreadStore;
|
||||
use crate::agent_profile::AgentProfile;
|
||||
|
@ -350,7 +351,6 @@ pub struct Thread {
|
|||
request_token_usage: Vec<TokenUsage>,
|
||||
cumulative_token_usage: TokenUsage,
|
||||
exceeded_window_error: Option<ExceededWindowError>,
|
||||
last_usage: Option<RequestUsage>,
|
||||
tool_use_limit_reached: bool,
|
||||
feedback: Option<ThreadFeedback>,
|
||||
message_feedback: HashMap<MessageId, ThreadFeedback>,
|
||||
|
@ -443,7 +443,6 @@ impl Thread {
|
|||
request_token_usage: Vec::new(),
|
||||
cumulative_token_usage: TokenUsage::default(),
|
||||
exceeded_window_error: None,
|
||||
last_usage: None,
|
||||
tool_use_limit_reached: false,
|
||||
feedback: None,
|
||||
message_feedback: HashMap::default(),
|
||||
|
@ -568,7 +567,6 @@ impl Thread {
|
|||
request_token_usage: serialized.request_token_usage,
|
||||
cumulative_token_usage: serialized.cumulative_token_usage,
|
||||
exceeded_window_error: None,
|
||||
last_usage: None,
|
||||
tool_use_limit_reached: serialized.tool_use_limit_reached,
|
||||
feedback: None,
|
||||
message_feedback: HashMap::default(),
|
||||
|
@ -875,10 +873,6 @@ impl Thread {
|
|||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
pub fn last_usage(&self) -> Option<RequestUsage> {
|
||||
self.last_usage
|
||||
}
|
||||
|
||||
pub fn tool_use_limit_reached(&self) -> bool {
|
||||
self.tool_use_limit_reached
|
||||
}
|
||||
|
@ -1658,9 +1652,7 @@ impl Thread {
|
|||
CompletionRequestStatus::UsageUpdated {
|
||||
amount, limit
|
||||
} => {
|
||||
let usage = RequestUsage { limit, amount: amount as i32 };
|
||||
|
||||
thread.last_usage = Some(usage);
|
||||
thread.update_model_request_usage(amount as u32, limit, cx);
|
||||
}
|
||||
CompletionRequestStatus::ToolUseLimitReached => {
|
||||
thread.tool_use_limit_reached = true;
|
||||
|
@ -1871,11 +1863,8 @@ impl Thread {
|
|||
LanguageModelCompletionEvent::StatusUpdate(
|
||||
CompletionRequestStatus::UsageUpdated { amount, limit },
|
||||
) => {
|
||||
this.update(cx, |thread, _cx| {
|
||||
thread.last_usage = Some(RequestUsage {
|
||||
limit,
|
||||
amount: amount as i32,
|
||||
});
|
||||
this.update(cx, |thread, cx| {
|
||||
thread.update_model_request_usage(amount as u32, limit, cx);
|
||||
})?;
|
||||
continue;
|
||||
}
|
||||
|
@ -2757,6 +2746,20 @@ impl Thread {
|
|||
}
|
||||
}
|
||||
|
||||
fn update_model_request_usage(&self, amount: u32, limit: UsageLimit, cx: &mut Context<Self>) {
|
||||
self.project.update(cx, |project, cx| {
|
||||
project.user_store().update(cx, |user_store, cx| {
|
||||
user_store.update_model_request_usage(
|
||||
ModelRequestUsage(RequestUsage {
|
||||
amount: amount as i32,
|
||||
limit,
|
||||
}),
|
||||
cx,
|
||||
)
|
||||
})
|
||||
});
|
||||
}
|
||||
|
||||
pub fn deny_tool_use(
|
||||
&mut self,
|
||||
tool_use_id: LanguageModelToolUseId,
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue