diff --git a/crates/assistant2/src/thread.rs b/crates/assistant2/src/thread.rs index 5facf752b4..c75ead9dd3 100644 --- a/crates/assistant2/src/thread.rs +++ b/crates/assistant2/src/thread.rs @@ -286,8 +286,7 @@ impl Thread { tool_use, action_log: cx.new(|_| ActionLog::new()), initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(), - // TODO: persist token usage? - cumulative_token_usage: TokenUsage::default(), + cumulative_token_usage: serialized.cumulative_token_usage, feedback: None, } } @@ -648,6 +647,7 @@ impl Thread { }) .collect(), initial_project_snapshot, + cumulative_token_usage: this.cumulative_token_usage.clone(), }) }) } diff --git a/crates/assistant2/src/thread_store.rs b/crates/assistant2/src/thread_store.rs index aa2e0aa971..474db2bcd8 100644 --- a/crates/assistant2/src/thread_store.rs +++ b/crates/assistant2/src/thread_store.rs @@ -16,7 +16,7 @@ use gpui::{ }; use heed::types::SerdeBincode; use heed::Database; -use language_model::{LanguageModelToolUseId, Role}; +use language_model::{LanguageModelToolUseId, Role, TokenUsage}; use project::Project; use prompt_store::PromptBuilder; use serde::{Deserialize, Serialize}; @@ -308,6 +308,8 @@ pub struct SerializedThread { pub messages: Vec, #[serde(default)] pub initial_project_snapshot: Option>, + #[serde(default)] + pub cumulative_token_usage: TokenUsage, } impl SerializedThread { @@ -390,6 +392,7 @@ impl LegacySerializedThread { updated_at: self.updated_at, messages: self.messages.into_iter().map(|msg| msg.upgrade()).collect(), initial_project_snapshot: self.initial_project_snapshot, + cumulative_token_usage: TokenUsage::default(), } } }