fix: record streaming token usage in v2 profile stats

crates/switchyard-components-v2/Cargo.toml

@@ -12,7 +12,9 @@ repository.workspace = true
 rust-version.workspace = true
 
 [dependencies]
+async-stream = "0.3"
 async-trait = "0.1"
+futures-util = "0.3"
 parking_lot = "0.12"
 rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots"] }

crates/switchyard-components-v2/src/profiles/latency_service/mod.rs

@@ -16,7 +16,6 @@ use std::time::Instant;
 
 use async_trait::async_trait;
 use parking_lot::RwLock;
-use switchyard_components::stats::usage_from_body;
 use switchyard_components::StatsAccumulator;
 use switchyard_core::{ChatResponse, LlmTarget, LlmTargetId, Result, SwitchyardError};
 
@@ -26,6 +25,7 @@ use self::selection::select_target;
 pub use self::selection::SelectedTarget;
 use crate::backend::{native_target_backend, TargetBackend};
 use crate::profile_stats_accumulator;
+use crate::stats::{record_usage_or_tap_stream, UsageAttribution};
 use crate::{
     profile_config, Profile, ProfileConfig, ProfileHooks, ProfileInput, ProfileResponse,
     RoutingMetadata,
@@ -329,16 +329,16 @@ impl Profile for LatencyServiceProfile {
                         Some(backend_latency_ms),
                         None,
                     )?;
-                    let total_latency_ms = profile_started_at.elapsed().as_secs_f64() * 1000.0;
-                    let routing_overhead_ms = (total_latency_ms - backend_latency_ms).max(0.0);
-                    let usage = response.body().map(usage_from_body).unwrap_or_default();
-                    self.stats.record_usage_after_success_attribution(
-                        stats_model,
-                        usage,
-                        Some(total_latency_ms),
-                        Some(routing_overhead_ms),
-                        None,
-                    )?;
+                    let response = record_usage_or_tap_stream(
+                        response,
+                        UsageAttribution::new(
+                            self.stats.clone(),
+                            stats_model,
+                            None,
+                            profile_started_at,
+                            backend_latency_ms,
+                        ),
+                    );
                     let processed = LatencyServiceProcessedRequest {
                         profile_input: input,
                         selected,

crates/switchyard-components-v2/src/profiles/latency_service/tests.rs

@@ -15,6 +15,7 @@ use std::time::Duration;
 
 use async_trait::async_trait;
 use futures_core::Stream;
+use futures_util::StreamExt;
 use parking_lot::RwLock;
 use serde_json::{json, Value};
 use switchyard_core::{BackendFormat, ChatRequest, LlmTargetId, ModelId, StreamEvent};
@@ -42,6 +43,7 @@ struct TestBackend {
 enum ResponseMode {
     Completion,
     OpenAiStream,
+    OpenAiStreamWithUsage,
 }
 
 struct OneEventStream {
@@ -86,6 +88,15 @@ impl ProfileBackend for TestBackend {
                     })))),
                 })))
             }
+            ResponseMode::OpenAiStreamWithUsage => {
+                Ok(ChatResponse::OpenAiStream(Box::pin(OneEventStream {
+                    event: Some(Ok(StreamEvent::Json(json!({
+                        "backend": self.name,
+                        "model": request.model(),
+                        "usage": {"prompt_tokens": 21, "completion_tokens": 11},
+                    })))),
+                })))
+            }
         }
     }
 }
