feat: count tokens from streaming (SSE) responses

src/app.rs

@@ -616,6 +616,24 @@ async fn handle_request(
         body_bytes
     };
 
+    // 3b. Try injecting stream_options so the provider returns token usage
+    // in SSE chunks. If the upstream rejects it (400 "unknown_parameter"),
+    // we retry without it in step 4b.
+    let (body_bytes, original_body) = match ensure_stream_options(&body_bytes, provider) {
+        Some(patched) => (Bytes::from(patched), Some(body_bytes)),
+        None => (body_bytes, None),
+    };
+    let retry_headers = if original_body.is_some() {
+        Some(headers.clone())
+    } else {
+        None
+    };
+    let retry_source = if original_body.is_some() {
+        Some(source.clone())
+    } else {
+        None
+    };
+
     // 4. Forward to upstream
     info!(
         method = %method,
@@ -669,6 +687,65 @@ async fn handle_request(
         })?,
     };
 
+    // 4b. If we injected stream_options and the upstream returned 400 with
+    // "unknown_parameter", the model doesn't support it — retry without.
+    let response = if let (Some(original_body), Some(retry_headers), Some(retry_source)) =
+        (original_body, retry_headers, retry_source)
+    {
+        if response.status() == StatusCode::BAD_REQUEST {
+            let (parts, err_body) = response.into_parts();
+            let err_bytes = err_body
+                .collect()
+                .await
+                .map(|b| b.to_bytes())
+                .unwrap_or_default();
+            if String::from_utf8_lossy(&err_bytes).contains("unknown_parameter") {
+                debug!(
+                    method = %method,
+                    path = %path,
+                    "Upstream rejected stream_options; retrying without it"
+                );
+                match retry_source {
+                    KeySource::Static { real_key } => state
+                        .proxy_client
+                        .forward(
+                            method.clone(),
+                            &uri,
+                            retry_headers,
+                            &real_key,
+                            original_body,
+                            provider,
+                        )
+                        .await
+                        .map_err(|e| {
+                            error!(error = %e, "Proxy error on stream_options retry");
+                            e.into_response()
+                        })?,
+                    KeySource::OAuth { provider_id } => forward_oauth_request(
+                        &state,
+                        method.clone(),
+                        &uri,
+                        retry_headers,
+                        original_body,
+                        provider,
+                        &provider_id,
+                    )
+                    .await
+                    .map_err(|e| {
+                        error!(error = %e, "OAuth proxy error on stream_options retry");
+                        error_response(StatusCode::BAD_GATEWAY, &format!("OAuth proxy error: {e}"))
+                    })?,
+                }
+            } else {
+                Response::from_parts(parts, Body::from(err_bytes))
+            }
+        } else {
+            response
+        }
+    } else {
+        response
+    };
+
     // 5. Response processing.
     // Non-streaming responses are buffered so we can (a) record upstream
     // token usage for stats and (b) optionally DLP-redact before sending.
