i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 17:37:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(openai): unify tier authorization and settlement

Browse Source
This commit is contained in:
MMEXA
2026-07-11 12:27:05 +08:00
parent b2f596b8f0
commit 0b30cc6b0f
32 changed files with 3170 additions and 670 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/control/auth/gate.rs
+521 -271
View File
File diff suppressed because it is too large Load Diff
apps/aether-gateway/src/control/auth/mod.rs
+2 -2
View File
@@ -6,8 +6,8 @@ mod types;
pub(crate) use credentials::extract_requested_model;
pub(crate) use gate::{
request_model_local_rejection, should_buffer_request_for_local_auth,
trusted_auth_local_rejection, GatewayLocalAuthRejection,
execution_plan_balance_capacity_rejection, request_model_local_rejection,
should_buffer_request_for_local_auth, trusted_auth_local_rejection, GatewayLocalAuthRejection,
};
pub(crate) use resolution::{
refresh_execution_runtime_auth_context, resolve_execution_runtime_auth_context,
apps/aether-gateway/src/control/mod.rs
+2 -1
View File
@@ -8,7 +8,8 @@ mod public;
mod route;
pub(crate) use auth::{
extract_requested_model, refresh_execution_runtime_auth_context, request_model_local_rejection,
execution_plan_balance_capacity_rejection, extract_requested_model,
refresh_execution_runtime_auth_context, request_model_local_rejection,
resolve_execution_runtime_auth_context, should_buffer_request_for_local_auth,
trusted_auth_local_rejection, GatewayAdminPrincipalContext, GatewayControlAuthContext,
GatewayLocalAuthRejection,
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+8 -3
View File
@@ -5487,9 +5487,8 @@ async fn execute_stream_from_frame_stream(
telemetry = Some(frame_telemetry);
}
StreamFramePayload::Eof { summary } => {
if summary.is_some() {
stream_terminal_summary = summary;
}
stream_terminal_summary =
merge_stream_terminal_summary(stream_terminal_summary.take(), summary);
break;
}
StreamFramePayload::Error { error } => {
@@ -6296,12 +6295,14 @@ mod tests {
Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(runtime_usage),
model: Some("gpt-5.5".to_string()),
provider_actual_service_tier: Some("priority".to_string()),
unknown_event_count: 1,
..ExecutionStreamTerminalSummary::default()
}),
Some(ExecutionStreamTerminalSummary {
standardized_usage: Some(observed_usage),
response_id: Some("resp_123".to_string()),
provider_actual_service_tier: Some("default".to_string()),
observed_finish: true,
unknown_event_count: 2,
..ExecutionStreamTerminalSummary::default()
@@ -6316,6 +6317,10 @@ mod tests {
assert_eq!(usage.output_tokens, 137);
assert_eq!(merged.model.as_deref(), Some("gpt-5.5"));
assert_eq!(merged.response_id.as_deref(), Some("resp_123"));
assert_eq!(
merged.provider_actual_service_tier.as_deref(),
Some("default")
);
assert!(merged.observed_finish);
assert_eq!(merged.unknown_event_count, 3);
}
apps/aether-gateway/src/executor/candidate_loop.rs
+107 -33
View File
@@ -166,18 +166,31 @@ where
type Error = GatewayError;
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> {
prewarm_direct_reqwest_candidate_client(attempt.execution_plan());
let plan = attempt.execution_plan();
let report_context = attempt.report_context();
if let Some(response) = execution_plan_balance_capacity_response(
self.state,
self.trace_id,
self.decision,
plan,
report_context.as_ref(),
)
.await?
{
return Ok(Some(response));
}
prewarm_direct_reqwest_candidate_client(plan);
let _permit = acquire_upstream_execution_gate(self.state, self.trace_id).await?;
let upstream_execution_gate_held_started_at = std::time::Instant::now();
let mut response = execute_execution_runtime_sync(
self.state,
self.parts.uri.path(),
attempt.execution_plan().clone(),
plan.clone(),
self.trace_id,
self.decision,
self.plan_kind,
attempt.report_kind(),
attempt.report_context(),
report_context,
)
.await?;
observe_gateway_stage_ms(
@@ -187,10 +200,7 @@ where
.as_millis() as u64,
);
if let Some(response) = response.as_mut() {
attach_redaction_execution_candidate(
response,
attempt.execution_plan().candidate_id.as_deref(),
);
attach_redaction_execution_candidate(response, plan.candidate_id.as_deref());
}
Ok(response)
}
@@ -355,7 +365,14 @@ where
};
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
let execute_started_at = std::time::Instant::now();
let response = port.execute_attempt(&attempt).await?;
let response = match port.execute_attempt(&attempt).await {
Ok(response) => response,
Err(err) => {
let remaining = source.drain_execution_attempts().await?;
port.mark_unused_attempts(remaining).await?;
return Err(err);
}
};
observe_gateway_stage_ms(
"stream_candidate_execute",
execute_started_at.elapsed().as_millis() as u64,
@@ -450,6 +467,17 @@ where
candidate_index = candidate_index.as_str(),
"candidate loop attempting stream execution candidate"
);
if let Some(response) = execution_plan_balance_capacity_response(
self.state,
self.trace_id,
self.decision,
&plan,
report_context.as_ref(),
)
.await?
{
return Ok(Some(response));
}
prewarm_direct_reqwest_candidate_client(&plan);
let watchdog_plan = plan.clone();
let watchdog_report_context = report_context.clone();
@@ -523,45 +551,91 @@ fn prewarm_direct_reqwest_candidate_client(plan: &aether_contracts::ExecutionPla
);
}
async fn execution_plan_balance_capacity_response(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<Option<Response<Body>>, GatewayError> {
let rejection = match crate::control::execution_plan_balance_capacity_rejection(
state,
decision,
plan,
report_context,
)
.await
{
Ok(rejection) => rejection,
Err(err) => {
mark_unused_local_candidate(state, plan, report_context).await;
return Err(err);
}
};
let Some(rejection) = rejection else {
return Ok(None);
};
mark_unused_local_candidate(state, plan, report_context).await;
let mut response = crate::api::response::build_local_auth_rejection_response(
trace_id,
Some(decision),
&rejection,
)?;
attach_redaction_execution_candidate(&mut response, plan.candidate_id.as_deref());
Ok(Some(response))
}
pub(crate) async fn mark_unused_local_candidates<T>(state: &AppState, remaining: Vec<T>)
where
T: AiExecutionAttempt,
{
for plan_and_report in remaining {
let report_context = plan_and_report.report_context();
let metadata =
local_execution_candidate_metadata_from_report_context(report_context.as_ref());
if let Some(lease) = metadata.pool_key_lease.as_ref() {
if let Err(err) =
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), lease).await
{
warn!(
error = ?err,
"gateway candidate loop: failed to release unused pool key lease"
);
}
}
if should_skip_unused_persistence_from_metadata(&metadata) {
continue;
}
record_local_request_candidate_status(
mark_unused_local_candidate(
state,
plan_and_report.execution_plan(),
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Unused,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_ms: None,
finished_at_unix_ms: None,
},
)
.await;
}
}
async fn mark_unused_local_candidate(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
) {
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
if let Some(lease) = metadata.pool_key_lease.as_ref() {
if let Err(err) =
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), lease).await
{
warn!(
error = ?err,
"gateway candidate loop: failed to release unused pool key lease"
);
}
}
if should_skip_unused_persistence_from_metadata(&metadata) {
return;
}
record_local_request_candidate_status(
state,
plan,
report_context,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Unused,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
started_at_unix_ms: None,
finished_at_unix_ms: None,
},
)
.await;
}
fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool {
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
should_skip_unused_persistence_from_metadata(&metadata)
apps/aether-gateway/src/handlers/public/support/user_me_usage.rs
+6
View File
@@ -515,6 +515,9 @@ fn build_users_me_usage_record_payload(
if let Some(service_tier) = item.provider_service_tier() {
payload["service_tier"] = json!(service_tier);
}
if let Some(actual_service_tier) = item.provider_actual_service_tier() {
payload["actual_service_tier"] = json!(actual_service_tier);
}
if include_actual_cost {
payload["actual_cost"] = json!(round_to(item.actual_total_cost_usd, 6));
payload["rate_multiplier"] = json!(rate_multiplier);
@@ -582,6 +585,9 @@ fn build_users_me_usage_active_payload(item: &StoredRequestUsageAudit) -> serde_
if let Some(service_tier) = item.provider_service_tier() {
payload["service_tier"] = json!(service_tier);
}
if let Some(actual_service_tier) = item.provider_actual_service_tier() {
payload["actual_service_tier"] = json!(actual_service_tier);
}
payload
}
apps/aether-gateway/src/state/core.rs
+1
View File
@@ -709,6 +709,7 @@ impl AppState {
self.data.clear_minimal_candidate_selection_cache();
self.data.clear_routing_group_cache();
self.data.clear_provider_catalog_cache();
self.auth_request_cost_upper_bound_cache.clear();
self.routing_group_selection_cache.clear();
self.candidate_row_page_cache.clear();
self.candidate_page_cache.clear();