@@ -593,6 +604,54 @@ fn unknown_target_health_update_is_rejected() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+async fn streaming_response_records_usage_after_consumption() -> Result<()> {
+    let (profile, _calls) = profile_with_backend_modes(
+        vec![target("fast", "upstream-fast")?],
+        BTreeMap::new(),
+        BTreeMap::from([("fast", ResponseMode::OpenAiStreamWithUsage)]),
+        0,
+    )?;
+    profile.update_health(
+        LlmTargetId::from_static("fast"),
+        EndpointHealth::new(EndpointHealthStatus::Healthy),
+    )?;
+
+    let response = profile
+        .run(profile_input(ChatRequest::openai_chat(json!({
+            "model": "incoming-model",
+            "messages": [],
+        }))))
+        .await?;
+
+    // The call is attributed immediately; streaming usage is recorded lazily.
+    let before = profile.stats.snapshot()?;
+    assert_eq!(before.total_requests, 1);
+    assert_eq!(before.total_tokens.prompt, 0);
+
+    let mut stream = match response.response {
+        ChatResponse::OpenAiStream(stream) => stream,
+        _ => {
+            return Err(SwitchyardError::Other(
+                "expected OpenAI stream response".into(),
+            ))
+        }
+    };
+    while let Some(event) = stream.next().await {
+        event?;
+    }
+
+    let after = profile.stats.snapshot()?;
+    assert_eq!(after.total_tokens.prompt, 21);
+    assert_eq!(after.total_tokens.completion, 11);
+    let model = after
+        .models
+        .get("upstream-fast")
+        .ok_or_else(|| SwitchyardError::Other("model stats should be present".into()))?;
+    assert_eq!(model.calls, 1);
+    Ok(())
+}
+
 #[tokio::test]
 async fn poll_once_updates_cache_and_ignores_unknown_service_ids() -> Result<()> {
     let mut server = OneShotServer::json(

crates/switchyard-components-v2/src/profiles/llm_routing.rs

@@ -17,6 +17,7 @@ use switchyard_core::{
 
 use crate::backend::{native_target_backend, TargetBackend};
 use crate::profile_stats_accumulator;
+use crate::stats::{record_usage_or_tap_stream, UsageAttribution};
 use crate::{
     profile_config, Profile, ProfileConfig, ProfileHooks, ProfileInput, ProfileResponse,
     RoutingMetadata,
@@ -421,25 +422,25 @@ impl LlmRoutingProfile {
     fn record_success(
         &self,
         decision: &LlmRoutingDecision,
-        response: &ChatResponse,
-        total_latency_ms: f64,
+        response: ChatResponse,
+        started_at: Instant,
         backend_latency_ms: f64,
-    ) -> Result<()> {
+    ) -> Result<ChatResponse> {
         self.stats.record_success(
             decision.selected_model.as_str(),
             Some(backend_latency_ms),
             Some(decision.tier.as_str()),
         )?;
-        let routing_overhead_ms = (total_latency_ms - backend_latency_ms).max(0.0);
-        let usage = response.body().map(usage_from_body).unwrap_or_default();
-        self.stats.record_usage_after_success_attribution(
-            decision.selected_model.as_str(),
-            usage,
-            Some(total_latency_ms),
-            Some(routing_overhead_ms),
-            Some(decision.tier.as_str()),
-        )?;
-        Ok(())
+        Ok(record_usage_or_tap_stream(
+            response,
+            UsageAttribution::new(
+                self.stats.clone(),
+                decision.selected_model.as_str(),
+                Some(decision.tier.clone()),
+                started_at,
+                backend_latency_ms,
+            ),
+        ))
     }
 
     fn record_error(&self, decision: &LlmRoutingDecision) -> Result<()> {
@@ -489,11 +490,10 @@ impl Profile for LlmRoutingProfile {
         let (first_result, first_backend_latency_ms) = self.call_selected(&processed).await;
         match first_result {
             Ok(response) => {
-                let total_latency_ms = profile_started_at.elapsed().as_secs_f64() * 1000.0;
-                self.record_success(
+                let response = self.record_success(
                     &processed.decision,
-                    &response,
-                    total_latency_ms,
+                    response,
+                    profile_started_at,
                     first_backend_latency_ms,
                 )?;
                 let response = self.rprocess(&processed, response).await?;
@@ -507,11 +507,10 @@ impl Profile for LlmRoutingProfile {
                 let (retry_result, retry_backend_latency_ms) = self.call_selected(&retry).await;
                 match retry_result {
                     Ok(response) => {
-                        let total_latency_ms = profile_started_at.elapsed().as_secs_f64() * 1000.0;
-                        self.record_success(
+                        let response = self.record_success(
                             &retry.decision,
-                            &response,
-                            total_latency_ms,
+                            response,
+                            profile_started_at,
                             retry_backend_latency_ms,
                         )?;
                         let response = self.rprocess(&retry, response).await?;
@@ -1603,3 +1602,170 @@ fn normalize_reasoning_effort(request: &mut ChatRequest) {
         Value::String("high".to_string()),
     );
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use async_trait::async_trait;
+    use futures_util::StreamExt;
+    use serde_json::json;
+    use switchyard_core::{ModelId, StreamEvent};
+
+    use crate::backend::ProfileBackend;
+
+    use super::*;
+
+    // Stub backend used only to satisfy the profile struct; `record_success`
+    // never dispatches through it.
+    struct UnusedBackend;
+
+    #[async_trait]
+    impl ProfileBackend for UnusedBackend {
+        async fn call(&self, _request: &ChatRequest) -> Result<ChatResponse> {
+            Err(SwitchyardError::Other(
+                "backend should not be called".to_string(),
+            ))
+        }
+    }
+
+    fn target(id: &str, model: &str) -> Result<LlmTarget> {
+        let mut target = LlmTarget::new(LlmTargetId::new(id)?, ModelId::new(model)?);
+        target.format = BackendFormat::OpenAi;
+        Ok(target)
+    }
+
+    fn stub_backend(target: &LlmTarget) -> TargetBackend {
+        TargetBackend::new(target.clone(), Arc::new(UnusedBackend))
+    }
+
+    fn config(
+        strong: &LlmTarget,
+        weak: &LlmTarget,
+        classifier: &LlmTarget,
+    ) -> LlmRoutingProfileConfig {
+        LlmRoutingProfileConfig {
+            strong: strong.clone(),
+            weak: weak.clone(),
+            classifier: classifier.clone(),
+            fallback_target_on_evict: strong.id.clone(),
+            profile_name: PROFILE_GENERAL.to_string(),
+            classifier_min_confidence: 0.0,
+            classifier_fail_open: true,
+            classifier_recent_turn_window: 4,
+            classifier_max_tokens: 256,
+            alignment_min_confidence: DEFAULT_ALIGNMENT_MIN_CONFIDENCE,
+            default_tier: None,
+            tier_mapping: None,
+            classifier_system_prompt: None,
+            classifier_tool_name: None,
+            classifier_tool_description: None,
+            classifier_tool_parameters: None,
+        }
+    }
+
+    fn profile(stats: StatsAccumulator) -> Result<LlmRoutingProfile> {
+        let strong = target("strong", "frontier/model")?;
+        let weak = target("weak", "cheap/model")?;
+        let classifier = target("classifier", "cls/model")?;
+        let config = config(&strong, &weak, &classifier);
+        let policy = ClassifierPolicy::from_name(&config.profile_name)?;
+        Ok(LlmRoutingProfile {
+            policy,
+            strong_backend: stub_backend(&strong),
+            weak_backend: stub_backend(&weak),
+            classifier_backend: stub_backend(&classifier),
+            classifier_target: classifier,
+            fallback_target_on_evict: config.fallback_target_on_evict.clone(),
+            classifier_min_confidence: config.classifier_min_confidence,
+            classifier_fail_open: config.classifier_fail_open,
+            classifier_recent_turn_window: config.classifier_recent_turn_window,
+            classifier_max_tokens: config.classifier_max_tokens,
+            alignment_min_confidence: config.alignment_min_confidence,
+            default_tier: resolve_default_tier(config.default_tier.as_deref())?,
+            tier_mapping: config
+                .tier_mapping
+                .clone()
+                .unwrap_or_else(|| policy.default_tier_mapping()),
+            classifier_tool: LlmRoutingClassifierTool::from_config(policy, &config),
+            stats,
+        })
+    }
+
+    fn weak_decision() -> LlmRoutingDecision {
+        LlmRoutingDecision {
+            tier: TIER_WEAK.to_string(),
+            source: "unit_test".to_string(),
+            reason: "unit_test".to_string(),
+            policy_tier: None,
+            llm_recommended_tier: None,
+            confidence: None,
+            selected_target: "weak".to_string(),
+            selected_model: "cheap/model".to_string(),
+            signals: None,
+        }
+    }
+
+    #[tokio::test]
+    async fn record_success_taps_streaming_usage_on_consumption() -> Result<()> {
+        let stats = StatsAccumulator::new();
+        let profile = profile(stats.clone())?;
+
+        let events: Vec<Result<StreamEvent>> = vec![
+            Ok(StreamEvent::Json(json!({
+                "choices": [{"delta": {"content": "hi"}}],
+            }))),
+            Ok(StreamEvent::Json(json!({
+                "choices": [],
+                "usage": {"prompt_tokens": 23, "completion_tokens": 12},
+            }))),
+        ];
+        let response = ChatResponse::OpenAiStream(Box::pin(futures_util::stream::iter(events)));
+
+        let response = profile.record_success(&weak_decision(), response, Instant::now(), 4.0)?;
+
+        // Success is attributed synchronously; streaming usage waits for consume.
+        let before = stats.snapshot()?;
+        assert_eq!(before.total_requests, 1);
+        assert_eq!(before.total_tokens.prompt, 0);
+
+        let mut stream = match response {
+            ChatResponse::OpenAiStream(stream) => stream,
+            _ => return Err(SwitchyardError::Other("expected stream response".into())),
+        };
+        let mut forwarded = 0;
+        while let Some(event) = stream.next().await {
+            event?;
+            forwarded += 1;
+        }
+        assert_eq!(forwarded, 2, "every event must still reach the client");
+
+        let after = stats.snapshot()?;
+        assert_eq!(after.total_tokens.prompt, 23);
+        assert_eq!(after.total_tokens.completion, 12);
+        let tier = after
+            .tiers
+            .get(TIER_WEAK)
+            .ok_or_else(|| SwitchyardError::Other("weak tier stats should be present".into()))?;
+        assert_eq!(tier.calls, 1);
+        assert_eq!(tier.model, "cheap/model");
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn record_success_records_buffered_usage_immediately() -> Result<()> {
+        let stats = StatsAccumulator::new();
+        let profile = profile(stats.clone())?;
+
+        let response = ChatResponse::openai_completion(json!({
+            "usage": {"prompt_tokens": 7, "completion_tokens": 2},
+        }));
+        let _response = profile.record_success(&weak_decision(), response, Instant::now(), 4.0)?;
+
+        let snapshot = stats.snapshot()?;
+        assert_eq!(snapshot.total_requests, 1);
+        assert_eq!(snapshot.total_tokens.prompt, 7);
+        assert_eq!(snapshot.total_tokens.completion, 2);
+        Ok(())
+    }
+}