i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(usage): track OpenAI image SSE completion and usage estimates

Browse Source
This commit is contained in:
ZheFox
2026-05-18 01:32:45 +08:00
parent c024c782e4
commit 691ccaaa04
4 changed files with 511 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

116
apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs
View File

@@ -3,19 +3,20 @@ use std::io::Error as IoError;
use std::time::Instant;
use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTelemetry, RequestBody, ResponseBody, StreamFrame,
StreamFramePayload, StreamFrameType, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
ExecutionPlan, ExecutionResult, ExecutionStreamTerminalSummary, ExecutionTelemetry,
RequestBody, ResponseBody, StreamFrame, StreamFramePayload, StreamFrameType,
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
};
use axum::body::Bytes;
use base64::Engine as _;
use chrono::{FixedOffset, Utc};
use futures_util::stream::{self, BoxStream};
use futures_util::StreamExt;
use serde_json::{json, Value};
use serde_json::{json, Map, Value};
use tracing::debug;
use uuid::Uuid;
use crate::ai_serving::api::StreamingStandardTerminalObserver;
use crate::clock::current_unix_secs;
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::{
@@ -109,7 +110,7 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream(
Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err),
};
Ok(Some(ChatGptWebImageStream {
frame_stream: execution_result_frame_stream(&result),
frame_stream: execution_result_frame_stream(plan, &result, report_context),
report_context: report_context.cloned(),
}))
}
@@ -1479,9 +1480,12 @@ fn bytes_execution_result(
}
fn execution_result_frame_stream(
plan: &ExecutionPlan,
result: &ExecutionResult,
report_context: Option<&Value>,
) -> BoxStream<'static, Result<Bytes, IoError>> {
let body = execution_result_body_bytes_lossy(result);
let terminal_summary = chatgpt_web_stream_terminal_summary(plan, result, report_context, &body);
let mut frames = vec![
StreamFrame {
frame_type: StreamFrameType::Headers,
@@ -1520,7 +1524,7 @@ fn execution_result_frame_stream(
}),
},
});
frames.push(StreamFrame::eof());
frames.push(StreamFrame::eof_with_summary(terminal_summary));
stream::iter(
frames
.into_iter()
@@ -1529,6 +1533,81 @@ fn execution_result_frame_stream(
.boxed()
}
fn chatgpt_web_stream_terminal_summary(
plan: &ExecutionPlan,
result: &ExecutionResult,
report_context: Option<&Value>,
body: &[u8],
) -> Option<ExecutionStreamTerminalSummary> {
if !(200..300).contains(&result.status_code) || body.is_empty() {
return None;
}
let observer_context = chatgpt_web_stream_observer_context(plan, report_context);
let mut observer = StreamingStandardTerminalObserver::default();
let mut line_start = 0usize;
for (index, byte) in body.iter().enumerate() {
if *byte != b'\n' {
continue;
}
observer
.push_line(&observer_context, body[line_start..=index].to_vec())
.ok()?;
line_start = index.saturating_add(1);
}
if line_start < body.len() {
observer
.push_line(&observer_context, body[line_start..].to_vec())
.ok()?;
}
observer.finish(&observer_context).ok().flatten()
}
fn chatgpt_web_stream_observer_context(
plan: &ExecutionPlan,
report_context: Option<&Value>,
) -> Value {
let mut context = report_context
.cloned()
.filter(Value::is_object)
.unwrap_or_else(|| json!({}));
let object = context
.as_object_mut()
.expect("observer context should be an object");
object
.entry("provider_api_format".to_string())
.or_insert_with(|| Value::String(plan.provider_api_format.clone()));
object
.entry("client_api_format".to_string())
.or_insert_with(|| Value::String(plan.client_api_format.clone()));
object
.entry("model".to_string())
.or_insert_with(|| Value::String(plan.model_name.clone().unwrap_or_default()));
if !object.contains_key("image_request") {
if let Some(image_request) = chatgpt_web_image_request_context(plan) {
object.insert("image_request".to_string(), image_request);
}
}
context
}
fn chatgpt_web_image_request_context(plan: &ExecutionPlan) -> Option<Value> {
let body = plan.body.json_body.as_ref()?.as_object()?;
let mut image_request = Map::new();
image_request.insert(
"operation".to_string(),
Value::String("generate".to_string()),
);
for key in ["model", "size", "quality", "output_format"] {
if let Some(value) = body.get(key).and_then(Value::as_str).map(str::trim) {
if !value.is_empty() {
image_request.insert(key.to_string(), Value::String(value.to_string()));
}
}
}
Some(Value::Object(image_request))
}
fn telemetry(started_at: Instant, upstream_bytes: u64) -> ExecutionTelemetry {
let elapsed_ms = started_at.elapsed().as_millis() as u64;
ExecutionTelemetry {
@@ -2225,6 +2304,31 @@ data: [DONE]
assert!(decoded_data.contains("\"width\":2"));
assert!(decoded_data.contains("\"height\":3"));
assert!(text.contains("\"type\":\"eof\""));
let eof_frame = text
.lines()
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
.find(|frame| frame.get("type").and_then(Value::as_str) == Some("eof"))
.expect("eof frame should exist");
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("standardized_usage"))
.and_then(|usage| usage.get("dimensions"))
.and_then(|dimensions| dimensions.get("image_count"))
.and_then(Value::as_u64),
Some(1)
);
assert_eq!(
eof_frame
.get("payload")
.and_then(|payload| payload.get("summary"))
.and_then(|summary| summary.get("standardized_usage"))
.and_then(|usage| usage.get("dimensions"))
.and_then(|dimensions| dimensions.get("image_size"))
.and_then(Value::as_str),
Some("1024x1024")
);
handle.abort();
}

