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

fix(usage): 修复含失败信号的 pending 请求被误判为活跃的问题

Browse Source 
- pending/streaming 状态含 status_code >= 400 或 error_message 时视为 failed
- 候选状态派生改为 failed 优先于 stale pending/streaming
- 活跃请求轮询时保留失败信号字段的更新
- 图片流支持 partial_image 事件转发及从 response.completed 读取最终图片
This commit is contained in:
fawney19
2026-04-24 03:07:11 +08:00
parent 581bc03d4e
commit 5148370253
16 changed files with 394 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
apps/aether-gateway/src/ai_pipeline/finalize/internal/stream_rewrite.rs
View File

@@ -106,6 +106,7 @@ struct OpenAiImageStreamState {
buffered: Vec<u8>,
latest_image: Option<OpenAiImageFrame>,
emitted_partial_count: u64,
saw_upstream_partial: bool,
}
#[derive(Clone)]
@@ -166,12 +167,55 @@ impl OpenAiImageStreamState {
.or(event_name.as_deref())
.unwrap_or_default();
match event_type {
"response.image_generation_call.partial_image" => {
self.handle_image_generation_partial(report_context, &event)
}
"response.output_item.done" => self.handle_output_item_done(report_context, &event),
"response.completed" => self.handle_completed(report_context, &event),
_ => Ok(Vec::new()),
}
}
fn handle_image_generation_partial(
&mut self,
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, GatewayError> {
if requested_partial_images(report_context) == 0 {
return Ok(Vec::new());
}
let Some(result) = event
.get("partial_image_b64")
.or_else(|| event.get("b64_json"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
return Ok(Vec::new());
};
let partial_image_index = event
.get("partial_image_index")
.or_else(|| event.get("output_index"))
.and_then(Value::as_u64)
.unwrap_or(self.emitted_partial_count);
self.emitted_partial_count = self
.emitted_partial_count
.max(partial_image_index.saturating_add(1));
self.saw_upstream_partial = true;
self.latest_image = Some(OpenAiImageFrame {
b64_json: result.to_string(),
});
encode_json_sse(
Some(image_partial_event_name(report_context)),
&serde_json::json!({
"type": image_partial_event_name(report_context),
"b64_json": result,
"partial_image_index": partial_image_index,
}),
)
}
fn handle_output_item_done(
&mut self,
report_context: &Value,
@@ -193,7 +237,7 @@ impl OpenAiImageStreamState {
b64_json: result.to_string(),
});
if requested_partial_images(report_context) == 0 {
if requested_partial_images(report_context) == 0 || self.saw_upstream_partial {
return Ok(Vec::new());
}
@@ -218,6 +262,13 @@ impl OpenAiImageStreamState {
report_context: &Value,
event: &Value,
) -> Result<Vec<u8>, GatewayError> {
if self.latest_image.is_none() {
if let Some(result) = completed_response_image_result(event) {
self.latest_image = Some(OpenAiImageFrame {
b64_json: result.to_string(),
});
}
}
let Some(latest_image) = self.latest_image.clone() else {
return Ok(Vec::new());
};
@@ -244,6 +295,19 @@ impl OpenAiImageStreamState {
}
}
fn completed_response_image_result(event: &Value) -> Option<&str> {
event
.get("response")
.and_then(|value| value.get("output"))
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter(|item| item.get("type").and_then(Value::as_str) == Some("image_generation_call"))
.filter_map(|item| item.get("result").and_then(Value::as_str))
.map(str::trim)
.find(|value| !value.is_empty())
}
fn requested_partial_images(report_context: &Value) -> u64 {
report_context
.get("image_request")

85
apps/aether-gateway/src/ai_pipeline/finalize/tests_stream.rs
View File

@@ -172,6 +172,91 @@ fn openai_image_stream_rewriter_emits_completed_event_for_generate() {
assert!(rewriter.finish().expect("finish should succeed").is_empty());
}
#[test]
fn openai_image_stream_rewriter_maps_responses_partial_image_events() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
"image_request": {
"operation": "generate",
"partial_images": 1
}
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let partial = rewriter
.push_chunk(
concat!(
"event: response.image_generation_call.partial_image\n",
"data: {\"type\":\"response.image_generation_call.partial_image\",\"partial_image_index\":0,\"partial_image_b64\":\"cGFydGlhbA==\"}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let partial_text = utf8(partial);
assert!(partial_text.contains("event: image_generation.partial_image"));
assert!(partial_text.contains("\"type\":\"image_generation.partial_image\""));
assert!(partial_text.contains("\"b64_json\":\"cGFydGlhbA==\""));
assert!(partial_text.contains("\"partial_image_index\":0"));
assert!(!partial_text.contains("response.image_generation_call.partial_image"));
let done = rewriter
.push_chunk(
concat!(
"event: response.output_item.done\n",
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"result\":\"ZmluYWw=\"}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
assert!(done.is_empty());
let completed = rewriter
.push_chunk(
concat!(
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":4,\"output_tokens\":5,\"total_tokens\":9}}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let completed_text = utf8(completed);
assert!(completed_text.contains("event: image_generation.completed"));
assert!(completed_text.contains("\"type\":\"image_generation.completed\""));
assert!(completed_text.contains("\"b64_json\":\"ZmluYWw=\""));
assert!(completed_text.contains("\"total_tokens\":9"));
}
#[test]
fn openai_image_stream_rewriter_reads_final_image_from_completed_response_output() {
let report_context = json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"needs_conversion": false,
"image_request": {
"operation": "generate"
}
});
let mut rewriter =
maybe_build_local_stream_rewriter(Some(&report_context)).expect("rewriter should exist");
let completed = rewriter
.push_chunk(
concat!(
"event: response.completed\n",
"data: {\"type\":\"response.completed\",\"response\":{\"output\":[{\"type\":\"message\"},{\"type\":\"image_generation_call\",\"result\":\"ZnJvbV9vdXRwdXQ=\"}],\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n"
)
.as_bytes(),
)
.expect("rewrite should succeed");
let completed_text = utf8(completed);
assert!(completed_text.contains("event: image_generation.completed"));
assert!(completed_text.contains("\"b64_json\":\"ZnJvbV9vdXRwdXQ=\""));
assert!(completed_text.contains("\"total_tokens\":3"));
}
#[test]
fn openai_image_stream_rewriter_emits_partial_and_completed_events_for_edit() {
let report_context = json!({

45
apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/trace.rs
View File

@@ -76,6 +76,51 @@ async fn admin_monitoring_trace_request_returns_local_payload() {
assert_eq!(payload["candidates"][0]["status_code"], json!(502));
}
#[tokio::test]
async fn admin_monitoring_trace_final_status_prefers_failed_over_stale_pending() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate(
"cand-stale-pending",
"request-1",
0,
RequestCandidateStatus::Pending,
None,
None,
None,
),
sample_candidate(
"cand-failed",
"request-1",
1,
RequestCandidateStatus::Failed,
Some(101),
Some(33),
Some(502),
),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = AppState::new()
.expect("state should build")
.with_decision_trace_data_readers_for_tests(request_candidates, provider_catalog);
let context = request_context(http::Method::GET, "/api/admin/monitoring/trace/request-1");
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["final_status"], json!("failed"));
}
#[tokio::test]
async fn admin_monitoring_trace_request_keeps_format_conversion_disabled_candidates_visible() {
let mut format_disabled_candidate = sample_candidate(

17
apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs
View File

@@ -6,10 +6,11 @@ use crate::handlers::admin::shared::query_param_value;
use crate::GatewayError;
use aether_admin::observability::usage::{
admin_usage_bad_request_response, admin_usage_data_unavailable_response,
admin_usage_has_fallback, admin_usage_matches_search, admin_usage_matches_username,
admin_usage_parse_ids, admin_usage_parse_limit, admin_usage_parse_offset,
build_admin_usage_active_requests_response, build_admin_usage_records_response,
build_admin_usage_summary_stats_response_from_summary, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
admin_usage_has_fallback, admin_usage_is_failed, admin_usage_matches_search,
admin_usage_matches_username, admin_usage_parse_ids, admin_usage_parse_limit,
admin_usage_parse_offset, build_admin_usage_active_requests_response,
build_admin_usage_records_response, build_admin_usage_summary_stats_response_from_summary,
ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
};
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, UsageAuditKeywordSearchQuery, UsageAuditListQuery,
@@ -305,6 +306,14 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
})
.await?
};
let items = if requested_ids.is_some() {
items
} else {
items
.into_iter()
.filter(|item| !admin_usage_is_failed(item))
.collect::<Vec<_>>()
};
let api_key_names = admin_usage_api_key_names(state, &items).await?;
let provider_key_names = admin_usage_provider_key_names(state, &items).await?;

43
apps/aether-gateway/src/handlers/public/support/user_me_usage.rs
View File

@@ -404,6 +404,8 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
"rate_multiplier": item.settlement_rate_multiplier(),
"response_time_ms": item.response_time_ms,
"first_byte_time_ms": item.first_byte_time_ms,
"status_code": item.status_code,
"error_message": item.error_message,
"api_format": item.api_format,
"endpoint_api_format": item.endpoint_api_format,
"is_stream": item.is_stream,
@@ -435,6 +437,24 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
payload
}
fn users_me_usage_is_failed(item: &StoredRequestUsageAudit) -> bool {
let has_failure_signal = item.status_code.is_some_and(|value| value >= 400)
|| item
.error_message
.as_deref()
.is_some_and(|value| !value.trim().is_empty());
let status = item.status.trim().to_ascii_lowercase();
if status.is_empty() {
return has_failure_signal;
}
match status.as_str() {
"completed" | "cancelled" => false,
"pending" | "streaming" => has_failure_signal,
"failed" => true,
_ => false,
}
}
fn build_users_me_usage_summary_by_model(
rows: &[StoredUsageBreakdownSummaryRow],
include_actual_cost: bool,
@@ -997,6 +1017,15 @@ pub(super) async fn handle_users_me_usage_active_get(
},
};
let items = if ids.is_some() {
items
} else {
items
.into_iter()
.filter(|item| !users_me_usage_is_failed(item))
.collect::<Vec<_>>()
};
Json(json!({
"requests": items
.iter()
@@ -1233,7 +1262,8 @@ mod tests {
use super::{
build_users_me_usage_active_payload, build_users_me_usage_record_payload,
users_me_usage_client_is_stream, users_me_usage_upstream_is_stream,
users_me_usage_client_is_stream, users_me_usage_is_failed,
users_me_usage_upstream_is_stream,
};
fn sample_usage(status: &str) -> StoredRequestUsageAudit {
@@ -1310,6 +1340,17 @@ mod tests {
assert_eq!(payload["cache_creation_ephemeral_1h_input_tokens"], 6);
}
#[test]
fn user_usage_active_pending_with_failure_signal_is_not_active() {
let item = StoredRequestUsageAudit {
status_code: Some(503),
error_message: Some("upstream failed".to_string()),
..sample_usage("pending")
};
assert!(users_me_usage_is_failed(&item));
}
#[test]
fn user_usage_payloads_include_symmetric_stream_fields() {
let item = StoredRequestUsageAudit {

20
apps/aether-gateway/src/tests/control/admin/usage.rs
View File

@@ -764,6 +764,26 @@ async fn gateway_handles_admin_usage_active_locally_with_trusted_admin_principal
row.candidate_index = Some(1);
row
},
{
let mut row = sample_usage_row(
"usage-failed-pending",
"req-failed-pending",
Some("user-1"),
Some("key-1"),
Some("primary"),
"OpenAI",
"gpt-5",
"pending",
10,
0,
0.0,
0.0,
DAY_2_UNIX_SECS + 1,
);
row.status_code = Some(503);
row.error_message = Some("upstream failed".to_string());
row
},
sample_usage_row(
"usage-done",
"req-done",