assistant2: Serialize token usage (#27586)

We'll need this for detecting old long threads

Release Notes:

- N/A
This commit is contained in:
Agus Zubiaga 2025-03-27 10:38:08 -03:00 committed by GitHub
parent 76d3a9a0f0
commit f15a241d3e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 6 additions and 3 deletions

View file

@ -286,8 +286,7 @@ impl Thread {
tool_use, tool_use,
action_log: cx.new(|_| ActionLog::new()), action_log: cx.new(|_| ActionLog::new()),
initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(), initial_project_snapshot: Task::ready(serialized.initial_project_snapshot).shared(),
// TODO: persist token usage? cumulative_token_usage: serialized.cumulative_token_usage,
cumulative_token_usage: TokenUsage::default(),
feedback: None, feedback: None,
} }
} }
@ -648,6 +647,7 @@ impl Thread {
}) })
.collect(), .collect(),
initial_project_snapshot, initial_project_snapshot,
cumulative_token_usage: this.cumulative_token_usage.clone(),
}) })
}) })
} }

View file

@ -16,7 +16,7 @@ use gpui::{
}; };
use heed::types::SerdeBincode; use heed::types::SerdeBincode;
use heed::Database; use heed::Database;
use language_model::{LanguageModelToolUseId, Role}; use language_model::{LanguageModelToolUseId, Role, TokenUsage};
use project::Project; use project::Project;
use prompt_store::PromptBuilder; use prompt_store::PromptBuilder;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
@ -308,6 +308,8 @@ pub struct SerializedThread {
pub messages: Vec<SerializedMessage>, pub messages: Vec<SerializedMessage>,
#[serde(default)] #[serde(default)]
pub initial_project_snapshot: Option<Arc<ProjectSnapshot>>, pub initial_project_snapshot: Option<Arc<ProjectSnapshot>>,
#[serde(default)]
pub cumulative_token_usage: TokenUsage,
} }
impl SerializedThread { impl SerializedThread {
@ -390,6 +392,7 @@ impl LegacySerializedThread {
updated_at: self.updated_at, updated_at: self.updated_at,
messages: self.messages.into_iter().map(|msg| msg.upgrade()).collect(), messages: self.messages.into_iter().map(|msg| msg.upgrade()).collect(),
initial_project_snapshot: self.initial_project_snapshot, initial_project_snapshot: self.initial_project_snapshot,
cumulative_token_usage: TokenUsage::default(),
} }
} }
} }