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

Fix update of prompt usage count when using text threads (#35671)

Browse source 
Release Notes:

- Fixed update of prompt usage count when using agent text threads.

Co-authored-by: Oleksiy <oleksiy@zed.dev>
This commit is contained in:
Michael Sloan 2025-08-05 15:39:34 -06:00 • committed by GitHub
parent dd7fce3f5e
commit 42699411f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 29 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

32
crates/assistant_context/src/assistant_context.rs
View file

@ -9,9 +9,9 @@ use assistant_slash_command::{
SlashCommandResult, SlashCommandWorkingSet,
};
use assistant_slash_commands::FileCommandMetadata;
use client::{self, Client, proto, telemetry::Telemetry};
use client::{self, Client, ModelRequestUsage, RequestUsage, proto, telemetry::Telemetry};
use clock::ReplicaId;
use cloud_llm_client::CompletionIntent;
use cloud_llm_client::{CompletionIntent, CompletionRequestStatus, UsageLimit};
use collections::{HashMap, HashSet};
use fs::{Fs, RenameOptions};
use futures::{FutureExt, StreamExt, future::Shared};
@ -2080,7 +2080,18 @@ impl AssistantContext {
});
match event {
LanguageModelCompletionEvent::StatusUpdate { .. } => {}
LanguageModelCompletionEvent::StatusUpdate(status_update) => {
match status_update {
CompletionRequestStatus::UsageUpdated { amount, limit } => {
this.update_model_request_usage(
amount as u32,
limit,
cx,
);
}
_ => {}
}
}
LanguageModelCompletionEvent::StartMessage { .. } => {}
LanguageModelCompletionEvent::Stop(reason) => {
stop_reason = reason;
@ -2956,6 +2967,21 @@ impl AssistantContext {
summary.text = custom_summary;
cx.emit(ContextEvent::SummaryChanged);
}
fn update_model_request_usage(&self, amount: u32, limit: UsageLimit, cx: &mut App) {
let Some(project) = &self.project else {
return;
};
project.read(cx).user_store().update(cx, |user_store, cx| {
user_store.update_model_request_usage(
ModelRequestUsage(RequestUsage {
amount: amount as i32,
limit,
}),
cx,
)
});
}
}
#[derive(Debug, Default)]