@@ -681,21 +758,18 @@ async fn handle_request(
         .is_some_and(|ct| ct.contains("text/event-stream"));
 
     let response = if is_streaming {
-        if state.dlp_scanner.scan_responses() {
-            debug!(
-                method = %method,
-                path = %path,
-                virtual_key = %virtual_key,
-                "Streaming response (SSE) — wrapping with DLP SSE scanner"
-            );
-            let (parts, body) = response.into_parts();
-            let dlp_body =
-                crate::translate::wrap_body_with_dlp_sse_stream(body, state.dlp_scanner.clone());
-            Response::from_parts(parts, dlp_body)
-        } else {
-            trace!("Streaming response, DLP response scanning disabled");
-            response
-        }
+        trace!(
+            method = %method,
+            path = %path,
+            "Streaming response (SSE) — wrapping for DLP + token counting"
+        );
+        let (parts, body) = response.into_parts();
+        let dlp_body = crate::translate::wrap_body_with_dlp_sse_stream(
+            body,
+            state.dlp_scanner.clone(),
+            state.stats.clone(),
+        );
+        Response::from_parts(parts, dlp_body)
     } else {
         let (parts, body) = response.into_parts();
         let body_bytes = body
@@ -1009,6 +1083,31 @@ fn build_rewritten_uri(original: &Uri, new_path: &str) -> Result<Uri, String> {
         .map_err(|e| format!("failed to build rewritten URI: {e}"))
 }
 
+fn ensure_stream_options(body: &[u8], provider: Provider) -> Option<Vec<u8>> {
+    if !matches!(provider, Provider::Openai | Provider::Openrouter) {
+        return None;
+    }
+    let mut json: serde_json::Value = serde_json::from_slice(body).ok()?;
+    let obj = json.as_object_mut()?;
+    if obj.get("stream").and_then(serde_json::Value::as_bool) != Some(true) {
+        return None;
+    }
+    let already_set = obj
+        .get("stream_options")
+        .and_then(serde_json::Value::as_object)
+        .and_then(|so| so.get("include_usage"))
+        .and_then(serde_json::Value::as_bool)
+        == Some(true);
+    if already_set {
+        return None;
+    }
+    obj.insert(
+        "stream_options".to_string(),
+        serde_json::json!({"include_usage": true}),
+    );
+    serde_json::to_vec(&json).ok()
+}
+
 fn error_response(status: StatusCode, message: &str) -> Response {
     let body = serde_json::json!({ "error": message });
     (status, axum::Json(body)).into_response()

src/app/tests.rs

@@ -1971,3 +1971,58 @@ async fn test_email_message_content_endpoint_disabled_returns_not_found() {
     let resp = app.oneshot(req).await.unwrap();
     assert_eq!(resp.status(), StatusCode::NOT_FOUND);
 }
+
+// ========== ensure_stream_options Tests ==========
+
+#[test]
+fn test_ensure_stream_options_injects_for_openai() {
+    use crate::config::Provider;
+    let body = serde_json::json!({
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "hi"}],
+        "stream": true
+    });
+    let result = super::ensure_stream_options(body.to_string().as_bytes(), Provider::Openai);
+    assert!(result.is_some());
+    let parsed: serde_json::Value = serde_json::from_slice(&result.unwrap()).unwrap();
+    assert_eq!(
+        parsed["stream_options"]["include_usage"],
+        serde_json::json!(true)
+    );
+}
+
+#[test]
+fn test_ensure_stream_options_injects_for_openrouter() {
+    use crate::config::Provider;
+    let body = serde_json::json!({"stream": true, "messages": []});
+    let result = super::ensure_stream_options(body.to_string().as_bytes(), Provider::Openrouter);
+    assert!(result.is_some());
+}
+
+#[test]
+fn test_ensure_stream_options_skips_non_streaming() {
+    use crate::config::Provider;
+    let body = serde_json::json!({"model": "gpt-4o", "messages": []});
+    let result = super::ensure_stream_options(body.to_string().as_bytes(), Provider::Openai);
+    assert!(result.is_none());
+}
+
+#[test]
+fn test_ensure_stream_options_skips_anthropic() {
+    use crate::config::Provider;
+    let body = serde_json::json!({"stream": true, "messages": []});
+    let result = super::ensure_stream_options(body.to_string().as_bytes(), Provider::Anthropic);
+    assert!(result.is_none());
+}
+
+#[test]
+fn test_ensure_stream_options_preserves_existing() {
+    use crate::config::Provider;
+    let body = serde_json::json!({
+        "stream": true,
+        "stream_options": {"include_usage": true},
+        "messages": []
+    });
+    let result = super::ensure_stream_options(body.to_string().as_bytes(), Provider::Openai);
+    assert!(result.is_none());
+}

src/stats.rs

@@ -101,7 +101,12 @@ impl Stats {
         let Ok(json) = serde_json::from_slice::<Value>(body) else {
             return;
         };
-        let Some(usage) = json.get("usage") else {
+        // Top-level `usage` (OpenAI, Anthropic message_delta) or nested
+        // under `message.usage` (Anthropic message_start).
+        let usage = json
+            .get("usage")
+            .or_else(|| json.get("message").and_then(|m| m.get("usage")));
+        let Some(usage) = usage else {
             return;
         };