assistant2: Serialize token usage (#27586)

Agus Zubiaga created 9 months ago

We'll need this for detecting old long threads

Release Notes:

- N/A

Change summary

crates/assistant2/src/thread.rs       | 4 ++--
crates/assistant2/src/thread_store.rs | 5 ++++-
2 files changed, 6 insertions(+), 3 deletions(-)

Detailed changes

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(),
             })
         })
     }

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<SerializedMessage>,
     #[serde(default)]
     pub initial_project_snapshot: Option<Arc<ProjectSnapshot>>,
+    #[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(),
         }
     }
 }