87
apps/aether-gateway/src/execution_runtime/fallback.rs
View File

@@ -13,6 +13,16 @@ fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
)
}
fn openai_image_success_disables_local_success_failover(
plan: &ExecutionPlan,
status_code: u16,
) -> bool {
status_code == 200
&& plan
.provider_api_format
.eq_ignore_ascii_case("openai:image")
}
pub(crate) async fn should_retry_next_local_candidate_sync(
state: &AppState,
plan: &ExecutionPlan,
@@ -48,6 +58,10 @@ pub(crate) async fn analyze_local_candidate_failover_sync(
return LocalFailoverAnalysis::use_default();
}
if openai_image_success_disables_local_success_failover(plan, result.status_code) {
return LocalFailoverAnalysis::use_default();
}
resolve_local_failover_analysis_for_attempt(
state,
plan,
@@ -218,6 +232,10 @@ pub(crate) async fn resolve_local_candidate_failover_analysis_stream(
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverAnalysis {
if openai_image_success_disables_local_success_failover(plan, status_code) {
return LocalFailoverAnalysis::use_default();
}
resolve_local_failover_analysis_for_attempt(
state,
plan,
@@ -756,6 +774,75 @@ mod tests {
);
}
#[tokio::test]
async fn stream_success_failover_does_not_retry_openai_image_success() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"pattern": ".*"}
]
}
})));
let mut plan = sample_plan();
plan.provider_api_format = "openai:image".to_string();
assert!(
!should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_image_stream",
Some(&local_report_context),
200,
Some("{\"data\":[{\"b64_json\":\"aGVsbG8=\"}]}"),
)
.await,
"successful OpenAI image responses should not be retried by success failover rules"
);
}
#[tokio::test]
async fn sync_success_failover_does_not_retry_openai_image_success() {
let local_report_context = serde_json::json!({
"candidate_index": 0,
"retry_index": 0,
});
let state = build_state_with_provider_config(Some(serde_json::json!({
"failover_rules": {
"success_failover_patterns": [
{"pattern": ".*"}
]
}
})));
let mut plan = sample_plan();
plan.provider_api_format = "openai:image".to_string();
let result = ExecutionResult {
request_id: "req-1".to_string(),
candidate_id: None,
status_code: 200,
headers: Default::default(),
body: None,
telemetry: None,
error: None,
};
assert!(
!should_retry_next_local_candidate_sync(
&state,
&plan,
"openai_image_sync",
Some(&local_report_context),
&result,
Some("{\"data\":[{\"b64_json\":\"aGVsbG8=\"}]}")
)
.await,
"successful OpenAI image responses should not be retried by success failover rules"
);
}
#[test]
fn resolve_local_failover_policy_reads_provider_rules() {
let state = build_state_with_provider_config(Some(serde_json::json!({

165
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -991,6 +991,12 @@ fn build_sse_body_stream(
}
}
fn stream_chunk_contains_sse_done(chunk: &[u8]) -> bool {
std::str::from_utf8(chunk)
.ok()
.is_some_and(|text| text.lines().any(|line| line.trim() == "data: [DONE]"))
}
async fn next_stream_frame<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
@@ -2026,6 +2032,8 @@ async fn execute_stream_from_frame_stream(
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let mut client_visible_stream_completed =
stream_chunk_contains_sse_done(&prefetched_body_for_report);
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let reached_eof = initial_reached_eof;
@@ -2452,6 +2460,8 @@ async fn execute_stream_from_frame_stream(
);
let rewritten_chunk_len =
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&rewritten_chunk);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_disconnected",
@@ -2464,6 +2474,7 @@ async fn execute_stream_from_frame_stream(
downstream_dropped = true;
break;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
@@ -2578,6 +2589,8 @@ async fn execute_stream_from_frame_stream(
);
let rewritten_chunk_len =
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&rewritten_chunk);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_flush_disconnected",
@@ -2589,6 +2602,7 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes
.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
@@ -2635,6 +2649,8 @@ async fn execute_stream_from_frame_stream(
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
let chunk_completed_stream =
stream_chunk_contains_sse_done(&flushed_chunk);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
@@ -2646,6 +2662,7 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
client_visible_stream_completed |= chunk_completed_stream;
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
@@ -2755,6 +2772,18 @@ async fn execute_stream_from_frame_stream(
),
);
if downstream_dropped && client_visible_stream_completed && terminal_failure.is_none() {
debug!(
event_name = "execution_runtime_stream_downstream_closed_after_done",
log_type = "debug",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway treats downstream close after client-visible SSE DONE as completed"
);
downstream_dropped = false;
}
if downstream_dropped {
debug!(
event_name = "execution_runtime_stream_report_skipped",
@@ -3417,6 +3446,142 @@ mod tests {
);
}
#[tokio::test]
async fn image_stream_downstream_close_after_done_is_recorded_success() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let plan = ExecutionPlan {
request_id: "req-image-done-close-success".into(),
candidate_id: Some("cand-image-done-close-success".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/v1/images/generations".into(),
headers: BTreeMap::from([("accept".into(), "text/event-stream".into())]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-2",
"prompt": "draw a small image",
"stream": true
})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:image".into(),
model_name: Some("gpt-image-2".into()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.output_item.done\\ndata: {\\\"type\\\":\\\"response.output_item.done\\\",\\\"output_index\\\":0,\\\"item\\\":{\\\"id\\\":\\\"ig_1\\\",\\\"type\\\":\\\"image_generation_call\\\",\\\"result\\\":\\\"aGVsbG8=\\\"}}\\n\\nevent: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_1\\\",\\\"model\\\":\\\"gpt-image-2\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":null}}\\n\\n\"}}\n",
));
std::future::pending::<()>().await;
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-image-done-close-success",
&test_decision(),
"openai_chat_stream",
Some("openai_chat_stream_success".to_string()),
Some(json!({
"request_id": "req-image-done-close-success",
"candidate_id": "cand-image-done-close-success",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:image",
"client_api_format": "openai:chat",
"image_request": {
"size": "1024x1024",
"quality": "medium"
}
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
let mut body_stream = response.into_body().into_data_stream();
let mut body = Vec::new();
tokio::time::timeout(Duration::from_secs(1), async {
while !String::from_utf8_lossy(&body).contains("data: [DONE]") {
let chunk = body_stream
.next()
.await
.expect("body should yield until done")
.expect("chunk should be ok");
body.extend_from_slice(&chunk);
}
})
.await
.expect("final DONE should arrive");
drop(body_stream);
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
.list_by_request_id("req-image-done-close-success")
.await
.expect("request candidates should read");
if candidates
.first()
.is_some_and(|candidate| candidate.status == RequestCandidateStatus::Success)
{
break candidates;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("candidate should be marked success");
assert_eq!(candidates[0].status_code, Some(200));
let stored_usage = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let usage = usage_repository
.find_by_request_id("req-image-done-close-success")
.await
.expect("usage should read");
if usage
.as_ref()
.is_some_and(|usage| usage.status == "completed")
{
break usage.expect("completed usage should exist");
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("usage should be marked completed");
assert_eq!(stored_usage.status_code, Some(200));
assert!(stored_usage.total_tokens > 0);
}
#[tokio::test]
async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
) {

149
crates/aether-usage-runtime/src/write.rs
View File

@@ -619,6 +619,10 @@ fn build_terminal_usage_event_from_seed_impl(
}
}
if matches!(event_type, UsageEventType::Completed) {
apply_completed_image_usage_estimate(&mut data);
}
if matches!(event_type, UsageEventType::Cancelled) {
apply_cancelled_usage_estimate(&mut data);
}
@@ -2290,6 +2294,63 @@ fn apply_cancelled_usage_estimate(data: &mut UsageEventData) {
}
}
fn apply_completed_image_usage_estimate(data: &mut UsageEventData) {
if !usage_event_data_is_image(data) {
return;
}
if data
.response_body
.as_ref()
.and_then(extract_token_counts_from_value)
.is_some()
{
return;
}
let request_usage = data
.provider_request_body
.as_ref()
.or(data.request_body.as_ref())
.and_then(estimate_request_usage);
if positive_tokens(data.input_tokens) == 0 {
if let Some(usage) = request_usage.as_ref() {
data.input_tokens = Some(usage.input_tokens);
}
}
apply_cancelled_request_cache_estimate(data, request_usage.as_ref());
if positive_tokens(data.total_tokens) == 0 {
let total_tokens =
positive_tokens(data.input_tokens).saturating_add(positive_tokens(data.output_tokens));
if total_tokens > 0 {
data.total_tokens = Some(total_tokens);
}
}
}
fn usage_event_data_is_image(data: &UsageEventData) -> bool {
data.request_type
.as_deref()
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
|| data
.endpoint_kind
.as_deref()
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
|| data
.provider_endpoint_kind
.as_deref()
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
|| data
.endpoint_api_format
.as_deref()
.and_then(infer_endpoint_kind)
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
|| data
.api_format
.as_deref()
.and_then(infer_endpoint_kind)
.is_some_and(|value| value.eq_ignore_ascii_case("image"))
}
fn apply_cancelled_request_cache_estimate(
data: &mut UsageEventData,
request_usage: Option<&EstimatedRequestUsage>,
@@ -3642,6 +3703,94 @@ mod tests {
assert!(event.data.client_response_body.is_none());
}
#[test]
fn completed_image_usage_estimates_request_tokens_when_provider_usage_is_missing() {
let plan = ExecutionPlan {
request_id: "req-image-completed-estimate-1".to_string(),
candidate_id: Some("cand-image-completed-estimate-1".to_string()),
provider_name: Some("OpenAI".to_string()),
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
method: "POST".to_string(),
url: "https://example.com/v1/images/generations".to_string(),
headers: BTreeMap::new(),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-2",
"prompt": "draw a small red cube on a clean desk",
"size": "1024x1024",
"quality": "medium"
})),
stream: true,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:image".to_string(),
model_name: Some("gpt-image-2".to_string()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let mut standardized_usage = StandardizedUsage::new();
standardized_usage.request_count = 1;
standardized_usage
.dimensions
.insert("image_count".to_string(), json!(1));
standardized_usage
.dimensions
.insert("image_size".to_string(), json!("1024x1024"));
standardized_usage
.dimensions
.insert("image_quality".to_string(), json!("medium"));
let payload = GatewayStreamReportRequest {
trace_id: "trace-image-completed-estimate-1".to_string(),
report_kind: "openai_chat_stream_success".to_string(),
report_context: Some(json!({
"client_api_format": "openai:chat",
"provider_api_format": "openai:image",
"image_request": {
"size": "1024x1024",
"quality": "medium"
}
})),
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: Some(UsageBodyCaptureState::None),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(standardized_usage),
finish_reason: Some("stop".to_string()),
response_id: Some("resp_image_estimate_1".to_string()),
model: Some("gpt-image-2".to_string()),
observed_finish: true,
unknown_event_count: 0,
parser_error: None,
}),
telemetry: None,
};
let event =
build_stream_terminal_usage_event(&plan, payload.report_context.as_ref(), &payload)
.expect("usage event should build");
assert_eq!(event.event_type, UsageEventType::Completed);
assert!(event.data.input_tokens.unwrap_or_default() > 0);
assert_eq!(event.data.output_tokens.unwrap_or_default(), 0);
assert_eq!(event.data.total_tokens, event.data.input_tokens);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("dimensions"))
.and_then(|dimensions| dimensions.get("image_count"))
.and_then(Value::as_i64),
Some(1)
);
}
#[test]
fn stream_terminal_usage_prefers_more_complete_provider_chunks_usage() {
let plan = ExecutionPlan {