mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-09 10:57:03 +08:00
fix(codex): fence concurrent quota updates
This commit is contained in:
@@ -5,8 +5,8 @@ use std::time::{Duration, Instant};
|
||||
|
||||
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope, UPSTREAM_IS_STREAM_KEY};
|
||||
use aether_contracts::{
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionResult,
|
||||
ExecutionTelemetry,
|
||||
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan,
|
||||
ExecutionResponseObservation, ExecutionResult, ExecutionTelemetry,
|
||||
};
|
||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||
use aether_scheduler_core::{
|
||||
@@ -70,11 +70,12 @@ use crate::execution_runtime::{
|
||||
};
|
||||
use crate::log_ids::short_request_id;
|
||||
use crate::orchestration::{
|
||||
apply_local_execution_effect, build_local_error_flow_metadata, trace_upstream_response_body,
|
||||
with_error_flow_report_context, with_upstream_response_report_context,
|
||||
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
|
||||
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
|
||||
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
|
||||
apply_local_execution_effect, build_local_error_flow_metadata,
|
||||
spawn_local_oauth_success_effect, trace_upstream_response_body, with_error_flow_report_context,
|
||||
with_upstream_response_report_context, LocalAdaptiveRateLimitEffect,
|
||||
LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
|
||||
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalHealthSuccessEffect,
|
||||
LocalOAuthInvalidationEffect, LocalOAuthSuccessEffect, LocalPoolErrorEffect,
|
||||
};
|
||||
use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard;
|
||||
use crate::request_candidate_runtime::{
|
||||
@@ -1379,7 +1380,19 @@ async fn execute_direct_sync_runtime_candidate(
|
||||
candidate_started_unix_ms,
|
||||
event.status_code,
|
||||
event.ttfb_ms,
|
||||
)
|
||||
);
|
||||
spawn_local_oauth_success_effect(
|
||||
state_for_response_started.clone(),
|
||||
plan,
|
||||
report_context,
|
||||
LocalOAuthSuccessEffect {
|
||||
status_code: event.status_code,
|
||||
request_started_at_unix_ms: Some(
|
||||
event.response_observation.request_started_at_unix_ms,
|
||||
),
|
||||
request_order_id: Some(&event.response_observation.request_order_id),
|
||||
},
|
||||
);
|
||||
})
|
||||
.await
|
||||
.map_err(SyncExecutionFailure::from_transport);
|
||||
@@ -1483,12 +1496,25 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
OpenAiImageSyncProgressRecorder::new(state, plan, report_context, progress_snapshot);
|
||||
progress.record_connecting().await;
|
||||
|
||||
let request_started_at_unix_ms = current_request_candidate_unix_ms();
|
||||
let request_order_id = uuid::Uuid::now_v7().to_string();
|
||||
let response = send_request(plan, request_body)
|
||||
.await
|
||||
.map_err(SyncExecutionFailure::from_transport)?;
|
||||
let ttfb_ms = started_at.elapsed().as_millis() as u64;
|
||||
let response_headers_observed_at_unix_ms = current_request_candidate_unix_ms();
|
||||
let status_code = response.status_code();
|
||||
let headers = response.headers();
|
||||
spawn_local_oauth_success_effect(
|
||||
state.clone(),
|
||||
plan,
|
||||
report_context,
|
||||
LocalOAuthSuccessEffect {
|
||||
status_code,
|
||||
request_started_at_unix_ms: Some(request_started_at_unix_ms),
|
||||
request_order_id: Some(&request_order_id),
|
||||
},
|
||||
);
|
||||
progress.record_response_started(status_code, ttfb_ms).await;
|
||||
|
||||
let mut body_bytes = Vec::new();
|
||||
@@ -1569,6 +1595,11 @@ async fn execute_openai_image_sync_upstream_sse_candidate(
|
||||
candidate_id: plan.candidate_id.clone(),
|
||||
status_code,
|
||||
headers,
|
||||
response_observation: Some(ExecutionResponseObservation {
|
||||
request_started_at_unix_ms,
|
||||
response_headers_observed_at_unix_ms,
|
||||
request_order_id,
|
||||
}),
|
||||
body,
|
||||
telemetry: Some(ExecutionTelemetry {
|
||||
ttfb_ms: Some(ttfb_ms),
|
||||
@@ -2461,6 +2492,16 @@ async fn execute_execution_runtime_sync_impl(
|
||||
};
|
||||
let mut candidate_first_byte_elapsed_ms =
|
||||
calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result);
|
||||
let initial_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
||||
let mut provider_response_observation =
|
||||
result
|
||||
.response_observation
|
||||
.clone()
|
||||
.unwrap_or(ExecutionResponseObservation {
|
||||
request_started_at_unix_ms: candidate_started_unix_secs,
|
||||
response_headers_observed_at_unix_ms: initial_response_observed_at_unix_ms,
|
||||
request_order_id: uuid::Uuid::now_v7().to_string(),
|
||||
});
|
||||
let mut oauth_retry_attempted = false;
|
||||
let (
|
||||
result_error_type,
|
||||
@@ -2473,6 +2514,18 @@ async fn execute_execution_runtime_sync_impl(
|
||||
local_failover_response_text,
|
||||
local_failover_analysis,
|
||||
) = loop {
|
||||
spawn_local_oauth_success_effect(
|
||||
state.clone(),
|
||||
&plan,
|
||||
report_context.as_ref(),
|
||||
LocalOAuthSuccessEffect {
|
||||
status_code: result.status_code,
|
||||
request_started_at_unix_ms: Some(
|
||||
provider_response_observation.request_started_at_unix_ms,
|
||||
),
|
||||
request_order_id: Some(&provider_response_observation.request_order_id),
|
||||
},
|
||||
);
|
||||
let result_latency_ms = result
|
||||
.telemetry
|
||||
.as_ref()
|
||||
@@ -2534,10 +2587,15 @@ async fn execute_execution_runtime_sync_impl(
|
||||
result.status_code,
|
||||
local_failover_response_text.as_deref(),
|
||||
trace_id,
|
||||
report_context.as_ref(),
|
||||
Some(provider_response_observation.request_started_at_unix_ms),
|
||||
Some(&provider_response_observation.request_order_id),
|
||||
)
|
||||
.await
|
||||
{
|
||||
oauth_retry_attempted = true;
|
||||
let retry_started_at_unix_ms = current_request_candidate_unix_ms();
|
||||
let retry_request_order_id = uuid::Uuid::now_v7().to_string();
|
||||
match crate::execution_runtime::execute_execution_runtime_sync_plan(
|
||||
state,
|
||||
Some(trace_id),
|
||||
@@ -2546,6 +2604,16 @@ async fn execute_execution_runtime_sync_impl(
|
||||
.await
|
||||
{
|
||||
Ok(retry_result) => {
|
||||
let retry_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
||||
provider_response_observation = retry_result
|
||||
.response_observation
|
||||
.clone()
|
||||
.unwrap_or(ExecutionResponseObservation {
|
||||
request_started_at_unix_ms: retry_started_at_unix_ms,
|
||||
response_headers_observed_at_unix_ms:
|
||||
retry_response_observed_at_unix_ms,
|
||||
request_order_id: retry_request_order_id,
|
||||
});
|
||||
candidate_first_byte_elapsed_ms =
|
||||
calibrated_sync_candidate_first_byte_elapsed_ms(
|
||||
candidate_started_at,
|
||||
@@ -2594,6 +2662,13 @@ async fn execute_execution_runtime_sync_impl(
|
||||
local_failover_analysis,
|
||||
);
|
||||
};
|
||||
let mut report_context = attach_provider_response_headers_to_report_context(
|
||||
report_context,
|
||||
&headers,
|
||||
provider_response_observation.request_started_at_unix_ms,
|
||||
provider_response_observation.response_headers_observed_at_unix_ms,
|
||||
&provider_response_observation.request_order_id,
|
||||
);
|
||||
if result.status_code >= 400 {
|
||||
apply_local_execution_effect(
|
||||
state,
|
||||
@@ -2739,8 +2814,6 @@ async fn execute_execution_runtime_sync_impl(
|
||||
}
|
||||
let status_code = result.status_code;
|
||||
let has_body_bytes = body_base64.is_some();
|
||||
let mut report_context =
|
||||
attach_provider_response_headers_to_report_context(report_context, &headers);
|
||||
if (200..300).contains(&status_code) {
|
||||
seed_kiro_sync_simulated_cache_enabled(state, &plan, &mut report_context).await;
|
||||
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
|
||||
@@ -3231,6 +3304,8 @@ async fn execute_sync_via_remote_execution_runtime(
|
||||
candidate_started_unix_secs: u64,
|
||||
candidate_started_at: Instant,
|
||||
) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
|
||||
let remote_request_started_at_unix_ms = current_request_candidate_unix_ms();
|
||||
let remote_request_order_id = uuid::Uuid::now_v7().to_string();
|
||||
let response = match post_sync_plan_to_remote_execution_runtime(
|
||||
state,
|
||||
remote_execution_runtime_base_url,
|
||||
@@ -3299,11 +3374,19 @@ async fn execute_sync_via_remote_execution_runtime(
|
||||
));
|
||||
}
|
||||
|
||||
response
|
||||
.json()
|
||||
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
|
||||
let mut result = response
|
||||
.json::<ExecutionResult>()
|
||||
.await
|
||||
.map(RemoteSyncFallbackOutcome::Executed)
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
result
|
||||
.response_observation
|
||||
.get_or_insert(ExecutionResponseObservation {
|
||||
request_started_at_unix_ms: remote_request_started_at_unix_ms,
|
||||
response_headers_observed_at_unix_ms: remote_response_observed_at_unix_ms,
|
||||
request_order_id: remote_request_order_id,
|
||||
});
|
||||
Ok(RemoteSyncFallbackOutcome::Executed(result))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
Reference in New Issue
Block a user