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: 提前记录 pending 用量、优化流遥测时序与前端活跃请求发现机制

Browse Source 
- 将 record_pending 调用移至执行开始前（sync/stream 两路），确保请求在执行前即有 pending 记录
- stream_pump 在收到第一个数据块前优先 yield 遥测帧，保证 ttfb 早于 data 帧到达
- stream execution 增加 should_refresh_stream_usage_telemetry，在遥测帧携带新 ttfb/elapsed 时及时更新 record_stream_started
- access_log 对高频轮询路径（usage/active、usage/records 等）降级为 TRACE 日志，减少日志噪音
- 前端新增 reconcileActiveRequestDiscovery 工具函数及 discoverActiveRequests 逻辑，活跃请求发现与全局自动刷新解耦，空闲时降频为 5 秒扫描
- RequestDetailDrawer 调整：进行中请求不再自动开启轮询，由用户手动触发；刷新按钮 title 动态适配状态
This commit is contained in:
fawney19
2026-04-10 20:58:01 +08:00
parent d5b8583d6b
commit 46f1507d44
13 changed files with 846 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -70,6 +70,10 @@ pub(crate) async fn execute_execution_runtime_stream(
mut report_context: Option<serde_json::Value>,
) -> Result<Option<Response<Body>>, GatewayError> {
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
state
.usage_runtime
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
.await;
let plan_request_id_for_log = short_request_id(plan.request_id.as_str());
#[cfg(not(test))]
{
@@ -235,6 +239,19 @@ where
read_next_frame(lines).await
}
fn should_refresh_stream_usage_telemetry(
previous: Option<&ExecutionTelemetry>,
next: &ExecutionTelemetry,
) -> bool {
let previous_ttfb = previous.and_then(|telemetry| telemetry.ttfb_ms);
let previous_elapsed = previous.and_then(|telemetry| telemetry.elapsed_ms);
let next_ttfb = next.ttfb_ms;
let next_elapsed = next.elapsed_ms;
(next_ttfb.is_some() && next_ttfb != previous_ttfb)
|| (next_elapsed.is_some() && next_elapsed != previous_elapsed)
}
async fn probe_local_stream_success_failover_text<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
@@ -740,10 +757,6 @@ async fn execute_stream_from_frame_stream(
}
let candidate_started_unix_secs = current_request_candidate_unix_ms();
state
.usage_runtime
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
.await;
state
.usage_runtime
.record_stream_started(
@@ -794,7 +807,8 @@ async fn execute_stream_from_frame_stream(
tokio::spawn(async move {
let mut provider_buffered_body = provider_prefetched_body_for_report;
let mut buffered_body = prefetched_body_for_report;
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let reached_eof = initial_reached_eof;
let mut downstream_dropped = false;
let mut terminal_failure: Option<StreamFailureReport> = None;
@@ -928,7 +942,25 @@ async fn execute_stream_from_frame_stream(
StreamFramePayload::Telemetry {
telemetry: frame_telemetry,
} => {
telemetry = Some(frame_telemetry);
let should_refresh_stream_usage = should_refresh_stream_usage_telemetry(
usage_stream_telemetry.as_ref(),
&frame_telemetry,
);
telemetry = Some(frame_telemetry.clone());
if should_refresh_stream_usage {
state_for_report
.usage_runtime
.record_stream_started(
state_for_report.data.as_ref(),
&plan_for_report,
report_context_owned.as_ref(),
status_code,
&headers_for_report,
Some(&frame_telemetry),
)
.await;
usage_stream_telemetry = Some(frame_telemetry);
}
}
StreamFramePayload::Eof { .. } => {
break;

120
apps/aether-gateway/src/execution_runtime/stream_pump.rs
View File

@@ -42,6 +42,7 @@ pub(crate) fn build_direct_execution_frame_stream(
let mut upstream_bytes = 0u64;
let mut ttfb_ms = None;
let mut first_chunk_telemetry_emitted = false;
let mut bytes_stream = response.bytes_stream();
while let Some(item) = bytes_stream.next().await {
match item {
@@ -49,6 +50,26 @@ pub(crate) fn build_direct_execution_frame_stream(
if ttfb_ms.is_none() {
ttfb_ms = Some(started_at.elapsed().as_millis() as u64);
}
if !first_chunk_telemetry_emitted {
let telemetry_frame = StreamFrame {
frame_type: StreamFrameType::Telemetry,
payload: StreamFramePayload::Telemetry {
telemetry: ExecutionTelemetry {
ttfb_ms,
elapsed_ms: ttfb_ms,
upstream_bytes: Some(upstream_bytes),
},
},
};
match encode_stream_frame_ndjson(&telemetry_frame) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
return;
}
}
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
let frame = StreamFrame {
frame_type: StreamFrameType::Data,
@@ -217,4 +238,103 @@ mod tests {
"telemetry frame should include a measured ttfb"
);
}
#[tokio::test]
async fn direct_execution_frame_stream_emits_telemetry_before_first_data_frame() {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/stream",
post(|| async {
let body_stream = stream! {
tokio::time::sleep(Duration::from_millis(10)).await;
yield Ok::<Bytes, Infallible>(Bytes::from_static(b"data: hello\n\n"));
};
(
[(
header::CONTENT_TYPE,
HeaderValue::from_static("text/event-stream"),
)],
Body::from_stream(body_stream),
)
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let runtime = DirectSyncExecutionRuntime::new();
let execution = runtime
.execute_stream(ExecutionPlan {
request_id: "req-telemetry-order".to_string(),
candidate_id: Some("cand-telemetry-order".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: format!("http://{addr}/stream"),
headers: BTreeMap::new(),
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: true,
client_api_format: "openai:chat".to_string(),
provider_api_format: "openai:chat".to_string(),
model_name: Some("gpt-5".into()),
proxy: None,
tls_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
})
.await
.expect("stream execution should succeed");
let frames = build_direct_execution_frame_stream(execution)
.map(|item| item.expect("frame should encode"))
.collect::<Vec<_>>()
.await
.into_iter()
.map(|bytes| String::from_utf8(bytes.to_vec()).expect("frame should be utf8"))
.collect::<Vec<_>>();
server.abort();
let frame_types = frames
.iter()
.map(|line| {
serde_json::from_str::<Value>(line)
.expect("frame should parse")
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
.to_string()
})
.collect::<Vec<_>>();
let first_data_idx = frame_types
.iter()
.position(|kind| kind == "data")
.expect("data frame should exist");
let first_telemetry_idx = frame_types
.iter()
.position(|kind| kind == "telemetry")
.expect("telemetry frame should exist");
assert!(
first_telemetry_idx < first_data_idx,
"first telemetry frame should be emitted before the first data frame"
);
}
}

8
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -87,6 +87,10 @@ pub(crate) async fn execute_execution_runtime_sync(
let plan_request_id_for_log = short_request_id(plan_request_id);
let plan_candidate_id = plan.candidate_id.as_deref();
let candidate_started_unix_secs = current_request_candidate_unix_ms();
state
.usage_runtime
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
.await;
#[cfg(not(test))]
let result = {
match DirectSyncExecutionRuntime::new()
@@ -276,10 +280,6 @@ pub(crate) async fn execute_execution_runtime_sync(
return Ok(None);
}
state
.usage_runtime
.record_pending(state.data.as_ref(), &plan, report_context.as_ref())
.await;
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,

22
apps/aether-gateway/src/handlers/proxy/finalize.rs
View File

@@ -7,13 +7,13 @@ use crate::constants::{
};
use crate::control::GatewayControlDecision;
use crate::control::GatewayPublicRequestContext;
use crate::middleware::RequestLogEmitted;
use crate::middleware::{should_downgrade_access_log, RequestLogEmitted};
use crate::AppState;
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
use axum::body::{Body, Bytes};
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
use std::time::Instant;
use tracing::{info, warn};
use tracing::{info, trace, warn};
pub(super) fn request_wants_stream(
request_context: &GatewayPublicRequestContext,
@@ -113,6 +113,24 @@ pub(super) fn finalize_gateway_response(
elapsed_ms,
"gateway request failed"
);
} else if should_downgrade_access_log(method, path_and_query) {
trace!(
event_name = "http_request_completed",
log_type = "access",
status = "completed",
status_code,
trace_id = %trace_id,
request_id,
remote_addr = %remote_addr,
method = %method,
path = %path_and_query,
route_class,
execution_path,
dependency_reason = dependency_reason.as_str(),
local_execution_runtime_miss_reason = local_execution_runtime_miss_reason.as_str(),
elapsed_ms,
"gateway completed request"
);
} else {
info!(
event_name = "http_request_completed",

181
apps/aether-gateway/src/middleware/access_log.rs
View File

@@ -3,9 +3,10 @@ use std::time::Instant;
use axum::body::Body;
use axum::extract::Request;
use axum::http::header::{HeaderName, HeaderValue};
use axum::http::Method;
use axum::middleware::Next;
use axum::response::Response;
use tracing::{info, warn};
use tracing::{info, trace, warn};
use crate::constants::{
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
@@ -16,6 +17,41 @@ use crate::log_ids::short_request_id;
#[derive(Debug, Clone, Copy)]
pub(crate) struct RequestLogEmitted;
fn is_usage_detail_path(path: &str) -> bool {
let Some(detail_id) = path.strip_prefix("/api/admin/usage/") else {
return false;
};
!detail_id.is_empty()
&& !detail_id.contains('/')
&& !matches!(detail_id, "active" | "records" | "stats" | "heatmap")
}
pub(crate) fn should_downgrade_access_log(method: &Method, path: &str) -> bool {
if method != Method::GET {
return false;
}
let normalized_path = path.split('?').next().unwrap_or(path);
matches!(
normalized_path,
"/api/admin/usage/active"
| "/api/users/me/usage/active"
| "/api/admin/usage/records"
| "/api/admin/usage/stats"
| "/api/admin/usage/aggregation/stats"
| "/api/admin/usage/heatmap"
| "/api/admin/usage/cache-affinity/interval-timeline"
| "/api/admin/usage/cache-affinity/ttl-analysis"
| "/api/admin/usage/cache-affinity/hit-analysis"
| "/api/admin/users"
| "/api/admin/monitoring/cache/stats"
| "/api/admin/monitoring/cache/model-mapping/stats"
| "/api/admin/monitoring/cache/config"
| "/api/admin/monitoring/cache/redis-keys"
| "/api/admin/monitoring/cache/affinities"
) || is_usage_detail_path(normalized_path)
|| normalized_path.starts_with("/api/admin/monitoring/trace/")
}
pub(crate) async fn access_log_middleware(request: Request<Body>, next: Next) -> Response {
let started_at = Instant::now();
let method = request.method().clone();
@@ -25,18 +61,33 @@ pub(crate) async fn access_log_middleware(request: Request<Body>, next: Next) ->
.map(|value| value.as_str().to_string())
.unwrap_or_else(|| "/".to_string());
let trace_id = extract_or_generate_trace_id(request.headers());
info!(
event_name = "http_request_started",
log_type = "access",
status = "started",
trace_id = %trace_id,
request_id = "-",
method = %method,
path = %path,
route_class = "pending",
execution_path = "pending",
"gateway request started"
);
if should_downgrade_access_log(&method, &path) {
trace!(
event_name = "http_request_started",
log_type = "access",
status = "started",
trace_id = %trace_id,
request_id = "-",
method = %method,
path = %path,
route_class = "pending",
execution_path = "pending",
"gateway request started"
);
} else {
info!(
event_name = "http_request_started",
log_type = "access",
status = "started",
trace_id = %trace_id,
request_id = "-",
method = %method,
path = %path,
route_class = "pending",
execution_path = "pending",
"gateway request started"
);
}
let mut response = next.run(request).await;
if !response.headers().contains_key(TRACE_ID_HEADER) {
response.headers_mut().insert(
@@ -79,6 +130,21 @@ pub(crate) async fn access_log_middleware(request: Request<Body>, next: Next) ->
elapsed_ms,
"gateway request failed"
);
} else if should_downgrade_access_log(&method, &path) {
trace!(
event_name = "http_request_completed",
log_type = "access",
status = "completed",
status_code,
trace_id = %trace_id,
request_id,
method = %method,
path = %path,
route_class,
execution_path,
elapsed_ms,
"gateway completed request"
);
} else {
info!(
event_name = "http_request_completed",
@@ -101,19 +167,20 @@ pub(crate) async fn access_log_middleware(request: Request<Body>, next: Next) ->
#[cfg(test)]
mod tests {
use super::access_log_middleware;
use super::{access_log_middleware, should_downgrade_access_log};
use crate::constants::{
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER,
TRACE_ID_HEADER,
};
use axum::body::Body;
use axum::http::{Request, Response, StatusCode};
use axum::http::{Method, Request, Response, StatusCode};
use axum::routing::get;
use axum::Router;
use bytes::Bytes;
use futures_util::stream;
use std::sync::{Arc, Mutex};
use tower::ServiceExt;
use tracing_subscriber::filter::LevelFilter;
use tracing_subscriber::prelude::*;
#[derive(Clone, Default)]
@@ -416,4 +483,88 @@ mod tests {
assert_eq!(logs[1]["route_class"], "ai_public");
assert_eq!(logs[1]["execution_path"], "execution_runtime_stream");
}
#[tokio::test]
async fn access_log_downgrades_usage_active_polling_to_trace() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::TRACE),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/api/admin/usage/active",
get(|| async {
Response::builder()
.status(StatusCode::OK)
.header(CONTROL_ROUTE_CLASS_HEADER, "admin_proxy")
.header(EXECUTION_PATH_HEADER, "public_proxy_passthrough")
.body(Body::empty())
.expect("response should build")
}),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/api/admin/usage/active?ids=req-1")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs.len(), 2);
assert_eq!(logs[0]["level"], "TRACE");
assert_eq!(logs[0]["event_name"], "http_request_started");
assert_eq!(logs[1]["level"], "TRACE");
assert_eq!(logs[1]["event_name"], "http_request_completed");
}
#[test]
fn access_log_marks_usage_active_paths_as_high_frequency() {
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/active"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/active?ids=req-1"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/users/me/usage/active"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/records?limit=20"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/usage/123e4567-e89b-12d3-a456-426614174000?include_bodies=false"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/monitoring/trace/req-123?attempted_only=false"
));
assert!(should_downgrade_access_log(
&Method::GET,
"/api/admin/monitoring/cache/stats"
));
assert!(!should_downgrade_access_log(
&Method::DELETE,
"/api/admin/monitoring/cache/affinity/provider/key/model/openai:cli"
));
assert!(!should_downgrade_access_log(&Method::GET, "/v1/responses"));
}
}

4
apps/aether-gateway/src/middleware/mod.rs
View File

@@ -2,6 +2,8 @@ mod access_log;
mod frontdoor_cors;
mod strip_cf_headers;
pub(crate) use access_log::{access_log_middleware, RequestLogEmitted};
pub(crate) use access_log::{
access_log_middleware, should_downgrade_access_log, RequestLogEmitted,
};
pub(crate) use frontdoor_cors::frontdoor_cors_middleware;
pub use strip_cf_headers::strip_cf_headers_middleware;

301
apps/aether-gateway/src/tests/usage/direct.rs
View File

@@ -118,6 +118,159 @@ async fn gateway_records_usage_for_execution_runtime_sync_when_runtime_enabled()
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_records_pending_usage_before_execution_runtime_sync_result_arrives() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let execution_request_started = Arc::new(tokio::sync::Notify::new());
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
let upstream = Router::new().route(
"/api/internal/gateway/report-sync",
any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any({
let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response);
move |_request: Request| {
let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response);
async move {
execution_request_started.notify_one();
allow_execution_response.notified().await;
Json(json!({
"request_id": "req-usage-sync-pending-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-usage-sync-pending-123",
"usage": {
"input_tokens": 3,
"output_tokens": 5,
"total_tokens": 8
}
}
},
"telemetry": {
"elapsed_ms": 45
}
}))
}
}
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-usage-sync-pending")),
sample_local_openai_auth_snapshot(
"api-key-usage-sync-pending-123",
"user-usage-sync-pending-123",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway_state =
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
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_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;
let request_task = tokio::spawn({
let gateway_url = gateway_url.clone();
async move {
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-usage-sync-pending",
)
.header(TRACE_ID_HEADER, "req-usage-sync-pending-123")
.body("{\"model\":\"gpt-5\",\"messages\":[]}")
.send()
.await
.expect("request should succeed");
let status = response.status();
let body = response.text().await.expect("body should read");
(status, body)
}
});
execution_request_started.notified().await;
let mut pending = None;
for _ in 0..50 {
pending = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if pending
.as_ref()
.is_some_and(|stored| stored.status == "pending")
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let pending = pending.expect("pending usage should be recorded before sync result resolves");
assert_eq!(pending.status, "pending");
assert_eq!(pending.billing_status, "pending");
assert_eq!(pending.response_time_ms, None);
allow_execution_response.notify_one();
let (status, _body) = request_task.await.expect("request task should join");
assert_eq!(status, StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-sync-pending-123")
.await
.expect("usage lookup should succeed");
if stored.as_ref().is_some_and(|row| row.status == "completed") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.response_time_ms, Some(45));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -239,3 +392,151 @@ async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_records_pending_usage_before_execution_runtime_stream_headers_arrive() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let execution_request_started = Arc::new(tokio::sync::Notify::new());
let allow_execution_response = Arc::new(tokio::sync::Notify::new());
let upstream = Router::new().route(
"/api/internal/gateway/report-stream",
any(|_request: Request| async move { Json(json!({"ok": true})) }),
);
let execution_runtime = Router::new().route(
"/v1/execute/stream",
any({
let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response);
move |_request: Request| {
let execution_request_started = Arc::clone(&execution_request_started);
let allow_execution_response = Arc::clone(&allow_execution_response);
async move {
execution_request_started.notify_one();
allow_execution_response.notified().await;
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-pending-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
);
let mut response = Response::builder()
.status(StatusCode::OK)
.body(Body::from(frames))
.expect("response should build");
response.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}
}
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-stream-pending")),
sample_local_openai_auth_snapshot(
"api-key-usage-stream-pending-123",
"user-usage-stream-pending-123",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let gateway_state = build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
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),
usage_repository.clone(),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.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;
let request_task = tokio::spawn({
let gateway_url = gateway_url.clone();
async move {
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-stream-pending",
)
.header(TRACE_ID_HEADER, "req-usage-stream-pending-123")
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send()
.await
.expect("request should succeed");
let status = response.status();
let body = response.text().await.expect("stream body should read");
(status, body)
}
});
execution_request_started.notified().await;
let mut pending = None;
for _ in 0..50 {
pending = usage_repository
.find_by_request_id("req-usage-stream-pending-123")
.await
.expect("usage lookup should succeed");
if pending
.as_ref()
.is_some_and(|stored| stored.status == "pending")
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let pending = pending.expect("pending usage should be recorded before stream headers arrive");
assert_eq!(pending.status, "pending");
assert_eq!(pending.billing_status, "pending");
assert_eq!(pending.first_byte_time_ms, None);
assert_eq!(pending.response_time_ms, None);
allow_execution_response.notify_one();
let (status, _body) = request_task.await.expect("request task should join");
assert_eq!(status, StatusCode::OK);
let mut stored = None;
for _ in 0..50 {
stored = usage_repository
.find_by_request_id("req-usage-stream-pending-123")
.await
.expect("usage lookup should succeed");
if stored.as_ref().is_some_and(|row| row.status == "completed") {
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.first_byte_time_ms, Some(19));
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}