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

feat(stream): 在流式执行中引入 StreamingStandardTerminalObserver 用量采集

Browse Source 
- 在 execute_stream_from_frame_stream 中集成 StreamingStandardTerminalObserver，对流式响应逐行观察并在结束时合并终态摘要
- 新增 observe_stream_usage_bytes / finalize_stream_usage_observer / merge_stream_terminal_summary 辅助函数
- 更新 codex-cli 流式集成测试：补充 response.completed 含完整 usage 字段的 mock 数据，并断言用量写入 usage_repository
This commit is contained in:
fawney19
2026-04-24 11:34:41 +08:00
parent 35400b0c2d
commit f29649e3a8
2 changed files with 177 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -41,6 +41,7 @@ use self::execution_failures::{
use crate::ai_pipeline_api::{
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,
maybe_build_stream_response_rewriter, normalize_provider_private_report_context,
StreamingStandardTerminalObserver,
};
use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
@@ -204,6 +205,84 @@ fn append_stream_capture_bytes(
}
}
fn observe_stream_usage_bytes(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
buffered: &mut Vec<u8>,
chunk: &[u8],
) {
if chunk.is_empty() {
return;
}
buffered.extend_from_slice(chunk);
while let Some(line_end) = buffered.iter().position(|byte| *byte == b'\n') {
let line = buffered.drain(..=line_end).collect::<Vec<_>>();
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
buffered.clear();
break;
}
}
}
fn finalize_stream_usage_observer(
observer: &mut Option<StreamingStandardTerminalObserver>,
report_context: Option<&Value>,
buffered: &mut Vec<u8>,
) -> Option<ExecutionStreamTerminalSummary> {
let (Some(observer), Some(report_context)) = (observer.as_mut(), report_context) else {
return None;
};
if !buffered.is_empty() {
let line = std::mem::take(buffered);
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
}
}
match observer.finish(report_context) {
Ok(summary) => summary,
Err(err) => {
observer.disable_with_error(err.to_string());
observer.latest_summary().cloned()
}
}
}
fn merge_stream_terminal_summary(
mut current: Option<ExecutionStreamTerminalSummary>,
observed: Option<ExecutionStreamTerminalSummary>,
) -> Option<ExecutionStreamTerminalSummary> {
let Some(observed) = observed else {
return current;
};
let Some(current_summary) = current.as_mut() else {
return Some(observed);
};
if current_summary.standardized_usage.is_none() {
current_summary.standardized_usage = observed.standardized_usage;
}
if current_summary.finish_reason.is_none() {
current_summary.finish_reason = observed.finish_reason;
}
if current_summary.response_id.is_none() {
current_summary.response_id = observed.response_id;
}
if current_summary.model.is_none() {
current_summary.model = observed.model;
}
current_summary.observed_finish |= observed.observed_finish;
if current_summary.parser_error.is_none() {
current_summary.parser_error = observed.parser_error;
}
current
}
async fn execute_in_process_stream(
state: &AppState,
plan: &ExecutionPlan,
@@ -1333,6 +1412,18 @@ async fn execute_stream_from_frame_stream(
} else {
maybe_build_stream_response_rewriter(normalized_stream_report_context_owned.as_ref())
};
let stream_usage_report_context =
normalized_stream_report_context_owned.clone().or_else(|| {
Some(serde_json::json!({
"provider_api_format": plan_for_report.provider_api_format.as_str(),
"client_api_format": plan_for_report.client_api_format.as_str(),
}))
});
let mut stream_usage_observer = stream_usage_report_context
.as_ref()
.filter(|_| !sync_json_stream_bridge_active_for_report)
.map(|_| StreamingStandardTerminalObserver::default());
let mut stream_usage_observer_buffered = Vec::new();
append_stream_capture_bytes(
&mut provider_buffered_body,
&provider_prefetched_body_for_report,
@@ -1382,6 +1473,17 @@ async fn execute_stream_from_frame_stream(
let replay_chunk = normalized_prefetched_chunk
.as_deref()
.unwrap_or(provider_prefetched_body_for_report.as_slice());
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
replay_chunk,
);
}
if terminal_failure.is_none() {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
if let Err(err) = rewriter.push_chunk(replay_chunk) {
@@ -1496,6 +1598,17 @@ async fn execute_stream_from_frame_stream(
} else {
chunk
};
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
&normalized_chunk,
);
}
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
{
match rewriter.push_chunk(&normalized_chunk) {
@@ -1600,6 +1713,17 @@ async fn execute_stream_from_frame_stream(
if let Some(normalizer) = private_stream_normalizer.as_mut() {
match normalizer.finish() {
Ok(normalized_chunk) if !normalized_chunk.is_empty() => {
if let (Some(observer), Some(report_context)) = (
stream_usage_observer.as_mut(),
stream_usage_report_context.as_ref(),
) {
observe_stream_usage_bytes(
observer,
report_context,
&mut stream_usage_observer_buffered,
&normalized_chunk,
);
}
let rewritten_chunk = if let Some(rewriter) = local_stream_rewriter.as_mut()
{
match rewriter.push_chunk(&normalized_chunk) {
@@ -1753,6 +1877,15 @@ async fn execute_stream_from_frame_stream(
drop(tx);
stream_terminal_summary = merge_stream_terminal_summary(
stream_terminal_summary,
finalize_stream_usage_observer(
&mut stream_usage_observer,
stream_usage_report_context.as_ref(),
&mut stream_usage_observer_buffered,
),
);
if downstream_dropped {
debug!(
event_name = "execution_runtime_stream_report_skipped",

50
apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs
View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request,
Response, Router, StatusCode, TRACE_ID_HEADER,
Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -10,6 +10,7 @@ use aether_data::repository::auth::{
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
@@ -19,6 +20,7 @@ use aether_data_contracts::repository::candidates::{
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::usage::UsageReadRepository;
use sha2::{Digest, Sha256};
#[tokio::test]
@@ -381,7 +383,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
});
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\"}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"event: response.completed\\ndata: {\\\"type\\\":\\\"response.completed\\\",\\\"response\\\":{\\\"id\\\":\\\"resp_codex_cli_stream_local_123\\\",\\\"object\\\":\\\"response\\\",\\\"model\\\":\\\"gpt-5.4\\\",\\\"status\\\":\\\"completed\\\",\\\"usage\\\":{\\\"input_tokens\\\":1,\\\"output_tokens\\\":2,\\\"total_tokens\\\":3}}}\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":41}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
@@ -416,6 +418,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
vec![sample_provider_catalog_key()],
));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (refresh_url, refresh_handle) = start_server(refresh).await;
@@ -429,15 +432,24 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
]);
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url.clone())
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_system_config_values_for_tests([(
"request_record_level".to_string(),
json!("base"),
)]),
)
.with_oauth_refresh_coordinator_for_tests(oauth_refresh);
.with_oauth_refresh_coordinator_for_tests(oauth_refresh)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -457,7 +469,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
"event: response.completed\ndata: {\"type\":\"response.completed\"}\n\n"
"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_codex_cli_stream_local_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n"
);
let seen_refresh_request = seen_refresh
@@ -511,6 +523,32 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
let mut stored_usage = None;
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(60);
loop {
stored_usage = usage_repository
.find_by_request_id("trace-codex-cli-stream-local-123")
.await
.expect("usage lookup should succeed");
if stored_usage
.as_ref()
.is_some_and(|usage| usage.status == "completed")
{
break;
}
assert!(
tokio::time::Instant::now() < deadline,
"usage should reach completed status"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored_usage = stored_usage.expect("usage should be recorded");
assert_eq!(stored_usage.total_tokens, 3);
assert!(stored_usage.request_body.is_none());
assert!(stored_usage.provider_request_body.is_none());
assert!(stored_usage.response_body.is_none());
assert!(stored_usage.client_response_body.is_none());
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(
seen_report.lock().expect("mutex should lock").is_none(),