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

Merge remote-tracking branch 'origin/pr/538'

Browse Source
This commit is contained in:
fawney19
2026-05-21 23:12:42 +08:00
parent d6eb41aa78 cadc45c5b8
commit 0e6d7539ad
16 changed files with 422 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3,6 +3,7 @@ use std::io::Error as IoError;
use std::sync::Arc;
use std::time::{Duration, Instant};
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
ExecutionTelemetry,
@@ -1034,7 +1035,7 @@ fn resolve_openai_image_sync_total_timeout_ms(plan: &ExecutionPlan) -> u64 {
fn report_context_upstream_is_stream(report_context: Option<&Value>) -> bool {
report_context
.and_then(|value| value.get("upstream_is_stream"))
.and_then(|value| value.get(UPSTREAM_IS_STREAM_KEY))
.and_then(Value::as_bool)
.unwrap_or(false)
}

75
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -1,5 +1,6 @@
use aether_ai_serving::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
UPSTREAM_IS_STREAM_KEY,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{
@@ -471,11 +472,24 @@ fn should_skip_unused_persistence_from_metadata(
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some()
}
fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration {
fn resolve_stream_candidate_watchdog_timeout(
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Duration {
let upstream_is_stream = report_context
.and_then(|context| context.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or(true);
let timeout_ms = plan
.timeouts
.as_ref()
.and_then(|timeouts| timeouts.first_byte_ms.or(timeouts.total_ms))
.and_then(|timeouts| {
if upstream_is_stream {
timeouts.first_byte_ms.or(timeouts.total_ms)
} else {
timeouts.total_ms.or(timeouts.first_byte_ms)
}
})
.unwrap_or(DEFAULT_STREAM_CANDIDATE_WATCHDOG_TIMEOUT_MS)
.max(1);
Duration::from_millis(timeout_ms)
@@ -493,7 +507,7 @@ where
Fut:
std::future::Future<Output = Result<Option<Response<Body>>, GatewayError>> + Send + 'static,
{
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan);
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context);
let candidate_started_unix_ms = current_unix_ms();
let mut join_handle = tokio::spawn(execute());
match timeout(timeout_duration, &mut join_handle).await {
@@ -655,19 +669,22 @@ mod tests {
#[test]
fn stream_candidate_watchdog_prefers_first_byte_timeout() {
let timeout =
resolve_stream_candidate_watchdog_timeout(&test_plan(Some(ExecutionTimeouts {
let report_context = json!({"upstream_is_stream": true});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(12_345),
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})));
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(12_345));
}
#[test]
fn stream_candidate_watchdog_uses_default_when_timeouts_missing() {
let timeout = resolve_stream_candidate_watchdog_timeout(&test_plan(None));
let timeout = resolve_stream_candidate_watchdog_timeout(&test_plan(None), None);
assert_eq!(
timeout,
@@ -675,6 +692,50 @@ mod tests {
);
}
#[test]
fn stream_candidate_watchdog_prefers_total_timeout_when_upstream_non_stream() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
total_ms: Some(599_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(599_000));
}
#[test]
fn stream_candidate_watchdog_falls_back_to_first_byte_when_upstream_non_stream_lacks_total() {
let report_context = json!({"upstream_is_stream": false});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(300_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(300_000));
}
#[test]
fn stream_candidate_watchdog_defaults_to_streaming_when_flag_missing() {
let report_context = json!({});
let timeout = resolve_stream_candidate_watchdog_timeout(
&test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(12_345),
total_ms: Some(90_000),
..ExecutionTimeouts::default()
})),
Some(&report_context),
);
assert_eq!(timeout, Duration::from_millis(12_345));
}
#[test]
fn unused_persistence_skips_pool_internal_candidates() {
assert!(should_skip_unused_persistence(Some(&json!({

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

@@ -1,5 +1,6 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_ai_serving::UPSTREAM_IS_STREAM_KEY;
use aether_billing::{
normalize_input_tokens_for_billing, normalize_total_input_context_for_cache_hit_rate,
};
@@ -314,7 +315,7 @@ fn users_me_usage_upstream_is_stream(item: &StoredRequestUsageAudit) -> bool {
item.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("upstream_is_stream"))
.and_then(|metadata| metadata.get(UPSTREAM_IS_STREAM_KEY))
.and_then(serde_json::Value::as_bool)
.or_else(|| users_me_usage_headers_stream_flag(item.response_headers.as_ref()))
.or_else(|| users_me_usage_infer_upstream_stream_from_captured_bodies(item))