From d2ab86e838f53ef960cdb8df4104fe445e469b45 Mon Sep 17 00:00:00 2001 From: Tanuj Nayak Date: Tue, 28 Jul 2026 18:47:58 -0700 Subject: [PATCH 1/3] [BUG](foundation-api): Parse search usage records --- rust/agent/src/lib.rs | 2 +- rust/agent/src/tool.rs | 18 ++++--- .../src/agent_tools/subagent_search_tool.rs | 19 ++++--- rust/foundation-api/src/routes/agent/mod.rs | 29 +++++----- .../src/routes/subagent_search/events.rs | 54 +++++++++++++++++-- .../src/routes/subagent_search/mod.rs | 12 ++--- .../routes/subagent_search/tests/00_unit.rs | 46 +++++++++++++--- .../subagent_search/tests/01_integration.rs | 18 ++++--- 8 files changed, 144 insertions(+), 54 deletions(-) diff --git a/rust/agent/src/lib.rs b/rust/agent/src/lib.rs index d1e503dafcc..673c5191214 100644 --- a/rust/agent/src/lib.rs +++ b/rust/agent/src/lib.rs @@ -24,7 +24,7 @@ pub use inference::{ UnknownAnthropicModel, }; pub use provider::ProviderFormat; -pub use tool::{DynTool, Tool, ToolCallMetadata, ToolSet}; +pub use tool::{DynTool, SubagentUsage, Tool, ToolCallMetadata, ToolSet}; pub use tools::weather::{GetWeatherTool, TemperatureUnit}; pub use trajectory::{ Action, ActionBuilder, ActionItem, Call, Entry, Observation, ObservationBuilder, diff --git a/rust/agent/src/tool.rs b/rust/agent/src/tool.rs index a00af898c5f..527278dfa37 100644 --- a/rust/agent/src/tool.rs +++ b/rust/agent/src/tool.rs @@ -28,13 +28,17 @@ use crate::provider::ProviderFormat; pub enum ToolCallMetadata { /// Token usage reported by the deep-research subagent behind /// Foundation's `subagent_search` tool. - SubagentUsage { - model: String, - input_tokens: u64, - output_tokens: u64, - cache_read_tokens: u64, - cache_write_tokens: u64, - }, + SubagentUsage { usages: Vec }, +} + +/// Token usage from one model used by a deep-research subagent. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct SubagentUsage { + pub model: String, + pub input_tokens: u64, + pub output_tokens: u64, + pub cache_read_tokens: u64, + pub cache_write_tokens: u64, } /// THE trait you implement to define a tool. diff --git a/rust/foundation-api/src/agent_tools/subagent_search_tool.rs b/rust/foundation-api/src/agent_tools/subagent_search_tool.rs index 923d4ccfdca..3b00463e31c 100644 --- a/rust/foundation-api/src/agent_tools/subagent_search_tool.rs +++ b/rust/foundation-api/src/agent_tools/subagent_search_tool.rs @@ -10,7 +10,7 @@ use async_trait::async_trait; use schemars::JsonSchema; use serde::Deserialize; -use chroma_agent::{AgentError, Tool, ToolCallMetadata}; +use chroma_agent::{AgentError, SubagentUsage, Tool, ToolCallMetadata}; use crate::routes::subagent_search::{subagent_search_text, SubagentSearchCreds}; @@ -77,12 +77,17 @@ impl Tool for SubagentSearchTool { .await .map_err(|err| AgentError::Tool(err.to_string()))?; - let metadata = usage.map(|usage| ToolCallMetadata::SubagentUsage { - model: usage.model, - input_tokens: usage.input_tokens, - output_tokens: usage.output_tokens, - cache_read_tokens: usage.cache_read_tokens, - cache_write_tokens: usage.cache_write_tokens, + let metadata = (!usage.is_empty()).then(|| ToolCallMetadata::SubagentUsage { + usages: usage + .into_iter() + .map(|usage| SubagentUsage { + model: usage.model, + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cache_read_tokens: usage.cache_read_tokens, + cache_write_tokens: usage.cache_write_tokens, + }) + .collect(), }); Ok((text, metadata)) diff --git a/rust/foundation-api/src/routes/agent/mod.rs b/rust/foundation-api/src/routes/agent/mod.rs index e5330d71930..82a1a70ea86 100644 --- a/rust/foundation-api/src/routes/agent/mod.rs +++ b/rust/foundation-api/src/routes/agent/mod.rs @@ -342,28 +342,27 @@ fn extract_subagent_usages(observation: &Observation) -> Vec { .iter() .filter_map(|item| { let ObservationItem::ToolResult { - metadata: - Some(chroma_agent::ToolCallMetadata::SubagentUsage { - model, - input_tokens, - output_tokens, - cache_read_tokens, - cache_write_tokens, - }), + metadata: Some(chroma_agent::ToolCallMetadata::SubagentUsage { usages }), .. } = item else { return None; }; - Some(InferenceUsage { - model: model.clone(), - input_tokens: *input_tokens, - output_tokens: *output_tokens, - cache_read_tokens: *cache_read_tokens, - cache_write_tokens: *cache_write_tokens, - }) + Some( + usages + .iter() + .map(|usage| InferenceUsage { + model: usage.model.clone(), + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cache_read_tokens: usage.cache_read_tokens, + cache_write_tokens: usage.cache_write_tokens, + }) + .collect::>(), + ) }) + .flatten() .collect() } diff --git a/rust/foundation-api/src/routes/subagent_search/events.rs b/rust/foundation-api/src/routes/subagent_search/events.rs index f3daf3d913c..4fe79d3f8f8 100644 --- a/rust/foundation-api/src/routes/subagent_search/events.rs +++ b/rust/foundation-api/src/routes/subagent_search/events.rs @@ -21,7 +21,7 @@ use std::sync::LazyLock; pub(crate) enum AgentEvent { Action(ActionData), Observation(ObservationData), - Usage(UsageData), + Usage(Vec), Done, Error(ErrorData), Unknown, @@ -41,7 +41,16 @@ impl AgentEvent { Some("observation") => { from_data(data()).map_or(AgentEvent::Unknown, AgentEvent::Observation) } - Some("usage") => from_data(data()).map_or(AgentEvent::Unknown, AgentEvent::Usage), + Some("usage") => match from_data::(data()) { + Some(usage) => AgentEvent::Usage(usage.into_records()), + None => { + tracing::warn!( + payload = %data(), + "dropping malformed search-agent usage event" + ); + AgentEvent::Unknown + } + }, Some("done") => AgentEvent::Done, Some("error") => from_data(data()).map_or(AgentEvent::Unknown, AgentEvent::Error), _ => AgentEvent::Unknown, @@ -102,6 +111,45 @@ pub(crate) struct UsageData { pub cache_write_tokens: u64, } +/// The usage envelope emitted by search-agent-research. +/// +/// Current search agents report one record per model under `usage_records`; +/// retain the original flat shape for compatibility with older deployments. +#[derive(Debug, Clone, Deserialize)] +struct UsageEventData { + #[serde(default)] + usage_records: Vec, + #[serde(default)] + model: Option, + #[serde(default)] + input_tokens: u64, + #[serde(default)] + output_tokens: u64, + #[serde(default)] + cache_read_tokens: u64, + #[serde(default)] + cache_write_tokens: u64, +} + +impl UsageEventData { + fn into_records(self) -> Vec { + if !self.usage_records.is_empty() { + return self.usage_records; + } + + self.model + .map(|model| UsageData { + model, + input_tokens: self.input_tokens, + output_tokens: self.output_tokens, + cache_read_tokens: self.cache_read_tokens, + cache_write_tokens: self.cache_write_tokens, + }) + .into_iter() + .collect() + } +} + impl ActionData { /// The `text` of this action's last `user_text` tool, if any — the agent /// "speaking" to the user. @@ -165,7 +213,7 @@ pub(crate) enum SubagentResultError { #[derive(Debug, Clone, PartialEq)] pub(crate) struct SubagentSearchResult { pub documents: Vec, - pub usage: Option, + pub usages: Vec, } /// Matches one `…` diff --git a/rust/foundation-api/src/routes/subagent_search/mod.rs b/rust/foundation-api/src/routes/subagent_search/mod.rs index d432224cb36..f46d584f80c 100644 --- a/rust/foundation-api/src/routes/subagent_search/mod.rs +++ b/rust/foundation-api/src/routes/subagent_search/mod.rs @@ -359,12 +359,12 @@ pub(crate) async fn subagent_search_text( creds: SubagentSearchCreds, query: String, ui_origin: Option<&str>, -) -> Result<(String, Option), SubagentResultError> { +) -> Result<(String, Vec), SubagentResultError> { let tenant = creds.chroma_tenant.clone(); let result = collect_subagent_search_final(http, url, creds, query).await?; Ok(( format_ranked_documents(&result.documents, ui_origin, &tenant), - result.usage, + result.usages, )) } @@ -421,7 +421,7 @@ pub(crate) async fn collect_subagent_search_final( // Keep the last action's `user_text` — the agent's final answer. let mut final_answer: Option = None; - let mut usage: Option = None; + let mut usages = Vec::new(); let mut saw_done = false; while let Some(item) = stream.next().await { let raw = item.map_err(SubagentResultError::Stream)?; @@ -438,8 +438,8 @@ pub(crate) async fn collect_subagent_search_final( saw_done = true; break; } - AgentEvent::Usage(event_usage) => { - usage = Some(event_usage); + AgentEvent::Usage(event_usages) => { + usages.extend(event_usages); } AgentEvent::Observation(_) | AgentEvent::Unknown => {} } @@ -456,7 +456,7 @@ pub(crate) async fn collect_subagent_search_final( .as_deref() .map(parse_ranked_documents) .unwrap_or_default(), - usage, + usages, }) } diff --git a/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs b/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs index 1397c7a30ab..6957eb81167 100644 --- a/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs +++ b/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs @@ -42,13 +42,45 @@ fn agent_event_parses_each_kind() { AgentEvent::parse( &json!({"type":"usage","data":{"model":"scout","input_tokens":123,"output_tokens":456}}).to_string() ), - AgentEvent::Usage(UsageData { - model, - input_tokens: 123, - output_tokens: 456, - cache_read_tokens: 0, - cache_write_tokens: 0, - }) if model == "scout" + AgentEvent::Usage(usages) + if usages == vec![UsageData { + model: "scout".to_string(), + input_tokens: 123, + output_tokens: 456, + cache_read_tokens: 0, + cache_write_tokens: 0, + }] + )); + assert!(matches!( + AgentEvent::parse( + &json!({ + "type": "usage", + "data": { + "usage_records": [ + {"model":"scout","input_tokens":123,"output_tokens":456}, + {"model":"max","input_tokens":7,"output_tokens":8,"cache_read_tokens":9} + ] + } + }) + .to_string() + ), + AgentEvent::Usage(usages) + if usages == vec![ + UsageData { + model: "scout".to_string(), + input_tokens: 123, + output_tokens: 456, + cache_read_tokens: 0, + cache_write_tokens: 0, + }, + UsageData { + model: "max".to_string(), + input_tokens: 7, + output_tokens: 8, + cache_read_tokens: 9, + cache_write_tokens: 0, + }, + ] )); assert!(matches!( AgentEvent::parse(&json!({"type":"done","data":{}}).to_string()), diff --git a/rust/foundation-api/src/routes/subagent_search/tests/01_integration.rs b/rust/foundation-api/src/routes/subagent_search/tests/01_integration.rs index d205f22586d..db8ec4c8024 100644 --- a/rust/foundation-api/src/routes/subagent_search/tests/01_integration.rs +++ b/rust/foundation-api/src/routes/subagent_search/tests/01_integration.rs @@ -24,7 +24,7 @@ async fn streams_and_collects_final_from_mocked_sse() { "data: {\"type\":\"action\",\"data\":{\"tools\":[{\"name\":\"search\"}],\"params\":[{\"query\":\"rag\"}]}}\n\n", "data: {\"type\":\"observation\",\"data\":{\"sources\":[\"a\"]}}\n\n", "data: {\"type\":\"action\",\"data\":{\"tools\":[{\"name\":\"user_text\"}],\"params\":[{\"text\":\"Relevant to rag.\"}]}}\n\n", - "data: {\"type\":\"usage\",\"data\":{\"model\":\"scout\",\"input_tokens\":123,\"output_tokens\":456}}\n\n", + "data: {\"type\":\"usage\",\"data\":{\"usage_records\":[{\"model\":\"scout\",\"input_tokens\":123,\"output_tokens\":456},{\"model\":\"max\",\"input_tokens\":7,\"output_tokens\":8,\"cache_read_tokens\":9}]}}\n\n", "data: {\"type\":\"done\",\"data\":{}}\n\n", ); let mock = server @@ -69,12 +69,14 @@ async fn streams_and_collects_final_from_mocked_sse() { justification: "Relevant to rag.".to_string(), }] ); - let usage = result.usage.expect("usage should be present"); - assert_eq!(usage.model, "scout"); - assert_eq!(usage.input_tokens, 123); - assert_eq!(usage.output_tokens, 456); - assert_eq!(usage.cache_read_tokens, 0); - assert_eq!(usage.cache_write_tokens, 0); + assert_eq!(result.usages.len(), 2); + assert_eq!(result.usages[0].model, "scout"); + assert_eq!(result.usages[0].input_tokens, 123); + assert_eq!(result.usages[0].output_tokens, 456); + assert_eq!(result.usages[0].cache_read_tokens, 0); + assert_eq!(result.usages[0].cache_write_tokens, 0); + assert_eq!(result.usages[1].model, "max"); + assert_eq!(result.usages[1].cache_read_tokens, 9); assert_eq!(mock.calls(), 2); } @@ -180,7 +182,7 @@ async fn answer_with_no_documents_emits_empty_result_then_done() { .await .expect("empty results are Ok"); assert!(result.documents.is_empty()); - assert!(result.usage.is_none()); + assert!(result.usages.is_empty()); assert_eq!(mock.calls(), 2); } From 37f1721eb4727b21cca66b085f2727298e6815ff Mon Sep 17 00:00:00 2001 From: Tanuj Nayak Date: Wed, 29 Jul 2026 00:13:51 -0700 Subject: [PATCH 2/3] [BUG](foundation-api): Reject invalid usage payloads --- .../src/routes/subagent_search/events.rs | 57 ++++++++++--------- .../routes/subagent_search/tests/00_unit.rs | 10 ++++ 2 files changed, 41 insertions(+), 26 deletions(-) diff --git a/rust/foundation-api/src/routes/subagent_search/events.rs b/rust/foundation-api/src/routes/subagent_search/events.rs index 4fe79d3f8f8..a159b97e71a 100644 --- a/rust/foundation-api/src/routes/subagent_search/events.rs +++ b/rust/foundation-api/src/routes/subagent_search/events.rs @@ -116,37 +116,42 @@ pub(crate) struct UsageData { /// Current search agents report one record per model under `usage_records`; /// retain the original flat shape for compatibility with older deployments. #[derive(Debug, Clone, Deserialize)] -struct UsageEventData { - #[serde(default)] - usage_records: Vec, - #[serde(default)] - model: Option, - #[serde(default)] - input_tokens: u64, - #[serde(default)] - output_tokens: u64, - #[serde(default)] - cache_read_tokens: u64, - #[serde(default)] - cache_write_tokens: u64, +#[serde(untagged)] +enum UsageEventData { + Records { + usage_records: Vec, + }, + Legacy { + model: String, + #[serde(default)] + input_tokens: u64, + #[serde(default)] + output_tokens: u64, + #[serde(default)] + cache_read_tokens: u64, + #[serde(default)] + cache_write_tokens: u64, + }, } impl UsageEventData { fn into_records(self) -> Vec { - if !self.usage_records.is_empty() { - return self.usage_records; - } - - self.model - .map(|model| UsageData { + match self { + Self::Records { usage_records } => usage_records, + Self::Legacy { model, - input_tokens: self.input_tokens, - output_tokens: self.output_tokens, - cache_read_tokens: self.cache_read_tokens, - cache_write_tokens: self.cache_write_tokens, - }) - .into_iter() - .collect() + input_tokens, + output_tokens, + cache_read_tokens, + cache_write_tokens, + } => vec![UsageData { + model, + input_tokens, + output_tokens, + cache_read_tokens, + cache_write_tokens, + }], + } } } diff --git a/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs b/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs index 6957eb81167..bc42bf745a4 100644 --- a/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs +++ b/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs @@ -98,6 +98,16 @@ fn agent_event_parses_each_kind() { assert!(matches!(AgentEvent::parse("not json"), AgentEvent::Unknown)); } +#[test] +fn malformed_usage_events_are_unknown() { + for data in [json!({}), json!({"usage_record": []})] { + assert!(matches!( + AgentEvent::parse(&json!({"type": "usage", "data": data}).to_string()), + AgentEvent::Unknown + )); + } +} + #[test] fn action_user_text_takes_last_and_detects_answer_only() { // A tool-call action with no user_text. From f06846bffc09feb718a5ba784ddb072780c2867b Mon Sep 17 00:00:00 2001 From: Tanuj Nayak Date: Wed, 29 Jul 2026 00:15:57 -0700 Subject: [PATCH 3/3] [BUG](foundation-api): Require usage records --- .../src/routes/subagent_search/events.rs | 38 ++----------------- .../routes/subagent_search/tests/00_unit.rs | 19 +++------- 2 files changed, 9 insertions(+), 48 deletions(-) diff --git a/rust/foundation-api/src/routes/subagent_search/events.rs b/rust/foundation-api/src/routes/subagent_search/events.rs index a159b97e71a..376234e4574 100644 --- a/rust/foundation-api/src/routes/subagent_search/events.rs +++ b/rust/foundation-api/src/routes/subagent_search/events.rs @@ -113,45 +113,15 @@ pub(crate) struct UsageData { /// The usage envelope emitted by search-agent-research. /// -/// Current search agents report one record per model under `usage_records`; -/// retain the original flat shape for compatibility with older deployments. +/// Search agents report one record per model under `usage_records`. #[derive(Debug, Clone, Deserialize)] -#[serde(untagged)] -enum UsageEventData { - Records { - usage_records: Vec, - }, - Legacy { - model: String, - #[serde(default)] - input_tokens: u64, - #[serde(default)] - output_tokens: u64, - #[serde(default)] - cache_read_tokens: u64, - #[serde(default)] - cache_write_tokens: u64, - }, +struct UsageEventData { + usage_records: Vec, } impl UsageEventData { fn into_records(self) -> Vec { - match self { - Self::Records { usage_records } => usage_records, - Self::Legacy { - model, - input_tokens, - output_tokens, - cache_read_tokens, - cache_write_tokens, - } => vec![UsageData { - model, - input_tokens, - output_tokens, - cache_read_tokens, - cache_write_tokens, - }], - } + self.usage_records } } diff --git a/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs b/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs index bc42bf745a4..d5802702e32 100644 --- a/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs +++ b/rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs @@ -38,19 +38,6 @@ fn agent_event_parses_each_kind() { ), AgentEvent::Observation(_) )); - assert!(matches!( - AgentEvent::parse( - &json!({"type":"usage","data":{"model":"scout","input_tokens":123,"output_tokens":456}}).to_string() - ), - AgentEvent::Usage(usages) - if usages == vec![UsageData { - model: "scout".to_string(), - input_tokens: 123, - output_tokens: 456, - cache_read_tokens: 0, - cache_write_tokens: 0, - }] - )); assert!(matches!( AgentEvent::parse( &json!({ @@ -100,7 +87,11 @@ fn agent_event_parses_each_kind() { #[test] fn malformed_usage_events_are_unknown() { - for data in [json!({}), json!({"usage_record": []})] { + for data in [ + json!({}), + json!({"usage_record": []}), + json!({"model": "scout", "input_tokens": 123, "output_tokens": 456}), + ] { assert!(matches!( AgentEvent::parse(&json!({"type": "usage", "data": data}).to_string()), AgentEvent::Unknown