From 3b3e1d142f9acd18ecc05a0aced38647ff5403eb Mon Sep 17 00:00:00 2001 From: Filip Kujawa Date: Mon, 29 Jun 2026 15:36:51 -0700 Subject: [PATCH 1/2] fix(agents): roll subagent token usage and cost into parent session Subagents spawned via the delegate tool run in their own sessions, so their token usage and cost never reached the parent session - under-reporting cost-per-outcome and per-session budgets whenever delegation occurred. Fold a finished subagent's lifetime usage/cost into the parent session's accumulated_* totals at the run_subagent_task chokepoint (covers sync and async delegate), without touching the parent's usage (context window) columns so compaction triggers are unaffected. Accumulated totals are now incremented via a single atomic SQL statement (SessionManager::add_accumulated_usage), and update_session_metrics uses the same path, so a background subagent rolling up concurrently with the parent's reply loop cannot clobber it. --- crates/goose/src/agents/reply_parts.rs | 27 ++-- crates/goose/src/agents/subagent_handler.rs | 43 +++++- crates/goose/src/session/session_manager.rs | 151 ++++++++++++++++++++ 3 files changed, 204 insertions(+), 17 deletions(-) diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 3e7f6f816b60..573a15e569df 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -526,13 +526,10 @@ impl Agent { let manager = self.config.session_manager.clone(); let session = manager.get_session(session_id, false).await?; - let accumulated_usage = session.accumulated_usage + usage.usage; - - let accumulated_cost = session + let cost_delta = session .provider_name .as_deref() - .and_then(|pn| self.accumulate_cost(session.accumulated_cost, usage, pn)) - .or(session.accumulated_cost); + .and_then(|pn| self.estimate_chunk_cost(usage, pn)); let current_usage = if is_compaction_usage { // After compaction: summary output becomes new input context @@ -542,30 +539,28 @@ impl Agent { usage.usage }; + // Accumulated totals must be incremented (not set) so a background + // subagent rolling up concurrently is not clobbered; the context window + // has a single writer and is set outright. manager .update(session_id) .schedule_id(schedule_id) .usage(current_usage) - .accumulated_usage(accumulated_usage) - .accumulated_cost(accumulated_cost) .apply() .await?; + manager + .add_accumulated_usage(session_id, usage.usage, cost_delta) + .await?; + Ok(()) } - fn accumulate_cost( - &self, - existing: Option, - usage: &ProviderUsage, - provider_name: &str, - ) -> Option { + fn estimate_chunk_cost(&self, usage: &ProviderUsage, provider_name: &str) -> Option { let canonical = crate::providers::canonical::maybe_get_canonical_model(provider_name, &usage.model)?; - let chunk_cost = canonical.cost.estimate_cost(&usage.usage)?; - - Some(existing.unwrap_or(0.0) + chunk_cost) + canonical.cost.estimate_cost(&usage.usage) } } diff --git a/crates/goose/src/agents/subagent_handler.rs b/crates/goose/src/agents/subagent_handler.rs index 6c37debbb9b7..155a2195fde8 100644 --- a/crates/goose/src/agents/subagent_handler.rs +++ b/crates/goose/src/agents/subagent_handler.rs @@ -47,7 +47,15 @@ pub struct SubagentRunParams { pub async fn run_subagent_task(params: SubagentRunParams) -> Result { let return_last_only = params.return_last_only; - let (messages, final_output) = get_agent_messages(params).await.map_err(|e| { + let session_manager = params.config.session_manager.clone(); + let parent_session_id = params.task_config.parent_session_id.clone(); + let subagent_session_id = params.session_id.clone(); + + let result = get_agent_messages(params).await; + + roll_up_usage_to_parent(&session_manager, &parent_session_id, &subagent_session_id).await; + + let (messages, final_output) = result.map_err(|e| { ErrorData::new( ErrorCode::INTERNAL_ERROR, format!("Failed to execute task: {}", e), @@ -62,6 +70,39 @@ pub async fn run_subagent_task(params: SubagentRunParams) -> Result String { if return_last_only { messages diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index ed64b8d1eb53..36ae24719dab 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -389,6 +389,20 @@ impl SessionManager { SessionUpdateBuilder::new(self, id.to_string()) } + /// Atomically add a delta to `accumulated_usage`/`accumulated_cost` (never the + /// `usage` context-window columns), so a reply loop and a concurrent subagent + /// roll-up don't clobber each other. `None` cost is a no-op. + pub async fn add_accumulated_usage( + &self, + session_id: &str, + usage: Usage, + cost: Option, + ) -> Result<()> { + self.storage + .add_accumulated_usage(session_id, usage, cost) + .await + } + async fn apply_update_inner(&self, builder: SessionUpdateBuilder<'_>) -> Result<()> { self.storage.apply_update(builder).await } @@ -1560,6 +1574,44 @@ impl SessionStorage { Ok(()) } + async fn add_accumulated_usage( + &self, + session_id: &str, + usage: Usage, + cost: Option, + ) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?; + let result = sqlx::query( + "UPDATE sessions SET \ + accumulated_total_tokens = COALESCE(accumulated_total_tokens, 0) + COALESCE(?, 0), \ + accumulated_input_tokens = COALESCE(accumulated_input_tokens, 0) + COALESCE(?, 0), \ + accumulated_output_tokens = COALESCE(accumulated_output_tokens, 0) + COALESCE(?, 0), \ + accumulated_cache_read_tokens = COALESCE(accumulated_cache_read_tokens, 0) + COALESCE(?, 0), \ + accumulated_cache_write_tokens = COALESCE(accumulated_cache_write_tokens, 0) + COALESCE(?, 0), \ + accumulated_cost = CASE WHEN ? IS NULL THEN accumulated_cost ELSE COALESCE(accumulated_cost, 0) + ? END, \ + updated_at = datetime('now') \ + WHERE id = ?", + ) + .bind(usage.total_tokens) + .bind(usage.input_tokens) + .bind(usage.output_tokens) + .bind(usage.cache_read_input_tokens) + .bind(usage.cache_write_input_tokens) + .bind(cost) + .bind(cost) + .bind(session_id) + .execute(&mut *tx) + .await?; + + if result.rows_affected() == 0 { + return Err(anyhow::anyhow!("Session not found: {}", session_id)); + } + + tx.commit().await?; + Ok(()) + } + async fn get_conversation(&self, session_id: &str) -> Result { let pool = self.pool().await?; let rows = sqlx::query_as::<_, (String, String, i64, Option, Option)>( @@ -3621,4 +3673,103 @@ mod tests { assert_eq!(loaded.usage, usage); assert_eq!(loaded.accumulated_usage, accumulated_usage); } + + #[tokio::test] + async fn test_add_accumulated_usage_folds_into_accumulated_not_context() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + let parent = sm + .create_session( + PathBuf::from("/tmp/test"), + "Parent".to_string(), + SessionType::User, + GooseMode::default(), + ) + .await + .unwrap(); + + let parent_context = Usage::new(Some(1000), Some(200), Some(1200)); + let parent_accumulated = Usage::new(Some(5000), Some(800), Some(5800)); + sm.update(&parent.id) + .usage(parent_context) + .accumulated_usage(parent_accumulated) + .accumulated_cost(Some(0.50)) + .apply() + .await + .unwrap(); + + let subagent_usage = + Usage::new(Some(3000), Some(400), Some(3400)).with_cache_tokens(Some(2000), Some(100)); + sm.add_accumulated_usage(&parent.id, subagent_usage, Some(0.25)) + .await + .unwrap(); + + let loaded = sm.get_session(&parent.id, false).await.unwrap(); + + assert_eq!( + loaded.usage, parent_context, + "context-window usage must not absorb subagent tokens" + ); + assert_eq!(loaded.accumulated_usage.total_tokens, Some(5800 + 3400)); + assert_eq!(loaded.accumulated_usage.input_tokens, Some(5000 + 3000)); + assert_eq!(loaded.accumulated_usage.output_tokens, Some(800 + 400)); + assert_eq!(loaded.accumulated_usage.cache_read_input_tokens, Some(2000)); + assert_eq!(loaded.accumulated_usage.cache_write_input_tokens, Some(100)); + assert_eq!(loaded.accumulated_cost, Some(0.75)); + + sm.add_accumulated_usage(&parent.id, Usage::new(Some(1), Some(1), Some(2)), None) + .await + .unwrap(); + let loaded = sm.get_session(&parent.id, false).await.unwrap(); + assert_eq!(loaded.accumulated_usage.total_tokens, Some(5800 + 3400 + 2)); + assert_eq!( + loaded.accumulated_cost, + Some(0.75), + "a None cost estimate must leave accumulated cost unchanged" + ); + } + + // Concurrent increments must all survive (lost-update regression). + #[tokio::test] + async fn test_add_accumulated_usage_concurrent_increments_are_not_lost() { + let temp_dir = TempDir::new().unwrap(); + let sm = Arc::new(SessionManager::new(temp_dir.path().to_path_buf())); + let parent = sm + .create_session( + PathBuf::from("/tmp/test"), + "Parent".to_string(), + SessionType::User, + GooseMode::default(), + ) + .await + .unwrap(); + sm.storage().pool().await.unwrap(); + + let writers = 50; + let mut handles = Vec::new(); + for _ in 0..writers { + let sm = Arc::clone(&sm); + let id = parent.id.clone(); + handles.push(tokio::spawn(async move { + sm.add_accumulated_usage(&id, Usage::new(Some(10), Some(2), Some(12)), Some(0.01)) + .await + .unwrap(); + })); + } + for handle in handles { + handle.await.unwrap(); + } + + let loaded = sm.get_session(&parent.id, false).await.unwrap(); + assert_eq!(loaded.accumulated_usage.total_tokens, Some(12 * writers)); + assert_eq!(loaded.accumulated_usage.input_tokens, Some(10 * writers)); + assert_eq!(loaded.accumulated_usage.output_tokens, Some(2 * writers)); + let cost = loaded.accumulated_cost.unwrap(); + assert!( + (cost - 0.01 * writers as f64).abs() < 1e-9, + "expected ~{}, got {}", + 0.01 * writers as f64, + cost + ); + } } From e4bd435dec054efe44556d2f1674ce5a0a0203fc Mon Sep 17 00:00:00 2001 From: Filip Kujawa Date: Tue, 30 Jun 2026 21:23:02 -0700 Subject: [PATCH 2/2] refactor(session): fold parent usage update into a single UPDATE Address review: replace the two-statement parent metrics update (set usage + add_accumulated_usage) with one SessionManager::update_usage_metrics that sets the context-window columns and increments the accumulated totals/cost in a single atomic UPDATE. add_accumulated_usage remains for the subagent roll-up, which only touches accumulated columns. --- crates/goose/src/agents/reply_parts.rs | 21 ++-- crates/goose/src/session/session_manager.rs | 120 ++++++++++++++++++++ 2 files changed, 130 insertions(+), 11 deletions(-) diff --git a/crates/goose/src/agents/reply_parts.rs b/crates/goose/src/agents/reply_parts.rs index 573a15e569df..ebebed9ac06e 100644 --- a/crates/goose/src/agents/reply_parts.rs +++ b/crates/goose/src/agents/reply_parts.rs @@ -539,18 +539,17 @@ impl Agent { usage.usage }; - // Accumulated totals must be incremented (not set) so a background - // subagent rolling up concurrently is not clobbered; the context window - // has a single writer and is set outright. + // Set the context window outright (single writer) while incrementing the + // accumulated totals, so a subagent rolling up concurrently is not + // clobbered - both in one statement. manager - .update(session_id) - .schedule_id(schedule_id) - .usage(current_usage) - .apply() - .await?; - - manager - .add_accumulated_usage(session_id, usage.usage, cost_delta) + .update_usage_metrics( + session_id, + schedule_id, + current_usage, + usage.usage, + cost_delta, + ) .await?; Ok(()) diff --git a/crates/goose/src/session/session_manager.rs b/crates/goose/src/session/session_manager.rs index 36ae24719dab..7d2e7193a564 100644 --- a/crates/goose/src/session/session_manager.rs +++ b/crates/goose/src/session/session_manager.rs @@ -403,6 +403,29 @@ impl SessionManager { .await } + /// Set the context-window `usage` outright and increment the accumulated + /// totals in a single statement. The accumulated columns are incremented (not + /// set) so a concurrent subagent roll-up is not clobbered. `cost_delta` of + /// `None` leaves `accumulated_cost` unchanged. + pub async fn update_usage_metrics( + &self, + session_id: &str, + schedule_id: Option, + current_usage: Usage, + accumulated_delta: Usage, + cost_delta: Option, + ) -> Result<()> { + self.storage + .update_usage_metrics( + session_id, + schedule_id, + current_usage, + accumulated_delta, + cost_delta, + ) + .await + } + async fn apply_update_inner(&self, builder: SessionUpdateBuilder<'_>) -> Result<()> { self.storage.apply_update(builder).await } @@ -1612,6 +1635,58 @@ impl SessionStorage { Ok(()) } + async fn update_usage_metrics( + &self, + session_id: &str, + schedule_id: Option, + current_usage: Usage, + accumulated_delta: Usage, + cost_delta: Option, + ) -> Result<()> { + let pool = self.pool().await?; + let mut tx = pool.begin_with("BEGIN IMMEDIATE").await?; + let result = sqlx::query( + "UPDATE sessions SET \ + schedule_id = ?, \ + total_tokens = ?, \ + input_tokens = ?, \ + output_tokens = ?, \ + cache_read_tokens = ?, \ + cache_write_tokens = ?, \ + accumulated_total_tokens = COALESCE(accumulated_total_tokens, 0) + COALESCE(?, 0), \ + accumulated_input_tokens = COALESCE(accumulated_input_tokens, 0) + COALESCE(?, 0), \ + accumulated_output_tokens = COALESCE(accumulated_output_tokens, 0) + COALESCE(?, 0), \ + accumulated_cache_read_tokens = COALESCE(accumulated_cache_read_tokens, 0) + COALESCE(?, 0), \ + accumulated_cache_write_tokens = COALESCE(accumulated_cache_write_tokens, 0) + COALESCE(?, 0), \ + accumulated_cost = CASE WHEN ? IS NULL THEN accumulated_cost ELSE COALESCE(accumulated_cost, 0) + ? END, \ + updated_at = datetime('now') \ + WHERE id = ?", + ) + .bind(schedule_id) + .bind(current_usage.total_tokens) + .bind(current_usage.input_tokens) + .bind(current_usage.output_tokens) + .bind(current_usage.cache_read_input_tokens) + .bind(current_usage.cache_write_input_tokens) + .bind(accumulated_delta.total_tokens) + .bind(accumulated_delta.input_tokens) + .bind(accumulated_delta.output_tokens) + .bind(accumulated_delta.cache_read_input_tokens) + .bind(accumulated_delta.cache_write_input_tokens) + .bind(cost_delta) + .bind(cost_delta) + .bind(session_id) + .execute(&mut *tx) + .await?; + + if result.rows_affected() == 0 { + return Err(anyhow::anyhow!("Session not found: {}", session_id)); + } + + tx.commit().await?; + Ok(()) + } + async fn get_conversation(&self, session_id: &str) -> Result { let pool = self.pool().await?; let rows = sqlx::query_as::<_, (String, String, i64, Option, Option)>( @@ -3729,6 +3804,51 @@ mod tests { ); } + #[tokio::test] + async fn test_update_usage_metrics_sets_context_and_increments_accumulated() { + let temp_dir = TempDir::new().unwrap(); + let sm = SessionManager::new(temp_dir.path().to_path_buf()); + let parent = sm + .create_session( + PathBuf::from("/tmp/test"), + "Parent".to_string(), + SessionType::User, + GooseMode::default(), + ) + .await + .unwrap(); + + let first = Usage::new(Some(1000), Some(200), Some(1200)); + sm.update_usage_metrics(&parent.id, None, first, first, Some(0.10)) + .await + .unwrap(); + + let second = Usage::new(Some(1500), Some(300), Some(1800)); + sm.update_usage_metrics( + &parent.id, + Some("sched-1".to_string()), + second, + second, + None, + ) + .await + .unwrap(); + + let loaded = sm.get_session(&parent.id, false).await.unwrap(); + assert_eq!( + loaded.usage, second, + "context window is set to the latest snapshot" + ); + assert_eq!(loaded.accumulated_usage.total_tokens, Some(1200 + 1800)); + assert_eq!(loaded.accumulated_usage.input_tokens, Some(1000 + 1500)); + assert_eq!(loaded.schedule_id, Some("sched-1".to_string())); + assert_eq!( + loaded.accumulated_cost, + Some(0.10), + "None cost_delta must leave accumulated cost unchanged" + ); + } + // Concurrent increments must all survive (lost-update regression). #[tokio::test] async fn test_add_accumulated_usage_concurrent_increments_are_not_lost() {