tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

fix: persist accumulated cost in session DB to survive reload (#9191)

Browse Source 
Signed-off-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
Douwe Osinga
2026-05-14 12:57:06 -04:00
committed by GitHub
parent c938a46417
commit 4671053985
13 changed files with 100 additions and 192 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/reply.rs
+1
View File
@@ -164,6 +164,7 @@ pub async fn get_token_state(session_manager: &SessionManager, session_id: &str)
accumulated_input_tokens: session.accumulated_input_tokens.unwrap_or(0),
accumulated_output_tokens: session.accumulated_output_tokens.unwrap_or(0),
accumulated_total_tokens: session.accumulated_total_tokens.unwrap_or(0),
accumulated_cost: session.accumulated_cost,
})
.inspect_err(|e| {
tracing::warn!(
crates/goose/src/acp/server.rs
+1
View File
@@ -4293,6 +4293,7 @@ print(\"hello, world\")
accumulated_total_tokens,
accumulated_input_tokens,
accumulated_output_tokens,
accumulated_cost: None,
schedule_id: None,
recipe: None,
user_recipe_values: None,
crates/goose/src/agents/reply_parts.rs
+27
View File
@@ -527,6 +527,12 @@ impl Agent {
let accumulated_output =
accumulate(session.accumulated_output_tokens, usage.usage.output_tokens);
let accumulated_cost = session
.provider_name
.as_deref()
.and_then(|pn| self.accumulate_cost(session.accumulated_cost, usage, pn))
.or(session.accumulated_cost);
let (current_total, current_input, current_output) = if is_compaction_usage {
// After compaction: summary output becomes new input context
let new_input = usage.usage.output_tokens;
@@ -548,11 +554,32 @@ impl Agent {
.accumulated_total_tokens(accumulated_total)
.accumulated_input_tokens(accumulated_input)
.accumulated_output_tokens(accumulated_output)
.accumulated_cost(accumulated_cost)
.apply()
.await?;
Ok(())
}
fn accumulate_cost(
&self,
existing: Option<f64>,
usage: &ProviderUsage,
provider_name: &str,
) -> Option<f64> {
let canonical =
crate::providers::canonical::maybe_get_canonical_model(provider_name, &usage.model)?;
let input_price = canonical.cost.input?;
let output_price = canonical.cost.output?;
let input_tokens = usage.usage.input_tokens.unwrap_or(0) as f64;
let output_tokens = usage.usage.output_tokens.unwrap_or(0) as f64;
let chunk_cost = (input_tokens * input_price + output_tokens * output_price) / 1_000_000.0;
Some(existing.unwrap_or(0.0) + chunk_cost)
}
}
/// Check whether a tool should be callable by an app based on MCP Apps visibility metadata.
crates/goose/src/conversation/message.rs
+1
View File
@@ -1026,6 +1026,7 @@ pub struct TokenState {
pub accumulated_input_tokens: i32,
pub accumulated_output_tokens: i32,
pub accumulated_total_tokens: i32,
pub accumulated_cost: Option<f64>,
}
#[cfg(test)]
crates/goose/src/session/session_manager.rs
+35 -2
View File
@@ -19,7 +19,7 @@ use std::sync::{Arc, LazyLock};
use tracing::{info, warn};
use utoipa::ToSchema;
pub const CURRENT_SCHEMA_VERSION: i32 = 12;
pub const CURRENT_SCHEMA_VERSION: i32 = 13;
pub const SESSIONS_FOLDER: &str = "sessions";
pub const DB_NAME: &str = "sessions.db";
@@ -72,6 +72,7 @@ pub struct Session {
pub accumulated_total_tokens: Option<i32>,
pub accumulated_input_tokens: Option<i32>,
pub accumulated_output_tokens: Option<i32>,
pub accumulated_cost: Option<f64>,
pub schedule_id: Option<String>,
pub recipe: Option<Recipe>,
pub user_recipe_values: Option<HashMap<String, String>>,
@@ -101,6 +102,7 @@ pub struct SessionUpdateBuilder<'a> {
accumulated_total_tokens: Option<Option<i32>>,
accumulated_input_tokens: Option<Option<i32>>,
accumulated_output_tokens: Option<Option<i32>>,
accumulated_cost: Option<Option<f64>>,
schedule_id: Option<Option<String>>,
recipe: Option<Option<Recipe>>,
user_recipe_values: Option<Option<HashMap<String, String>>>,
@@ -135,6 +137,7 @@ impl<'a> SessionUpdateBuilder<'a> {
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
accumulated_cost: None,
schedule_id: None,
recipe: None,
user_recipe_values: None,
@@ -213,6 +216,11 @@ impl<'a> SessionUpdateBuilder<'a> {
self
}
pub fn accumulated_cost(mut self, cost: Option<f64>) -> Self {
self.accumulated_cost = Some(cost);
self
}
pub fn schedule_id(mut self, schedule_id: Option<String>) -> Self {
self.schedule_id = Some(schedule_id);
self
@@ -490,6 +498,7 @@ impl Default for Session {
accumulated_total_tokens: None,
accumulated_input_tokens: None,
accumulated_output_tokens: None,
accumulated_cost: None,
schedule_id: None,
recipe: None,
user_recipe_values: None,
@@ -557,6 +566,7 @@ impl sqlx::FromRow<'_, sqlx::sqlite::SqliteRow> for Session {
accumulated_total_tokens: row.try_get("accumulated_total_tokens")?,
accumulated_input_tokens: row.try_get("accumulated_input_tokens")?,
accumulated_output_tokens: row.try_get("accumulated_output_tokens")?,
accumulated_cost: row.try_get("accumulated_cost").ok().flatten(),
schedule_id: row.try_get("schedule_id")?,
recipe,
user_recipe_values,
@@ -666,6 +676,7 @@ impl SessionStorage {
accumulated_total_tokens INTEGER,
accumulated_input_tokens INTEGER,
accumulated_output_tokens INTEGER,
accumulated_cost REAL,
schedule_id TEXT,
recipe_json TEXT,
user_recipe_values_json TEXT,
@@ -786,9 +797,10 @@ impl SessionStorage {
id, name, user_set_name, session_type, working_dir, created_at, updated_at, extension_data,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
accumulated_cost,
schedule_id, recipe_json, user_recipe_values_json,
provider_name, model_config_json, goose_mode
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&session.id)
@@ -805,6 +817,7 @@ impl SessionStorage {
.bind(session.accumulated_total_tokens)
.bind(session.accumulated_input_tokens)
.bind(session.accumulated_output_tokens)
.bind(session.accumulated_cost)
.bind(&session.schedule_id)
.bind(recipe_json)
.bind(user_recipe_values_json)
@@ -1087,6 +1100,19 @@ impl SessionStorage {
.await?;
}
}
13 => {
let has_accumulated_cost = sqlx::query_scalar::<_, i32>(
"SELECT COUNT(*) FROM pragma_table_info('sessions') WHERE name = 'accumulated_cost'",
)
.fetch_one(&mut **tx)
.await?
> 0;
if !has_accumulated_cost {
sqlx::query("ALTER TABLE sessions ADD COLUMN accumulated_cost REAL")
.execute(&mut **tx)
.await?;
}
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}
@@ -1147,6 +1173,7 @@ impl SessionStorage {
SELECT id, working_dir, name, description, user_set_name, session_type, created_at, updated_at, extension_data,
total_tokens, input_tokens, output_tokens,
accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens,
accumulated_cost,
schedule_id, recipe_json, user_recipe_values_json,
provider_name, model_config_json, goose_mode,
archived_at, project_id
@@ -1207,6 +1234,7 @@ impl SessionStorage {
builder.accumulated_output_tokens,
"accumulated_output_tokens"
);
add_update!(builder.accumulated_cost, "accumulated_cost");
add_update!(builder.schedule_id, "schedule_id");
add_update!(builder.recipe, "recipe_json");
add_update!(builder.user_recipe_values, "user_recipe_values_json");
@@ -1259,6 +1287,9 @@ impl SessionStorage {
if let Some(aot) = builder.accumulated_output_tokens {
q = q.bind(aot);
}
if let Some(ac) = builder.accumulated_cost {
q = q.bind(ac);
}
if let Some(sid) = builder.schedule_id {
q = q.bind(sid);
}
@@ -1445,6 +1476,7 @@ impl SessionStorage {
SELECT s.id, s.working_dir, s.name, s.description, s.user_set_name, s.session_type, s.created_at, s.updated_at, s.extension_data,
s.total_tokens, s.input_tokens, s.output_tokens,
s.accumulated_total_tokens, s.accumulated_input_tokens, s.accumulated_output_tokens,
s.accumulated_cost,
s.schedule_id, s.recipe_json, s.user_recipe_values_json,
s.provider_name, s.model_config_json, s.goose_mode,
s.archived_at, s.project_id,
@@ -1564,6 +1596,7 @@ impl SessionStorage {
.accumulated_total_tokens(import.accumulated_total_tokens)
.accumulated_input_tokens(import.accumulated_input_tokens)
.accumulated_output_tokens(import.accumulated_output_tokens)
.accumulated_cost(import.accumulated_cost)
.schedule_id(import.schedule_id)
.recipe(import.recipe)
.user_recipe_values(import.user_recipe_values);