diff --git a/src/app.rs b/src/app.rs index 496f393..c062cd3 100644 --- a/src/app.rs +++ b/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 { .map_err(|e| format!("failed to build rewritten URI: {e}")) } +fn ensure_stream_options(body: &[u8], provider: Provider) -> Option> { + 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() diff --git a/src/app/tests.rs b/src/app/tests.rs index e3520a1..bb2f895 100644 --- a/src/app/tests.rs +++ b/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()); +} diff --git a/src/stats.rs b/src/stats.rs index f9c166b..7e80dd0 100644 --- a/src/stats.rs +++ b/src/stats.rs @@ -101,7 +101,12 @@ impl Stats { let Ok(json) = serde_json::from_slice::(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; }; diff --git a/src/translate.rs b/src/translate.rs index 5763615..cd5f652 100644 --- a/src/translate.rs +++ b/src/translate.rs @@ -1,4 +1,5 @@ use crate::dlp::DlpScanner; +use crate::stats::Stats; use axum::body::Body; use bytes::{Bytes, BytesMut}; use futures_util::Stream; @@ -18,7 +19,14 @@ pub enum TranslateError { } /// Fields that are compatible between chat/completions and responses API. -const PASSTHROUGH_FIELDS: &[&str] = &["model", "stream", "temperature", "top_p", "stop"]; +const PASSTHROUGH_FIELDS: &[&str] = &[ + "model", + "stream", + "stream_options", + "temperature", + "top_p", + "stop", +]; /// Fields that must be stripped from chat/completions requests (not supported by responses API). const STRIP_FIELDS: &[&str] = &[ @@ -507,16 +515,18 @@ pub fn redact_sse_data_line(line: &str, scanner: &DlpScanner) -> String { ) } -/// Stream adapter that applies DLP redaction to SSE data lines. +/// Stream adapter that applies DLP redaction to SSE data lines and +/// extracts upstream token usage for stats. pub struct DlpSseStream { inner: Pin> + Send>>, buffer: BytesMut, scanner: Arc, + stats: Arc, output_buffer: Vec, } impl DlpSseStream { - pub fn new(body: Body, scanner: Arc) -> Self { + pub fn new(body: Body, scanner: Arc, stats: Arc) -> Self { use futures_util::StreamExt; use http_body_util::BodyStream; @@ -531,6 +541,7 @@ impl DlpSseStream { inner: Box::pin(stream), buffer: BytesMut::new(), scanner, + stats, output_buffer: Vec::new(), } } @@ -549,6 +560,16 @@ impl DlpSseStream { continue; } + // Extract token usage from SSE data events. This is a no-op + // for the vast majority of events that don't carry a `usage` + // key; only the 1–3 events per stream that do will actually + // update the atomic counters. + if let Some(json_str) = line.strip_prefix("data: ") { + if !json_str.starts_with("[DONE]") { + self.stats.record_tokens_from_usage(json_str.as_bytes()); + } + } + let redacted = redact_sse_data_line(&line, &self.scanner); self.output_buffer.extend_from_slice(redacted.as_bytes()); self.output_buffer.extend_from_slice(b"\n"); @@ -578,6 +599,11 @@ impl Stream for DlpSseStream { let remaining = std::mem::take(&mut this.buffer); let line = String::from_utf8_lossy(&remaining).trim().to_string(); if !line.is_empty() { + if let Some(json_str) = line.strip_prefix("data: ") { + if !json_str.starts_with("[DONE]") { + this.stats.record_tokens_from_usage(json_str.as_bytes()); + } + } let redacted = redact_sse_data_line(&line, &this.scanner); return Poll::Ready(Some(Ok(Bytes::from(format!("{redacted}\n"))))); } @@ -590,9 +616,14 @@ impl Stream for DlpSseStream { } } -/// Wrap a Body in a DlpSseStream for streaming DLP redaction. -pub fn wrap_body_with_dlp_sse_stream(body: Body, scanner: Arc) -> Body { - Body::from_stream(DlpSseStream::new(body, scanner)) +/// Wrap a Body in a DlpSseStream for streaming DLP redaction and +/// token usage extraction. +pub fn wrap_body_with_dlp_sse_stream( + body: Body, + scanner: Arc, + stats: Arc, +) -> Body { + Body::from_stream(DlpSseStream::new(body, scanner, stats)) } #[cfg(test)] @@ -705,6 +736,7 @@ mod tests { "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "stream": true, + "stream_options": {"include_usage": true}, "temperature": 0.7, "top_p": 0.9, "stop": ["\n"] @@ -714,6 +746,10 @@ mod tests { assert_eq!(parsed["model"], "gpt-4o-mini"); assert_eq!(parsed["stream"], true); + assert_eq!( + parsed["stream_options"], + serde_json::json!({"include_usage": true}) + ); assert_eq!(parsed["temperature"], 0.7); assert_eq!(parsed["top_p"], 0.9); assert_eq!(parsed["stop"], serde_json::json!(["\n"])); @@ -992,4 +1028,76 @@ mod tests { "Non-data lines pass through unchanged" ); } + + #[test] + fn test_sse_stream_counts_openai_usage() { + let stats = Arc::new(Stats::new(None)); + let scanner = Arc::new(DlpScanner::new(&[], false).unwrap()); + let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ + data: {\"choices\":[],\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\ + data: [DONE]\n"; + let body = Body::from(sse); + let wrapped = wrap_body_with_dlp_sse_stream(body, scanner, stats.clone()); + + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + use http_body_util::BodyExt; + wrapped.collect().await.unwrap(); + }); + + let snap = stats.snapshot(); + assert_eq!(snap.prompt_tokens_total, 10); + assert_eq!(snap.completion_tokens_total, 5); + assert_eq!(snap.total_tokens_total, 15); + } + + #[test] + fn test_sse_stream_counts_anthropic_usage() { + let stats = Arc::new(Stats::new(None)); + let scanner = Arc::new(DlpScanner::new(&[], false).unwrap()); + let sse = "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":7}}}\n\ + data: {\"type\":\"content_block_delta\",\"delta\":{\"text\":\"hi\"}}\n\ + data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":3}}\n"; + let body = Body::from(sse); + let wrapped = wrap_body_with_dlp_sse_stream(body, scanner, stats.clone()); + + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + use http_body_util::BodyExt; + wrapped.collect().await.unwrap(); + }); + + let snap = stats.snapshot(); + assert_eq!(snap.prompt_tokens_total, 7); + assert_eq!(snap.completion_tokens_total, 3); + assert_eq!(snap.total_tokens_total, 10); + } + + #[test] + fn test_sse_stream_ignores_events_without_usage() { + let stats = Arc::new(Stats::new(None)); + let scanner = Arc::new(DlpScanner::new(&[], false).unwrap()); + let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\ + data: {\"choices\":[{\"delta\":{\"content\":\" world\"}}]}\n\ + data: [DONE]\n"; + let body = Body::from(sse); + let wrapped = wrap_body_with_dlp_sse_stream(body, scanner, stats.clone()); + + let rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + rt.block_on(async { + use http_body_util::BodyExt; + wrapped.collect().await.unwrap(); + }); + + let snap = stats.snapshot(); + assert_eq!(snap.prompt_tokens_total, 0); + assert_eq!(snap.completion_tokens_total, 0); + assert_eq!(snap.total_tokens_total, 0); + } }