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:
Michael Sloan 2025-06-20 15:28:48 -06:00 committed by GitHub
parent e0c0b6f95d
commit 7e801dccb0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 188 additions and 211 deletions

View file

@ -8,27 +8,22 @@ mod telemetry;
#[cfg(any(test, feature = "test-support"))]
pub mod fake_provider;
use anyhow::{Context as _, Result};
use anyhow::Result;
use client::Client;
use futures::FutureExt;
use futures::{StreamExt, future::BoxFuture, stream::BoxStream};
use gpui::{AnyElement, AnyView, App, AsyncApp, SharedString, Task, Window};
use http_client::http::{HeaderMap, HeaderValue};
use icons::IconName;
use parking_lot::Mutex;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize, de::DeserializeOwned};
use std::fmt;
use std::ops::{Add, Sub};
use std::str::FromStr as _;
use std::sync::Arc;
use std::time::Duration;
use thiserror::Error;
use util::serde::is_default;
use zed_llm_client::{
CompletionRequestStatus, MODEL_REQUESTS_USAGE_AMOUNT_HEADER_NAME,
MODEL_REQUESTS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
};
use zed_llm_client::CompletionRequestStatus;
pub use crate::model::*;
pub use crate::rate_limiter::*;
@ -106,32 +101,6 @@ pub enum StopReason {
Refusal,
}
#[derive(Debug, Clone, Copy)]
pub struct RequestUsage {
pub limit: UsageLimit,
pub amount: i32,
}
impl RequestUsage {
pub fn from_headers(headers: &HeaderMap<HeaderValue>) -> Result<Self> {
let limit = headers
.get(MODEL_REQUESTS_USAGE_LIMIT_HEADER_NAME)
.with_context(|| {
format!("missing {MODEL_REQUESTS_USAGE_LIMIT_HEADER_NAME:?} header")
})?;
let limit = UsageLimit::from_str(limit.to_str()?)?;
let amount = headers
.get(MODEL_REQUESTS_USAGE_AMOUNT_HEADER_NAME)
.with_context(|| {
format!("missing {MODEL_REQUESTS_USAGE_AMOUNT_HEADER_NAME:?} header")
})?;
let amount = amount.to_str()?.parse::<i32>()?;
Ok(Self { limit, amount })
}
}
#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize, Default)]
pub struct TokenUsage {
#[serde(default, skip_serializing_if = "is_default")]