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

Refactor usage body capture and stream terminal reporting

Browse Source
This commit is contained in:
fawney19
2026-04-18 17:48:21 +08:00
parent 569242d72f
commit 3363592751
36 changed files with 2673 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/src/ai_pipeline/adaptation/mod.rs
View File

@@ -14,4 +14,5 @@ pub(crate) use private_envelope::{
maybe_normalize_provider_private_sync_report_payload,
normalize_provider_private_report_context, normalize_provider_private_response_value,
provider_private_response_allows_sync_finalize, transform_provider_private_stream_line,
ProviderPrivateStreamNormalizer,
};

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

@@ -14,7 +14,9 @@ use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
use self::contracts::ExecutionRuntimeAuthContext;
pub(crate) use self::adaptation::maybe_build_provider_private_stream_normalizer;
pub(crate) use self::adaptation::{
maybe_build_provider_private_stream_normalizer, ProviderPrivateStreamNormalizer,
};
pub(crate) use self::finalize::common::LocalCoreSyncFinalizeOutcome;
pub(crate) use self::finalize::internal::{
maybe_build_stream_response_rewriter, maybe_build_sync_finalize_outcome,

9
apps/aether-gateway/src/ai_pipeline_api.rs
View File

@@ -34,10 +34,11 @@ pub(crate) use aether_ai_pipeline::api::{
GatewayControlSyncDecisionResponse, LocalCoreSyncErrorKind, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
LocalStandardSpec, LocalStreamPlanAndReport, LocalSyncPlanAndReport,
EXECUTION_RUNTIME_STREAM_DECISION_ACTION, EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
GEMINI_FILES_DOWNLOAD_PLAN_KIND, GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CONTENT_PLAN_KIND,
OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND, OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_VIDEO_CANCEL_SYNC_PLAN_KIND,
OPENAI_VIDEO_CONTENT_PLAN_KIND, OPENAI_VIDEO_DELETE_SYNC_PLAN_KIND,
OPENAI_VIDEO_REMIX_SYNC_PLAN_KIND,
};
pub(crate) fn parse_direct_request_body(

65
apps/aether-gateway/src/data/state/integrations.rs
View File

@@ -16,8 +16,10 @@ use aether_data_contracts::repository::settlement::{StoredUsageSettlement, Usage
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
use aether_usage_runtime::{
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRequestRecordLevel,
UsageRuntimeAccess, UsageSettlementWriter,
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEvent, UsageRecordWriter,
UsageRequestRecordLevel, UsageRuntimeAccess, UsageSettlementWriter,
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
use aether_video_tasks_core::StoredVideoTaskReadSide;
use async_trait::async_trait;
@@ -29,6 +31,8 @@ use crate::provider_transport::ProviderTransportSnapshotSource;
const REQUEST_RECORD_LEVEL_KEY: &str = "request_record_level";
const LEGACY_REQUEST_LOG_LEVEL_KEY: &str = "request_log_level";
const MAX_REQUEST_BODY_SIZE_KEY: &str = "max_request_body_size";
const MAX_RESPONSE_BODY_SIZE_KEY: &str = "max_response_body_size";
fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestRecordLevel {
let Some(value) = value.and_then(Value::as_str).map(str::trim) else {
@@ -47,6 +51,14 @@ fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestR
}
}
fn usage_body_capture_limit_from_value(value: Option<&Value>, default: usize) -> Option<usize> {
match value.and_then(Value::as_u64) {
Some(0) => None,
Some(limit) => usize::try_from(limit).ok().filter(|limit| *limit > 0),
None => Some(default),
}
}
#[async_trait]
impl RequestAuditReader for GatewayDataState {
async fn find_request_usage_audit_by_request_id(
@@ -196,14 +208,28 @@ impl UsageRuntimeAccess for GatewayDataState {
GatewayDataState::usage_worker_runner(self)
}
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
let value = GatewayDataState::find_system_config_value(self, REQUEST_RECORD_LEVEL_KEY)
.await?
.or(
GatewayDataState::find_system_config_value(self, LEGACY_REQUEST_LOG_LEVEL_KEY)
.await?,
);
Ok(usage_request_record_level_from_value(value.as_ref()))
let max_request_body_size =
GatewayDataState::find_system_config_value(self, MAX_REQUEST_BODY_SIZE_KEY).await?;
let max_response_body_size =
GatewayDataState::find_system_config_value(self, MAX_RESPONSE_BODY_SIZE_KEY).await?;
Ok(UsageBodyCapturePolicy {
record_level: usage_request_record_level_from_value(value.as_ref()),
max_request_body_bytes: usage_body_capture_limit_from_value(
max_request_body_size.as_ref(),
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
),
max_response_body_bytes: usage_body_capture_limit_from_value(
max_response_body_size.as_ref(),
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
),
})
}
}
@@ -327,4 +353,35 @@ mod tests {
assert_eq!(level, UsageRequestRecordLevel::Full);
}
#[tokio::test]
async fn usage_runtime_access_reads_body_capture_limits_from_system_config() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([
("max_request_body_size".to_string(), json!(1234)),
("max_response_body_size".to_string(), json!(5678)),
]);
let policy = UsageRuntimeAccess::body_capture_policy(&state)
.await
.expect("body capture policy should read");
assert_eq!(policy.record_level, UsageRequestRecordLevel::Full);
assert_eq!(policy.max_request_body_bytes, Some(1234));
assert_eq!(policy.max_response_body_bytes, Some(5678));
}
#[tokio::test]
async fn usage_runtime_access_treats_zero_body_capture_limit_as_unbounded() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([
("max_request_body_size".to_string(), json!(0)),
("max_response_body_size".to_string(), json!(0)),
]);
let policy = UsageRuntimeAccess::body_capture_policy(&state)
.await
.expect("body capture policy should read");
assert_eq!(policy.max_request_body_bytes, None);
assert_eq!(policy.max_response_body_bytes, None);
}
}

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

@@ -1,14 +1,20 @@
use std::collections::{BTreeMap, VecDeque};
use std::io::Error as IoError;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry, StreamFrame, StreamFramePayload};
use aether_contracts::{
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame,
StreamFramePayload,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use aether_scheduler_core::{
parse_request_candidate_report_context, SchedulerRequestCandidateStatusUpdate,
};
use aether_usage_runtime::{
build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed,
build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed,
UsageBodyCapturePolicy, UsageRequestRecordLevel, UsageRuntimeAccess,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
use async_stream::stream;
use axum::body::{Body, Bytes};
@@ -102,6 +108,27 @@ fn record_stream_terminal_usage(
);
}
fn append_stream_capture_bytes(
buffer: &mut Vec<u8>,
chunk: &[u8],
max_bytes: usize,
truncated: &mut bool,
) {
if chunk.is_empty() || max_bytes == 0 {
return;
}
if buffer.len() >= max_bytes {
*truncated = true;
return;
}
let remaining = max_bytes - buffer.len();
let keep_len = remaining.min(chunk.len());
buffer.extend_from_slice(&chunk[..keep_len]);
if keep_len < chunk.len() {
*truncated = true;
}
}
async fn execute_in_process_stream(
state: &AppState,
plan: &ExecutionPlan,
@@ -165,7 +192,10 @@ pub(crate) async fn execute_execution_runtime_stream(
#[cfg(not(test))]
{
let execution = match execute_in_process_stream(state, &plan).await {
Ok(execution) => execution,
Ok(mut execution) => {
apply_stream_summary_report_context(&mut execution, report_context.as_ref());
execution
}
Err(err) => {
info!(
event_name = "stream_execution_runtime_unavailable",
@@ -221,7 +251,10 @@ pub(crate) async fn execute_execution_runtime_stream(
.unwrap_or_default();
if remote_execution_runtime_base_url.trim().is_empty() {
let execution = match execute_in_process_stream(state, &plan).await {
Ok(execution) => execution,
Ok(mut execution) => {
apply_stream_summary_report_context(&mut execution, report_context.as_ref());
execution
}
Err(err) => {
info!(
event_name = "stream_execution_runtime_unavailable",
@@ -513,6 +546,7 @@ async fn execute_stream_from_frame_stream(
));
};
let mut buffered_frames = VecDeque::new();
let mut stream_terminal_summary: Option<ExecutionStreamTerminalSummary> = None;
if status_code == 200 {
let success_probe_text =
probe_local_stream_success_failover_text(&mut buffered_frames, &mut lines).await?;
@@ -1038,7 +1072,8 @@ async fn execute_stream_from_frame_stream(
} => {
prefetched_telemetry = Some(frame_telemetry);
}
StreamFramePayload::Eof { .. } => {
StreamFramePayload::Eof { summary } => {
stream_terminal_summary = summary;
reached_eof = true;
break;
}
@@ -1125,13 +1160,51 @@ async fn execute_stream_from_frame_stream(
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
let emit_passthrough_sse_terminal_error =
skip_direct_finalize_prefetch && response_headers_indicate_sse(&headers);
let body_capture_policy = match UsageRuntimeAccess::body_capture_policy(state.data.as_ref())
.await
{
Ok(policy) => policy,
Err(err) => {
warn!(
event_name = "stream_execution_body_capture_policy_read_failed",
log_type = "ops",
trace_id = %trace_id,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = %err,
fallback_request_body_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
"gateway failed to read body capture policy; falling back to default stream capture limits"
);
UsageBodyCapturePolicy::default()
}
};
let max_stream_body_buffer_bytes = if matches!(
body_capture_policy.record_level,
UsageRequestRecordLevel::Basic
) {
0
} else {
body_capture_policy
.max_response_body_bytes
.unwrap_or(DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES)
};
tokio::spawn(async move {
const MAX_STREAM_BODY_BUFFER_BYTES: usize = 256 * 1024; // 256KB
let mut provider_buffered_body: VecDeque<u8> = provider_prefetched_body_for_report.into();
let mut buffered_body: VecDeque<u8> = prefetched_body_for_report.into();
let mut provider_buffered_body = Vec::new();
let mut buffered_body = Vec::new();
let mut provider_body_truncated = false;
let mut client_body_truncated = false;
append_stream_capture_bytes(
&mut provider_buffered_body,
&provider_prefetched_body_for_report,
max_stream_body_buffer_bytes,
&mut provider_body_truncated,
);
append_stream_capture_bytes(
&mut buffered_body,
&prefetched_body_for_report,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let reached_eof = initial_reached_eof;
@@ -1193,13 +1266,12 @@ async fn execute_stream_from_frame_stream(
continue;
}
provider_buffered_body.extend(chunk.iter().copied());
if provider_buffered_body.len() > MAX_STREAM_BODY_BUFFER_BYTES {
let tail_start =
provider_buffered_body.len() - MAX_STREAM_BODY_BUFFER_BYTES;
provider_buffered_body.drain(..tail_start);
provider_body_truncated = true;
}
append_stream_capture_bytes(
&mut provider_buffered_body,
&chunk,
max_stream_body_buffer_bytes,
&mut provider_body_truncated,
);
let normalized_chunk = if let Some(normalizer) =
private_stream_normalizer.as_mut()
{
@@ -1256,12 +1328,12 @@ async fn execute_stream_from_frame_stream(
continue;
}
buffered_body.extend(&rewritten_chunk);
if buffered_body.len() > MAX_STREAM_BODY_BUFFER_BYTES {
let tail_start = buffered_body.len() - MAX_STREAM_BODY_BUFFER_BYTES;
buffered_body.drain(..tail_start);
client_body_truncated = true;
}
append_stream_capture_bytes(
&mut buffered_body,
&rewritten_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_disconnected",
@@ -1293,7 +1365,8 @@ async fn execute_stream_from_frame_stream(
usage_stream_telemetry = Some(frame_telemetry);
}
}
StreamFramePayload::Eof { .. } => {
StreamFramePayload::Eof { summary } => {
stream_terminal_summary = summary;
break;
}
StreamFramePayload::Error { error } => {
@@ -1355,7 +1428,12 @@ async fn execute_stream_from_frame_stream(
normalized_chunk
};
if !rewritten_chunk.is_empty() {
buffered_body.extend(&rewritten_chunk);
append_stream_capture_bytes(
&mut buffered_body,
&rewritten_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_flush_disconnected",
@@ -1394,7 +1472,12 @@ async fn execute_stream_from_frame_stream(
if let Some(rewriter) = local_stream_rewriter.as_mut() {
match rewriter.finish() {
Ok(flushed_chunk) if !flushed_chunk.is_empty() => {
buffered_body.extend(&flushed_chunk);
append_stream_capture_bytes(
&mut buffered_body,
&flushed_chunk,
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
@@ -1435,12 +1518,12 @@ async fn execute_stream_from_frame_stream(
if let Some(failure) = terminal_failure.as_ref() {
match encode_terminal_sse_error_event(failure) {
Ok(error_event) => {
buffered_body.extend(error_event.iter().copied());
if buffered_body.len() > MAX_STREAM_BODY_BUFFER_BYTES {
let tail_start = buffered_body.len() - MAX_STREAM_BODY_BUFFER_BYTES;
buffered_body.drain(..tail_start);
client_body_truncated = true;
}
append_stream_capture_bytes(
&mut buffered_body,
error_event.as_ref(),
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
if tx.send(Ok(error_event)).await.is_err() {
warn!(
event_name = "stream_execution_downstream_terminal_error_disconnected",
@@ -1490,18 +1573,26 @@ async fn execute_stream_from_frame_stream(
report_context: report_context_owned.clone(),
status_code: 499,
headers: headers_for_report.clone(),
provider_body_base64: (!provider_body_truncated
&& !provider_buffered_body.is_empty())
.then(|| {
base64::engine::general_purpose::STANDARD
.encode(provider_buffered_body.make_contiguous())
provider_body_base64: (!provider_buffered_body.is_empty()).then(|| {
base64::engine::general_purpose::STANDARD.encode(&provider_buffered_body)
}),
client_body_base64: (!client_body_truncated && !buffered_body.is_empty()).then(
|| {
base64::engine::general_purpose::STANDARD
.encode(buffered_body.make_contiguous())
},
),
provider_body_state: Some(if provider_body_truncated {
UsageBodyCaptureState::Truncated
} else if provider_buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
client_body_base64: (!buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
client_body_state: Some(if client_body_truncated {
UsageBodyCaptureState::Truncated
} else if buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
terminal_summary: stream_terminal_summary.clone(),
telemetry: telemetry.clone(),
},
true,
@@ -1533,11 +1624,7 @@ async fn execute_stream_from_frame_stream(
report_context_owned.as_ref(),
&headers_for_report,
telemetry.clone(),
if provider_body_truncated {
&[]
} else {
provider_buffered_body.make_contiguous()
},
&provider_buffered_body,
candidate_started_unix_secs_for_report,
failure,
)
@@ -1551,14 +1638,25 @@ async fn execute_stream_from_frame_stream(
report_context: report_context_owned.clone(),
status_code,
headers: headers_for_report.clone(),
provider_body_base64: (!provider_body_truncated && !provider_buffered_body.is_empty())
.then(|| {
base64::engine::general_purpose::STANDARD
.encode(provider_buffered_body.make_contiguous())
}),
client_body_base64: (!client_body_truncated && !buffered_body.is_empty()).then(|| {
base64::engine::general_purpose::STANDARD.encode(buffered_body.make_contiguous())
provider_body_base64: (!provider_buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&provider_buffered_body)),
provider_body_state: Some(if provider_body_truncated {
UsageBodyCaptureState::Truncated
} else if provider_buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
client_body_base64: (!buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
client_body_state: Some(if client_body_truncated {
UsageBodyCaptureState::Truncated
} else if buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
terminal_summary: stream_terminal_summary,
telemetry: telemetry.clone(),
};
apply_local_execution_effect(
@@ -1656,6 +1754,15 @@ async fn execute_stream_from_frame_stream(
)?))
}
fn apply_stream_summary_report_context(
execution: &mut DirectUpstreamStreamExecution,
report_context: Option<&Value>,
) {
if let Some(report_context) = report_context.cloned() {
execution.stream_summary_report_context = report_context;
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;

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

@@ -1,15 +1,20 @@
use std::io::Error as IoError;
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionTelemetry, StreamFrame,
StreamFramePayload, StreamFrameType,
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionStreamTerminalSummary,
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType,
};
use async_stream::stream;
use axum::body::Bytes;
use base64::Engine as _;
use futures_util::{Stream, StreamExt};
use serde_json::Value;
use tracing::warn;
use crate::ai_pipeline_api::{
maybe_build_provider_private_stream_normalizer, normalize_provider_private_report_context,
StreamingStandardTerminalObserver,
};
use crate::execution_runtime::ndjson::encode_stream_frame_ndjson;
use crate::execution_runtime::transport::DirectUpstreamResponse;
use crate::execution_runtime::DirectUpstreamStreamExecution;
@@ -23,10 +28,35 @@ pub(crate) fn build_direct_execution_frame_stream(
candidate_id: _,
status_code,
headers,
provider_api_format,
stream_summary_report_context,
response,
started_at,
} = execution;
let mut observer_context = stream_summary_report_context;
if observer_context
.get("provider_api_format")
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default()
.is_empty()
{
if let Some(object) = observer_context.as_object_mut() {
object.insert(
"provider_api_format".to_string(),
Value::String(provider_api_format.clone()),
);
}
}
let normalized_observer_context =
normalize_provider_private_report_context(Some(&observer_context))
.unwrap_or_else(|| observer_context.clone());
let mut private_stream_normalizer =
maybe_build_provider_private_stream_normalizer(Some(&observer_context));
let mut stream_terminal_observer = StreamingStandardTerminalObserver::default();
let mut observer_buffered = Vec::new();
match encode_headers_frame(status_code, headers) {
Ok(frame) => yield Ok(frame),
Err(err) => {
@@ -58,6 +88,13 @@ pub(crate) fn build_direct_execution_frame_stream(
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
@@ -105,6 +142,13 @@ pub(crate) fn build_direct_execution_frame_stream(
first_chunk_telemetry_emitted = true;
}
upstream_bytes += chunk.len() as u64;
observe_stream_chunk(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
chunk.as_ref(),
);
match encode_data_frame(&chunk) {
Ok(frame) => yield Ok(frame),
Err(err) => {
@@ -135,6 +179,12 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
}
let summary = finalize_stream_terminal_summary(
&mut stream_terminal_observer,
&normalized_observer_context,
private_stream_normalizer.as_mut(),
&mut observer_buffered,
);
match encode_telemetry_frame(
ttfb_ms,
@@ -147,7 +197,7 @@ pub(crate) fn build_direct_execution_frame_stream(
return;
}
}
match encode_stream_frame_ndjson(&StreamFrame::eof()) {
match encode_stream_frame_ndjson(&StreamFrame::eof_with_summary(summary)) {
Ok(frame) => yield Ok(frame),
Err(err) => yield Err(err),
}
@@ -221,6 +271,83 @@ fn format_error_chain(err: &(dyn std::error::Error + 'static)) -> String {
message
}
fn observe_stream_chunk(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
private_stream_normalizer: Option<&mut crate::ai_pipeline::ProviderPrivateStreamNormalizer>,
observer_buffered: &mut Vec<u8>,
chunk: &[u8],
) {
let normalized = if let Some(normalizer) = private_stream_normalizer {
match normalizer.push_chunk(chunk) {
Ok(normalized) => normalized,
Err(err) => {
observer.disable_with_error(format!(
"failed to normalize provider private stream chunk: {err:?}"
));
return;
}
}
} else {
chunk.to_vec()
};
observe_normalized_bytes(observer, report_context, observer_buffered, &normalized);
}
fn finalize_stream_terminal_summary(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
private_stream_normalizer: Option<&mut crate::ai_pipeline::ProviderPrivateStreamNormalizer>,
observer_buffered: &mut Vec<u8>,
) -> Option<ExecutionStreamTerminalSummary> {
if let Some(normalizer) = private_stream_normalizer {
match normalizer.finish() {
Ok(flushed) => {
observe_normalized_bytes(observer, report_context, observer_buffered, &flushed)
}
Err(err) => observer.disable_with_error(format!(
"failed to flush provider private stream normalization: {err:?}"
)),
}
}
if !observer_buffered.is_empty() {
let line = std::mem::take(observer_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 observe_normalized_bytes(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
observer_buffered: &mut Vec<u8>,
normalized: &[u8],
) {
if normalized.is_empty() {
return;
}
observer_buffered.extend_from_slice(normalized);
while let Some(line_end) = observer_buffered.iter().position(|byte| *byte == b'\n') {
let line = observer_buffered.drain(..=line_end).collect::<Vec<_>>();
if let Err(err) = observer.push_line(report_context, line) {
observer.disable_with_error(err.to_string());
observer_buffered.clear();
break;
}
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;

16
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -18,6 +18,7 @@ use reqwest::header::{HeaderMap, HeaderName, HeaderValue};
use reqwest::redirect::Policy;
use reqwest::tls::Version;
use serde::Serialize;
use serde_json::json;
use serde_json::Value;
use thiserror::Error;
@@ -147,6 +148,8 @@ pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) candidate_id: Option<String>,
pub(crate) status_code: u16,
pub(crate) headers: BTreeMap<String, String>,
pub(crate) provider_api_format: String,
pub(crate) stream_summary_report_context: Value,
pub(crate) response: DirectUpstreamResponse,
pub(crate) started_at: Instant,
}
@@ -226,11 +229,15 @@ impl DirectSyncExecutionRuntime {
let status_code = response.status().as_u16();
let headers = collect_response_headers(response.headers());
let stream_summary_report_context = build_stream_summary_report_context(&plan);
Ok(DirectUpstreamStreamExecution {
request_id: plan.request_id,
candidate_id: plan.candidate_id,
status_code,
headers,
provider_api_format: plan.provider_api_format.clone(),
stream_summary_report_context,
response: DirectUpstreamResponse::Reqwest(response),
started_at,
})
@@ -310,11 +317,20 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
provider_api_format: plan.provider_api_format.clone(),
stream_summary_report_context: build_stream_summary_report_context(plan),
response: DirectUpstreamResponse::LocalTunnel(response),
started_at,
}))
}
fn build_stream_summary_report_context(plan: &ExecutionPlan) -> Value {
json!({
"provider_api_format": plan.provider_api_format,
"model": plan.model_name,
})
}
async fn execute_sync_plan_via_local_tunnel(
state: &AppState,
plan: &ExecutionPlan,

14
apps/aether-gateway/src/handlers/admin/observability/usage/detail_routes.rs
View File

@@ -3,8 +3,8 @@ use super::analytics::admin_usage_provider_key_names;
use super::replay::{
admin_usage_curl_headers, admin_usage_curl_url, admin_usage_headers_from_value,
admin_usage_id_from_action_path, admin_usage_id_from_detail_path,
admin_usage_resolve_body_value, admin_usage_resolve_request_preview_body,
admin_usage_resolve_request_preview_body_for_item, build_admin_usage_curl_response,
admin_usage_resolve_body_value, admin_usage_resolve_request_capture_body,
admin_usage_resolve_request_capture_body_for_item, build_admin_usage_curl_response,
build_admin_usage_detail_payload, build_admin_usage_replay_response,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
@@ -101,9 +101,9 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
.await?;
let body = provider_request_body
.or(request_body)
.unwrap_or_else(|| admin_usage_resolve_request_preview_body(&item, None));
.or_else(|| admin_usage_resolve_request_capture_body(&item, None));
return Ok(Some(attach_admin_audit_response(
build_admin_usage_curl_response(&item, url, headers_json, &headers, &body),
build_admin_usage_curl_response(&item, url, headers_json, &headers, body.as_ref()),
"admin_usage_curl_viewed",
"view_usage_curl_replay",
"usage_record",
@@ -119,8 +119,8 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
{
response = attach_admin_audit_response(
response,
"admin_usage_replay_preview_generated",
"preview_usage_replay",
"admin_usage_replay_plan_generated",
"generate_usage_replay_plan",
"usage_record",
&usage_id,
);
@@ -173,7 +173,7 @@ pub(super) async fn maybe_build_local_admin_usage_detail_response(
let provider_key_name = admin_usage_provider_key_name(&item, &provider_key_names);
let request_body =
admin_usage_resolve_request_preview_body_for_item(state, &item, None).await?;
admin_usage_resolve_request_capture_body_for_item(state, &item, None).await?;
let mut detail_item = item.clone();
if include_bodies {
detail_item.provider_request_body = admin_usage_resolve_body_value(

35
apps/aether-gateway/src/handlers/admin/observability/usage/replay.rs
View File

@@ -6,7 +6,7 @@ use aether_admin::observability::usage::{
};
use aether_data_contracts::repository::{
provider_catalog::StoredProviderCatalogEndpoint,
usage::{StoredRequestUsageAudit, UsageBodyField},
usage::{StoredRequestUsageAudit, UsageBodyCaptureState, UsageBodyField},
};
use axum::{
body::Body,
@@ -37,11 +37,11 @@ struct AdminUsageReplayRequest {
body_override: Option<serde_json::Value>,
}
pub(super) fn admin_usage_resolve_request_preview_body(
pub(super) fn admin_usage_resolve_request_capture_body(
item: &StoredRequestUsageAudit,
body_override: Option<serde_json::Value>,
) -> serde_json::Value {
aether_admin::observability::usage::admin_usage_resolve_request_preview_body(
) -> Option<serde_json::Value> {
aether_admin::observability::usage::admin_usage_resolve_request_capture_body(
item,
body_override,
)
@@ -53,6 +53,15 @@ pub(super) async fn admin_usage_resolve_body_value(
inline_body: Option<&Value>,
field: UsageBodyField,
) -> Result<Option<Value>, GatewayError> {
match item.body_state(field) {
Some(UsageBodyCaptureState::Disabled)
| Some(UsageBodyCaptureState::Unavailable)
| Some(UsageBodyCaptureState::None) => return Ok(None),
Some(UsageBodyCaptureState::Inline) | Some(UsageBodyCaptureState::Truncated) => {
return Ok(inline_body.cloned());
}
Some(UsageBodyCaptureState::Reference) | None => {}
}
let resolved_ref_body = match item.body_ref(field) {
Some(body_ref) => state.resolve_request_usage_body_ref(body_ref).await?,
None => None,
@@ -70,13 +79,13 @@ fn admin_usage_body_value_from_sources(
resolved_ref_body.or_else(|| inline_body.cloned())
}
pub(super) async fn admin_usage_resolve_request_preview_body_for_item(
pub(super) async fn admin_usage_resolve_request_capture_body_for_item(
state: &AdminAppState<'_>,
item: &StoredRequestUsageAudit,
body_override: Option<serde_json::Value>,
) -> Result<serde_json::Value, GatewayError> {
) -> Result<Option<serde_json::Value>, GatewayError> {
if let Some(body_override) = body_override {
return Ok(body_override);
return Ok(Some(body_override));
}
if let Some(body) = admin_usage_resolve_body_value(
state,
@@ -86,9 +95,9 @@ pub(super) async fn admin_usage_resolve_request_preview_body_for_item(
)
.await?
{
return Ok(body);
return Ok(Some(body));
}
Ok(admin_usage_resolve_request_preview_body(item, None))
Ok(admin_usage_resolve_request_capture_body(item, None))
}
pub(super) fn build_admin_usage_curl_response(
@@ -96,7 +105,7 @@ pub(super) fn build_admin_usage_curl_response(
url: Option<String>,
headers_json: Option<Value>,
headers: &BTreeMap<String, String>,
body: &Value,
body: Option<&Value>,
) -> Response<Body> {
aether_admin::observability::usage::build_admin_usage_curl_response(
item,
@@ -115,7 +124,7 @@ pub(super) fn build_admin_usage_detail_payload(
auth_api_key_reader_available: bool,
provider_key_name: Option<&str>,
include_bodies: bool,
request_body: Value,
request_body: Option<Value>,
default_headers: &BTreeMap<String, String>,
) -> Value {
aether_admin::observability::usage::build_admin_usage_detail_payload(
@@ -136,7 +145,7 @@ pub(super) fn build_admin_usage_replay_plan_response(
target_provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
target_endpoint: &StoredProviderCatalogEndpoint,
target_api_key_id: Option<String>,
request_body: Value,
request_body: Option<Value>,
url: &str,
headers: &BTreeMap<String, String>,
same_provider: bool,
@@ -306,7 +315,7 @@ pub(super) async fn build_admin_usage_replay_response(
let same_provider = item.provider_id.as_deref() == Some(target_provider.id.as_str());
let same_endpoint = item.provider_endpoint_id.as_deref() == Some(target_endpoint.id.as_str());
let request_body =
admin_usage_resolve_request_preview_body_for_item(state, &item, payload.body_override)
admin_usage_resolve_request_capture_body_for_item(state, &item, payload.body_override)
.await?;
let url = admin_usage_curl_url(state, &target_endpoint, &item);

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

@@ -1436,9 +1436,6 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
assert_eq!(payload["client_response_headers"]["X-Request-Id"], "req-1");
assert!(payload["metadata"]["trace_id"].is_null());
assert_eq!(payload["trace"]["trace_id"], "trace-123");
assert!(payload["metadata"]["request_preview_source"].is_null());
assert!(payload["metadata"]["original_request_body_available"].is_null());
assert!(payload["metadata"]["original_response_body_available"].is_null());
assert!(payload["metadata"]["candidate_id"].is_null());
assert!(payload["metadata"]["candidate_index"].is_null());
assert!(payload["metadata"]["key_name"].is_null());
@@ -1460,7 +1457,7 @@ async fn gateway_handles_admin_usage_detail_locally_with_trusted_admin_principal
assert_eq!(payload["body_capture"]["request"]["storage"], "inline");
assert_eq!(payload["body_capture"]["request"]["available"], true);
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
payload["body_capture"]["request"]["capture_source"],
"stored_original"
);
assert_eq!(
@@ -1605,13 +1602,10 @@ async fn gateway_handles_admin_usage_detail_with_ref_backed_bodies() {
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert!(payload["metadata"]["trace_id"].is_null());
assert_eq!(payload["trace"]["trace_id"], "trace-ref-detail");
assert!(payload["metadata"]["request_preview_source"].is_null());
assert!(payload["metadata"]["original_request_body_available"].is_null());
assert!(payload["metadata"]["original_response_body_available"].is_null());
assert_eq!(payload["body_capture"]["request"]["storage"], "reference");
assert_eq!(payload["body_capture"]["request"]["available"], true);
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
payload["body_capture"]["request"]["capture_source"],
"stored_reference"
);
assert_eq!(
@@ -1812,15 +1806,12 @@ async fn gateway_handles_admin_usage_replay_locally_with_trusted_admin_principal
payload["request_headers"]["Content-Type"],
"application/json"
);
assert_eq!(payload["request_body"]["model"], "gpt-5");
assert_eq!(payload["request_body"]["target_model"], "gpt-5-target");
assert_eq!(payload["request_body"]["request_type"], "chat");
assert_eq!(payload["request_body"]["api_format"], "openai:chat");
assert_eq!(payload["request_body"]["stream"], false);
assert_eq!(payload["original_request_body_available"], false);
assert!(payload["request_body"].is_null());
assert_eq!(payload["captured_request_body_available"], false);
assert_eq!(payload["request_body_available"], false);
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
"local_reconstruction"
payload["body_capture"]["request"]["capture_source"],
"legacy_unknown"
);
assert_eq!(payload["body_capture"]["request"]["storage"], "missing");
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -1909,9 +1900,10 @@ async fn gateway_handles_admin_usage_replay_with_ref_backed_request_body() {
payload["request_body"]["messages"][0]["content"],
"replay from ref"
);
assert_eq!(payload["original_request_body_available"], true);
assert_eq!(payload["captured_request_body_available"], true);
assert_eq!(payload["request_body_available"], true);
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
payload["body_capture"]["request"]["capture_source"],
"stored_reference"
);
assert_eq!(payload["body_capture"]["request"]["storage"], "reference");
@@ -1975,8 +1967,8 @@ async fn local_admin_usage_replay_attaches_explicit_audit() {
.get::<AdminAuditEvent>()
.cloned()
.expect("usage replay should attach audit");
assert_eq!(audit.event_name, "admin_usage_replay_preview_generated");
assert_eq!(audit.action, "preview_usage_replay");
assert_eq!(audit.event_name, "admin_usage_replay_plan_generated");
assert_eq!(audit.action, "generate_usage_replay_plan");
assert_eq!(audit.target_type, "usage_record");
assert_eq!(audit.target_id, "usage-audit-replay");
}
@@ -2063,11 +2055,11 @@ async fn gateway_handles_admin_usage_curl_locally_with_trusted_admin_principal()
assert_eq!(payload["body"]["model"], "gpt-5-target");
assert_eq!(payload["body"]["temperature"], 0.2);
assert_eq!(payload["body"]["stream"], false);
assert_eq!(payload["original_request_body_available"], true);
assert_eq!(payload["captured_request_body_available"], true);
assert_eq!(payload["body_capture"]["body_source"], "provider_request");
assert_eq!(payload["body_capture"]["request"]["storage"], "inline");
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
payload["body_capture"]["request"]["capture_source"],
"stored_original"
);
assert_eq!(
@@ -2164,11 +2156,11 @@ async fn gateway_handles_admin_usage_curl_with_ref_backed_provider_request_body(
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["body"]["model"], "gpt-5-target");
assert_eq!(payload["body"]["temperature"], 0.2);
assert_eq!(payload["original_request_body_available"], true);
assert_eq!(payload["captured_request_body_available"], true);
assert_eq!(payload["body_capture"]["body_source"], "provider_request");
assert_eq!(payload["body_capture"]["request"]["storage"], "reference");
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
payload["body_capture"]["request"]["capture_source"],
"stored_reference"
);
assert_eq!(

312
apps/aether-gateway/src/tests/usage/local.rs
View File

@@ -12,6 +12,7 @@ use super::{
UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
};
use crate::constants::LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER;
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
fn deep_nested_metadata(levels: usize) -> serde_json::Value {
let mut current = json!({"leaf": "value"});
@@ -379,6 +380,144 @@ async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage()
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
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(|_request: Request| async move {
Json(json!({
"request_id": "trace-openai-chat-local-report-sync-request-limit-123",
"status_code": 200,
"headers": {
"content-type": "application/json"
},
"body": {
"json_body": {
"id": "chatcmpl-local-report-sync-request-limit-123",
"object": "chat.completion",
"model": "gpt-5-upstream",
"choices": [],
"usage": {
"prompt_tokens": 2,
"completion_tokens": 3,
"total_tokens": 5
}
}
},
"telemetry": {
"elapsed_ms": 25
}
}))
}),
);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(
"sk-client-openai-local-report-sync-request-limit",
)),
sample_local_openai_auth_snapshot(
"api-key-openai-usage-local-request-limit-1",
"user-openai-usage-local-request-limit-1",
),
)]));
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_system_config_values_for_tests([(
"max_request_body_size".to_string(),
json!(128),
)]),
)
.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_body = serde_json::to_string(&json!({
"model": "gpt-5",
"messages": [{
"role": "user",
"content": "x".repeat(2048)
}]
}))
.expect("request should encode");
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-local-report-sync-request-limit",
)
.header(
TRACE_ID_HEADER,
"trace-openai-chat-local-report-sync-request-limit-123",
)
.body(request_body)
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
"trace-openai-chat-local-report-sync-request-limit-123",
"completed",
)
.await;
assert_eq!(stored_usage.total_tokens, 5);
assert_eq!(
stored_usage.request_body_state,
Some(UsageBodyCaptureState::Truncated)
);
assert_eq!(
stored_usage.provider_request_body_state,
Some(UsageBodyCaptureState::Truncated)
);
assert_eq!(
stored_usage
.request_body
.as_ref()
.and_then(|value| value.get("truncated"))
.and_then(|value| value.as_bool()),
Some(true)
);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_strips_request_and_response_bodies_when_request_record_level_is_base() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
@@ -1097,6 +1236,179 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
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(|_request: Request| async move {
let delta_chunk = format!(
"data: {{\"id\":\"chatcmpl-local-report-stream-truncated-123\",\"choices\":[{{\"index\":0,\"delta\":{{\"content\":\"{}\"}}}}]}}\n\n",
"x".repeat(2048)
);
let summary = json!({
"standardized_usage": {
"input_tokens": 2,
"output_tokens": 4,
"cache_creation_tokens": 0,
"cache_creation_ephemeral_5m_tokens": 0,
"cache_creation_ephemeral_1h_tokens": 0,
"cache_read_tokens": 0,
"reasoning_tokens": 0,
"cache_storage_token_hours": 0.0,
"request_count": 1,
"dimensions": {}
},
"finish_reason": "stop",
"response_id": "chatcmpl-local-report-stream-truncated-123",
"model": "gpt-5-upstream",
"observed_finish": true
});
let frames = [
json!({
"type": "headers",
"payload": {
"kind": "headers",
"status_code": 200,
"headers": {"content-type": "text/event-stream"}
}
}),
json!({
"type": "data",
"payload": {"kind": "data", "text": delta_chunk}
}),
json!({
"type": "data",
"payload": {"kind": "data", "text": "data: [DONE]\\n\\n"}
}),
json!({
"type": "telemetry",
"payload": {
"kind": "telemetry",
"telemetry": {"elapsed_ms": 31, "ttfb_ms": 11}
}
}),
json!({
"type": "eof",
"payload": {"kind": "eof", "summary": summary}
}),
]
.into_iter()
.map(|frame| serde_json::to_string(&frame).expect("frame should encode"))
.collect::<Vec<_>>()
.join("\n")
+ "\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 auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key(
"sk-client-openai-local-report-stream-truncated",
)),
sample_local_openai_auth_snapshot(
"api-key-openai-usage-local-stream-truncated-1",
"user-openai-usage-local-stream-truncated-1",
),
)]));
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_system_config_values_for_tests([(
"max_response_body_size".to_string(),
json!(128),
)]),
)
.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 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-local-report-stream-truncated",
)
.header(
TRACE_ID_HEADER,
"trace-openai-chat-local-report-stream-truncated-123",
)
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let body_text = response.text().await.expect("stream body should read");
assert!(body_text.contains("chatcmpl-local-report-stream-truncated-123"));
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
"trace-openai-chat-local-report-stream-truncated-123",
"completed",
)
.await;
assert_eq!(stored_usage.total_tokens, 6);
assert_eq!(
stored_usage.response_body_state,
Some(UsageBodyCaptureState::Truncated)
);
assert_eq!(
stored_usage.client_response_body_state,
Some(UsageBodyCaptureState::Truncated)
);
assert_eq!(
stored_usage
.response_body
.as_ref()
.and_then(|value| value.get("truncated"))
.and_then(|value| value.as_bool()),
Some(true)
);
gateway_handle.abort();
execution_runtime_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_records_failed_usage_when_all_local_claude_cli_candidates_are_skipped() {
fn sample_auth_snapshot(api_key_id: &str, user_id: &str) -> StoredAuthApiKeySnapshot {

9
apps/aether-gateway/src/usage/reporting/mod.rs
View File

@@ -144,7 +144,10 @@ pub(crate) async fn submit_stream_report(
status_code: payload.status_code,
headers: payload.headers.clone(),
provider_body_base64: payload.provider_body_base64.clone(),
provider_body_state: payload.provider_body_state,
client_body_base64: payload.client_body_base64.clone(),
client_body_state: payload.client_body_state,
terminal_summary: payload.terminal_summary.clone(),
telemetry: payload.telemetry.clone(),
};
if should_handle_local_stream_report(
@@ -637,7 +640,10 @@ mod tests {
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: None,
telemetry: None,
},
)
@@ -751,7 +757,10 @@ mod tests {
status_code: 200,
headers: sample_codex_paid_headers(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: None,
telemetry: None,
},
)