boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: optimize proxy streaming and database access

Browse Source
This commit is contained in:
boos
2026-06-07 00:00:24 +08:00
parent f3d2d50f9a
commit aaba6f152e
13 changed files with 960 additions and 492 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/proxy/mod.rs
+470 -381
View File
@@ -7,9 +7,11 @@
use crate::{
auth::api_key::{authorize_downstream_token, DownstreamAuthResult},
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
proxy::conductor::{
AttemptResult, ExecuteInput, ExecuteResult, ProxyConductor, SelectedChannel,
classify_failure, AttemptFailureAction, AttemptResult, ExecuteInput, ExecuteResult,
ProxyConductor, SelectedChannel, SuccessMetrics,
},
proxy::executors::{self, UpstreamRoute},
proxy::token_router::TokenRouter,
@@ -17,15 +19,16 @@ use crate::{
AppState,
};
use axum::{
body::Body,
extract::{Multipart, Query, State},
http::HeaderMap,
response::{IntoResponse, Response, Sse},
http::{header, HeaderMap},
response::{IntoResponse, Response},
routing::{get, post},
Json, Router,
};
use futures::{stream, FutureExt, StreamExt};
use futures::{FutureExt, StreamExt};
use serde_json::{json, Value};
use std::{convert::Infallible, sync::Arc, time::Instant};
use std::{future::Future, sync::Arc, time::Instant};
use tracing::info;
// ── Auth types ────────────────────────────────────────────────────────
@@ -111,14 +114,7 @@ async fn messages_handler(
let log_content = proxy_log_content_from_request(&body);
if canonical.stream {
let chunks =
execute_messages_stream_with_conductor(state, key_id, model, canonical, log_content)
.await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
execute_messages_stream_with_conductor(state, key_id, model, canonical, log_content).await
} else {
let result =
execute_messages_with_conductor(state, key_id, model, canonical, log_content).await?;
@@ -154,12 +150,7 @@ async fn responses_handler(
.unwrap_or(false);
if stream {
let chunks = execute_responses_stream_with_conductor(state, key_id, model, body).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
execute_responses_stream_with_conductor(state, key_id, model, body).await
} else {
let result = execute_responses_with_conductor(state, key_id, model, body).await?;
Ok(Json(result).into_response())
@@ -188,13 +179,7 @@ async fn chat_completions(
.unwrap_or(false);
if stream {
let chunks =
execute_chat_completions_stream_with_conductor(state, key_id, model, body).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
execute_chat_completions_stream_with_conductor(state, key_id, model, body).await
} else {
let result = execute_chat_completions_with_conductor(state, key_id, model, body).await?;
Ok(Json(result).into_response())
@@ -635,16 +620,36 @@ async fn check_auth(state: &AppState, headers: &HeaderMap) -> Result<Option<i64>
match result {
DownstreamAuthResult::Managed { key } => {
if !key.enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
UsageTracker::new(state.db.clone()).check_rate_limit(key.id).await?;
check_downstream_key_limits(&key)?;
Ok(Some(key.id))
}
DownstreamAuthResult::Global { .. } => Ok(None),
}
}
fn check_downstream_key_limits(key: &DownstreamApiKey) -> AppResult<()> {
if !key.enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
if let Some(max_requests) = key.max_requests {
if max_requests <= 0 || key.used_requests.unwrap_or(0) + 1 > max_requests {
return Err(AppError::RateLimitedWithRetry {
message: "下游密钥请求额度已用完".into(),
retry_after_secs: 60,
});
}
}
if let Some(max_cost) = key.max_cost {
if max_cost <= 0.0 || key.used_cost.unwrap_or(0.0) >= max_cost {
return Err(AppError::RateLimitedWithRetry {
message: "下游密钥成本额度已用完".into(),
retry_after_secs: 60,
});
}
}
Ok(())
}
/// Best-effort usage logging. Only fires when the request was authenticated
/// via a managed downstream key (key_id is Some). Failures are logged but
/// never bubbled — billing must not break a successful proxy response.
@@ -1014,6 +1019,317 @@ fn ensure_stream_usage(mut body: Value) -> Value {
body
}
#[derive(Debug)]
struct OpenedStream {
selected: SelectedChannel,
attempts: usize,
response: reqwest::Response,
}
#[derive(Debug)]
struct StreamOpenFailure {
selected: Option<SelectedChannel>,
attempts: usize,
status: Option<u16>,
raw_error_text: Option<String>,
}
async fn open_stream_with_failover<F, Fut>(
state: &AppState,
requested_model: &str,
mut open_attempt: F,
) -> Result<OpenedStream, StreamOpenFailure>
where
F: FnMut(AppState, SelectedChannel) -> Fut,
Fut: Future<Output = AppResult<reqwest::Response>>,
{
let conductor = build_conductor(state);
let mut selected = (conductor.select_channel)(requested_model.to_string(), None)
.await
.ok_or_else(|| StreamOpenFailure {
selected: None,
attempts: 0,
status: None,
raw_error_text: Some(format!("No route configured for model '{requested_model}'")),
})?;
let mut excluded = Vec::new();
let mut attempts = 0usize;
loop {
attempts += 1;
match open_attempt(state.clone(), selected.clone()).await {
Ok(response) => {
return Ok(OpenedStream {
selected,
attempts,
response,
});
}
Err(err) => {
let status = upstream_status_from_error(&err);
let action = classify_failure(status, Some(&err.to_string()));
if let Err(record_err) = TokenRouter::new(state.db.clone())
.record_failure(
selected.channel_id,
crate::proxy::conductor::FailureInfo {
action: action.as_str().to_string(),
status,
raw_error_text: Some(err.to_string()),
},
)
.await
{
tracing::warn!(%record_err, channel_id = selected.channel_id, "failed to record streaming proxy failure");
}
if matches!(action, AttemptFailureAction::RetrySameChannel) && attempts < 5 {
continue;
}
let raw_error_text = Some(err.to_string());
if matches!(action, AttemptFailureAction::Terminal | AttemptFailureAction::Stop) || attempts >= 5 {
return Err(StreamOpenFailure {
selected: Some(selected),
attempts,
status,
raw_error_text,
});
}
excluded.push(selected.channel_id);
selected = (conductor.select_next_channel)(
requested_model.to_string(),
excluded.clone(),
None,
)
.await
.ok_or_else(|| StreamOpenFailure {
selected: Some(selected.clone()),
attempts,
status,
raw_error_text,
})?;
}
}
}
}
fn sse_passthrough_response(
state: AppState,
opened: OpenedStream,
key_id: Option<i64>,
requested_model: String,
started_at: Instant,
log_content: ProxyLogContent,
usage_extractor: fn(&str) -> Option<(i64, i64)>,
) -> AppResult<Response> {
let selected = opened.selected;
let attempts = opened.attempts;
let billing_state = state.clone();
let stream = async_stream::stream! {
let mut usage = StreamUsageAccumulator::default();
let mut bytes_stream = opened.response.bytes_stream();
let mut stream_error: Option<String> = None;
while let Some(chunk) = bytes_stream.next().await {
let bytes = match chunk {
Ok(bytes) => bytes,
Err(err) => {
let message = format!("upstream stream error: {err}");
stream_error = Some(message.clone());
let event = format!(
"event: error\ndata: {}\n\n",
serde_json::json!({ "error": { "message": message } })
);
yield Ok::<bytes::Bytes, std::io::Error>(bytes::Bytes::from(event));
break;
}
};
usage.ingest_bytes(&bytes, usage_extractor);
yield Ok::<bytes::Bytes, std::io::Error>(bytes);
}
usage.finish(usage_extractor);
let prompt_tokens = usage.prompt_tokens;
let completion_tokens = usage.completion_tokens;
let cost = estimate_request_cost(&selected, prompt_tokens, completion_tokens);
if let Some(error) = stream_error {
if let Err(err) = TokenRouter::new(billing_state.db.clone())
.record_failure(
selected.channel_id,
crate::proxy::conductor::FailureInfo {
action: AttemptFailureAction::Failover.as_str().to_string(),
status: None,
raw_error_text: Some(error.clone()),
},
)
.await
{
tracing::warn!(%err, channel_id = selected.channel_id, "failed to record streaming proxy failure");
}
record_proxy_log_failure(
&billing_state,
key_id,
Some(&selected),
&requested_model,
None,
Some(elapsed_ms_since(started_at)),
Some(&error),
attempts,
true,
log_content,
)
.await;
} else {
if let Err(err) = TokenRouter::new(billing_state.db.clone())
.record_success(
selected.channel_id,
SuccessMetrics {
latency_ms: Some(started_at.elapsed().as_millis().min(u64::MAX as u128) as u64),
cost: Some(cost),
metadata: Default::default(),
},
)
.await
{
tracing::warn!(%err, channel_id = selected.channel_id, "failed to record streaming proxy success");
}
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&requested_model,
Some(elapsed_ms_since(started_at)),
prompt_tokens,
completion_tokens,
cost,
attempts,
true,
log_content,
)
.await;
record_billing(
&billing_state,
key_id,
&requested_model,
prompt_tokens,
completion_tokens,
cost,
)
.await;
}
};
Response::builder()
.header(header::CONTENT_TYPE, "text/event-stream; charset=utf-8")
.header(header::CACHE_CONTROL, "no-cache, no-transform")
.header(header::CONNECTION, "keep-alive")
.header("x-accel-buffering", "no")
.body(Body::from_stream(stream))
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))
}
#[derive(Default)]
struct StreamUsageAccumulator {
buffer: String,
prompt_tokens: i64,
completion_tokens: i64,
}
impl StreamUsageAccumulator {
fn ingest_bytes(
&mut self,
bytes: &bytes::Bytes,
usage_extractor: fn(&str) -> Option<(i64, i64)>,
) {
self.buffer.push_str(&String::from_utf8_lossy(bytes));
while let Some(newline) = self.buffer.find('\n') {
let line = self.buffer[..newline].trim_end_matches('\r').to_string();
self.buffer.drain(..=newline);
self.ingest_line(&line, usage_extractor);
}
if self.buffer.len() > 64 * 1024 {
self.buffer.clear();
}
}
fn finish(&mut self, usage_extractor: fn(&str) -> Option<(i64, i64)>) {
if !self.buffer.trim().is_empty() {
let line = std::mem::take(&mut self.buffer);
self.ingest_line(line.trim_end_matches('\r'), usage_extractor);
}
}
fn ingest_line(&mut self, line: &str, usage_extractor: fn(&str) -> Option<(i64, i64)>) {
if let Some((prompt, completion)) = usage_extractor(line) {
if prompt > 0 {
self.prompt_tokens = prompt;
}
if completion > 0 {
self.completion_tokens = completion;
}
}
}
}
fn openai_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let (prompt, completion) = extract_usage(&value);
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn anthropic_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let mut prompt = 0;
let mut completion = 0;
for holder in [
value.get("message"),
value.get("delta"),
Some(&value),
value.get("response"),
]
.into_iter()
.flatten()
{
let (candidate_prompt, candidate_completion) = extract_usage(holder);
if candidate_prompt > 0 {
prompt = candidate_prompt;
}
if candidate_completion > 0 {
completion = candidate_completion;
}
}
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn responses_stream_usage_from_line(line: &str) -> Option<(i64, i64)> {
let payload = line
.strip_prefix("data: ")
.or_else(|| line.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let usage_holder = value.get("response").unwrap_or(&value);
let (prompt, completion) = extract_usage(usage_holder);
(prompt != 0 || completion != 0).then_some((prompt, completion))
}
fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone();
let select_db = db.clone();
@@ -1079,10 +1395,12 @@ async fn execute_messages_with_conductor(
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let mut request = request.clone();
request.model = ctx.selected.actual_model.clone();
async move {
match executors::execute_anthropic_messages(
state,
selected_channel_to_route(ctx.selected),
request,
)
@@ -1174,108 +1492,57 @@ async fn execute_messages_stream_with_conductor(
model: String,
request: crate::transformers::canonical::CanonicalRequest,
log_content: ProxyLogContent,
) -> AppResult<Vec<String>> {
) -> Result<Response, AppError> {
let started_at = Instant::now();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let mut request = request.clone();
request.model = ctx.selected.actual_model.clone();
async move {
match executors::execute_anthropic_messages_stream(
selected_channel_to_route(ctx.selected),
request,
)
.await
{
Ok(chunks) => AttemptResult::Success {
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let request_for_attempt = request.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let mut request = request_for_attempt.clone();
async move {
request.model = selected.actual_model.clone();
executors::execute_anthropic_messages_stream_response(
state,
selected_channel_to_route(selected),
request,
)
.await
}
})
.await;
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.filter_map(|chunk| chunk.as_str().map(str::to_string))
.collect::<Vec<_>>(),
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
log_content,
anthropic_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&billing_state,
&state,
key_id,
selected.as_ref(),
&billing_model,
status,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
failure.raw_error_text.as_deref(),
failure.attempts,
true,
log_content,
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
};
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
0,
0,
cost,
attempts,
true,
log_content,
)
.await;
record_billing(&billing_state, key_id, &billing_model, 0, 0, cost).await;
Ok(chunks)
}
}
async fn execute_chat_completions_with_conductor(
@@ -1389,148 +1656,57 @@ async fn execute_chat_completions_stream_with_conductor(
key_id: Option<i64>,
model: String,
body: Value,
) -> AppResult<Vec<String>> {
) -> Result<Response, AppError> {
let started_at = Instant::now();
let request_log_body = body.clone();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let body =
ensure_stream_usage(body_with_selected_model(body.clone(), &ctx.selected));
async move {
match executors::execute_chat_completions_stream(
state,
selected_channel_to_route(ctx.selected),
body,
)
.await
{
Ok(chunks) => AttemptResult::Success {
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = ensure_stream_usage(body_with_selected_model(body_for_attempt.clone(), &selected));
async move {
executors::execute_chat_completions_stream_response(
state,
selected_channel_to_route(selected),
body,
)
.await
}
})
.await;
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
)
})
})
.collect::<AppResult<Vec<String>>>()?,
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
proxy_log_content_from_request(&request_log_body),
openai_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&billing_state,
&state,
key_id,
selected.as_ref(),
&billing_model,
status,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
failure.raw_error_text.as_deref(),
failure.attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
};
// Streaming responses don't include a usage block in the aggregate.
// OpenAI emits usage on the final `data: {...}` chunk only when
// `stream_options.include_usage = true` was set; parse it best-effort.
let (prompt, completion) = chunks
.iter()
.rev()
.find_map(|chunk| {
let payload = chunk
.strip_prefix("data: ")
.or_else(|| chunk.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
let (p, c) = extract_usage(&value);
if p == 0 && c == 0 {
None
} else {
Some((p, c))
}
})
.unwrap_or((0, 0));
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
cost,
attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
cost,
)
.await;
Ok(chunks)
}
}
async fn execute_responses_with_conductor(
@@ -1644,145 +1820,58 @@ async fn execute_responses_stream_with_conductor(
key_id: Option<i64>,
model: String,
body: Value,
) -> AppResult<Vec<String>> {
) -> Result<Response, AppError> {
let started_at = Instant::now();
let request_log_body = body.clone();
let conductor = build_conductor(&state);
let billing_state = state.clone();
let billing_model = model.clone();
let result = conductor
.execute(ExecuteInput {
requested_model: model,
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let body = body_with_selected_model(body.clone(), &ctx.selected);
async move {
match executors::execute_streaming_request(
selected_channel_to_route(ctx.selected),
"/responses",
body,
)
.await
{
Ok(chunks) => AttemptResult::Success {
response: Value::Array(chunks.into_iter().map(Value::String).collect()),
latency_ms: None,
cost: None,
},
Err(err) => AttemptResult::Failure {
action: crate::proxy::conductor::AttemptFailureAction::Failover,
status: upstream_status_from_error(&err),
raw_error_text: Some(err.to_string()),
},
}
}
.boxed()
}),
on_terminal_failure: None,
})
.await;
let body_for_attempt = body.clone();
let opened = open_stream_with_failover(&state, &model, move |state, selected| {
let body = body_with_selected_model(body_for_attempt.clone(), &selected);
async move {
executors::execute_streaming_response(
state,
selected_channel_to_route(selected),
"/responses",
body,
)
.await
}
})
.await;
let (selected, attempts, chunks) = match result {
ExecuteResult::Success {
selected,
response,
attempts,
} => match response {
Value::Array(chunks) => (
selected,
attempts,
chunks
.into_iter()
.map(|chunk| {
chunk.as_str().map(str::to_string).ok_or_else(|| {
AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
)
})
})
.collect::<AppResult<Vec<String>>>()?,
),
_ => {
return Err(AppError::UpstreamError(
"Invalid streaming response from conductor".into(),
))
}
},
ExecuteResult::Failure {
reason,
selected,
status,
raw_error_text,
attempts,
} => {
match opened {
Ok(opened) => sse_passthrough_response(
state,
opened,
key_id,
model,
started_at,
proxy_log_content_from_request(&request_log_body),
responses_stream_usage_from_line,
),
Err(failure) => {
record_proxy_log_failure(
&billing_state,
&state,
key_id,
selected.as_ref(),
&billing_model,
status,
failure.selected.as_ref(),
&model,
failure.status,
Some(elapsed_ms_since(started_at)),
raw_error_text.as_deref(),
attempts,
failure.raw_error_text.as_deref(),
failure.attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
return Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
status,
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
)));
Err(AppError::UpstreamError(format!(
"Proxy conductor failed after {} attempts (stream open, status: {:?}): {}",
failure.attempts,
failure.status,
failure
.raw_error_text
.unwrap_or_else(|| "no upstream error text".to_string())
)))
}
};
// Responses API streams `response.completed` events that contain a
// `usage` object with input_tokens / output_tokens. extract_usage already
// handles both naming conventions.
let (prompt, completion) = chunks
.iter()
.rev()
.find_map(|chunk| {
let payload = chunk.strip_prefix("data: ")?;
if payload.trim() == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
// Responses API nests usage under `response.usage`.
let usage_holder = value.get("response").unwrap_or(&value);
let (p, c) = extract_usage(usage_holder);
if p == 0 && c == 0 {
None
} else {
Some((p, c))
}
})
.unwrap_or((0, 0));
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
&selected,
&billing_model,
Some(elapsed_ms_since(started_at)),
prompt,
completion,
cost,
attempts,
true,
proxy_log_content_from_request(&request_log_body),
)
.await;
record_billing(
&billing_state,
key_id,
&billing_model,
prompt,
completion,
cost,
)
.await;
Ok(chunks)
}
}
#[cfg(test)]