[BUG](foundation-api): Parse search usage records

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,

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<SubagentUsage> },
+}
+
+/// 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.

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))

rust/foundation-api/src/routes/agent/mod.rs

@@ -342,28 +342,27 @@ fn extract_subagent_usages(observation: &Observation) -> Vec<InferenceUsage> {
         .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::<Vec<_>>(),
+            )
         })
+        .flatten()
         .collect()
 }
 

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<UsageData>),
     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::<UsageEventData>(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,20 @@ pub(crate) struct UsageData {
     pub cache_write_tokens: u64,
 }
 
+/// The usage envelope emitted by search-agent-research.
+///
+/// Search agents report one record per model under `usage_records`.
+#[derive(Debug, Clone, Deserialize)]
+struct UsageEventData {
+    usage_records: Vec<UsageData>,
+}
+
+impl UsageEventData {
+    fn into_records(self) -> Vec<UsageData> {
+        self.usage_records
+    }
+}
+
 impl ActionData {
     /// The `text` of this action's last `user_text` tool, if any — the agent
     /// "speaking" to the user.
@@ -165,7 +188,7 @@ pub(crate) enum SubagentResultError {
 #[derive(Debug, Clone, PartialEq)]
 pub(crate) struct SubagentSearchResult {
     pub documents: Vec<RankedDocument>,
-    pub usage: Option<UsageData>,
+    pub usages: Vec<UsageData>,
 }
 
 /// Matches one `<Document id=…><Justification>…</Justification></Document>`

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<UsageData>), SubagentResultError> {
+) -> Result<(String, Vec<UsageData>), 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<String> = None;
-    let mut usage: Option<UsageData> = 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,
     })
 }
 

rust/foundation-api/src/routes/subagent_search/tests/00_unit.rs

@@ -40,15 +40,34 @@ fn agent_event_parses_each_kind() {
     ));
     assert!(matches!(
         AgentEvent::parse(
-            &json!({"type":"usage","data":{"model":"scout","input_tokens":123,"output_tokens":456}}).to_string()
+            &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(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,
+                },
+                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()),
@@ -66,6 +85,20 @@ 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": []}),
+        json!({"model": "scout", "input_tokens": 123, "output_tokens": 456}),
+    ] {
+        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.

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\":\"<Document id=doc-1><Justification>Relevant to rag.</Justification></Document>\"}]}}\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);
 }