fix(codex): fence concurrent quota updates

This commit is contained in:
elky
2026-08-14 09:28:07 +08:00
parent f3a12c1008
commit 5b0c763086
65 changed files with 13009 additions and 738 deletions
@@ -9,12 +9,12 @@ use std::sync::{Arc, LazyLock, Mutex as StdMutex, OnceLock, RwLock as StdRwLock}
use std::time::{Duration, Instant};
use aether_contracts::{
ExecutionPlan, ExecutionResponseBodyMode, ExecutionResult, ExecutionTelemetry, ProxySnapshot,
ResolvedTransportProfile, ResponseBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
EXECUTION_RESPONSE_BODY_MODE_HEADER, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_BACKEND_REQWEST_RUSTLS, TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE,
TRANSPORT_HTTP_MODE_HTTP1_ONLY,
ExecutionPlan, ExecutionResponseBodyMode, ExecutionResponseObservation, ExecutionResult,
ExecutionTelemetry, ProxySnapshot, ResolvedTransportProfile, ResponseBody,
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER, EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER,
EXECUTION_REQUEST_HTTP1_ONLY_HEADER, EXECUTION_RESPONSE_BODY_MODE_HEADER,
TRANSPORT_BACKEND_BROWSER_WREQ, TRANSPORT_BACKEND_REQWEST_RUSTLS,
TRANSPORT_HTTP_MODE_H2C_PRIOR_KNOWLEDGE, TRANSPORT_HTTP_MODE_HTTP1_ONLY,
};
use aether_data::repository::proxy_nodes::ProxyNodeTrafficMutation;
use aether_http::{apply_http_client_config, HttpClientConfig};
@@ -691,6 +691,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) stream_precommit_committed: bool,
pub(crate) response: DirectUpstreamResponse,
pub(crate) started_at: Instant,
pub(crate) response_observation: ExecutionResponseObservation,
pub(crate) stream_first_byte_timeout: Option<Duration>,
pub(crate) upstream_target_permit: Option<UpstreamTargetAdmissionPermit>,
}
@@ -699,6 +700,7 @@ pub(crate) struct DirectUpstreamStreamExecution {
pub(crate) struct DirectSyncResponseStarted {
pub(crate) status_code: u16,
pub(crate) ttfb_ms: u64,
pub(crate) response_observation: ExecutionResponseObservation,
}
impl DirectSyncExecutionRuntime {
@@ -724,14 +726,23 @@ impl DirectSyncExecutionRuntime {
let body_bytes = build_request_body(plan)?;
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
with_non_stream_total_timeout(plan, async move {
let response = send_request_inner(plan, body_bytes, false).await?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
let status_code = response.status_code();
let headers = response.headers();
let response_observation = ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
};
on_response_started(DirectSyncResponseStarted {
status_code,
ttfb_ms,
response_observation: response_observation.clone(),
});
let (body_bytes, stream_ttfb_ms) =
response.bytes_with_stream_timeout(plan, started_at).await?;
@@ -752,6 +763,7 @@ impl DirectSyncExecutionRuntime {
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: Some(response_observation),
body,
telemetry: Some(ExecutionTelemetry {
ttfb_ms: stream_ttfb_ms.or(Some(ttfb_ms)),
@@ -776,6 +788,8 @@ impl DirectSyncExecutionRuntime {
);
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let response = send_request(plan, body_bytes).await?;
observe_gateway_stage_ms(
"direct_send_headers",
@@ -783,6 +797,7 @@ impl DirectSyncExecutionRuntime {
);
let status_code = response.status_code();
let headers = response.headers();
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
let stream_summary_report_context = build_stream_summary_report_context(plan);
@@ -797,6 +812,11 @@ impl DirectSyncExecutionRuntime {
stream_precommit_committed: false,
response: response.into_direct_upstream_response(),
started_at,
response_observation: ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
},
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
upstream_target_permit: None,
})
@@ -834,7 +854,7 @@ pub(crate) async fn execute_sync_plan_with_report_context(
}
if resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).is_some() {
return execute_sync_plan_via_local_tunnel(state, plan)
return execute_sync_plan_via_local_tunnel(state, plan, report_context)
.await
.map_err(|err| GatewayError::Internal(err.to_string()));
}
@@ -857,7 +877,24 @@ pub(crate) async fn execute_sync_plan_with_report_context(
Ok(None) => {}
Err(err) => return Err(GatewayError::Internal(err.to_string())),
}
match DirectSyncExecutionRuntime::new().execute_sync(plan).await {
let state_for_response_started = state.clone();
match DirectSyncExecutionRuntime::new()
.execute_sync_with_response_started(plan, move |event| {
crate::orchestration::spawn_local_oauth_success_effect(
state_for_response_started,
plan,
report_context,
crate::orchestration::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
{
Ok(result) => {
record_manual_proxy_request_outcome(state, plan, result.status_code).await;
Ok(result)
@@ -889,6 +926,8 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
plan.body.body_bytes_b64.is_some(),
)?;
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let response = state
.tunnel
.open_direct_relay_stream(
@@ -900,6 +939,7 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let status_code = response.status();
let headers = collect_tunnel_response_headers(response.headers());
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
Ok(Some(DirectUpstreamStreamExecution {
request_id: plan.request_id.clone(),
@@ -912,6 +952,11 @@ pub(crate) async fn execute_stream_plan_via_local_tunnel(
stream_precommit_committed: false,
response: DirectUpstreamResponse::LocalTunnel(response),
started_at,
response_observation: ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
},
stream_first_byte_timeout: resolve_stream_first_byte_timeout(plan),
upstream_target_permit: None,
}))
@@ -991,13 +1036,19 @@ fn manual_proxy_node_id(proxy: Option<&ProxySnapshot>) -> Option<String> {
async fn execute_sync_plan_via_local_tunnel(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
with_non_stream_total_timeout(plan, execute_sync_plan_via_local_tunnel_inner(state, plan)).await
with_non_stream_total_timeout(
plan,
execute_sync_plan_via_local_tunnel_inner(state, plan, report_context),
)
.await
}
async fn execute_sync_plan_via_local_tunnel_inner(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
let node_id = resolve_local_tunnel_node_id(state, plan.proxy.as_ref()).ok_or_else(|| {
ExecutionRuntimeTransportError::RelayError("local tunnel node unavailable".to_string())
@@ -1030,6 +1081,8 @@ async fn execute_sync_plan_via_local_tunnel_inner(
"gateway execution runtime local tunnel request prepared"
);
let started_at = Instant::now();
let request_started_at_unix_ms = crate::clock::current_unix_ms();
let request_order_id = uuid::Uuid::now_v7().to_string();
let mut response = state
.tunnel
.open_direct_relay_stream(
@@ -1040,8 +1093,24 @@ async fn execute_sync_plan_via_local_tunnel_inner(
.await
.map_err(ExecutionRuntimeTransportError::RelayError)?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let response_headers_observed_at_unix_ms = crate::clock::current_unix_ms();
let status_code = response.status();
let headers = collect_tunnel_response_headers(response.headers());
let response_observation = ExecutionResponseObservation {
request_started_at_unix_ms,
response_headers_observed_at_unix_ms,
request_order_id,
};
crate::orchestration::spawn_local_oauth_success_effect(
state.clone(),
plan,
report_context,
crate::orchestration::LocalOAuthSuccessEffect {
status_code,
request_started_at_unix_ms: Some(response_observation.request_started_at_unix_ms),
request_order_id: Some(&response_observation.request_order_id),
},
);
let proxy_timing = execution_header_for_log(&headers, "x-proxy-timing").unwrap_or("-");
let (body_bytes, stream_ttfb_ms) =
collect_local_tunnel_response_body(response, plan, started_at).await?;
@@ -1095,6 +1164,7 @@ async fn execute_sync_plan_via_local_tunnel_inner(
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: Some(response_observation),
body,
telemetry: Some(ExecutionTelemetry {
ttfb_ms: stream_ttfb_ms.or(Some(ttfb_ms)),
@@ -5605,6 +5675,8 @@ mod tests {
)
.await
.expect("headers should write");
socket.flush().await.expect("headers should flush");
tokio::time::sleep(std::time::Duration::from_millis(40)).await;
socket
.write_all(b"b\r\ndata: one\n\n\r\n")
.await
@@ -5634,12 +5706,34 @@ mod tests {
let body = result
.body
.clone()
.and_then(|body| body.body_bytes_b64)
.and_then(|body| base64::engine::general_purpose::STANDARD.decode(body).ok())
.expect("stream body should be captured as bytes");
let body = String::from_utf8(body).expect("stream body should be utf8");
assert!(body.contains("data: one"));
assert!(body.contains("data: two"));
let observation = result
.response_observation
.expect("stream sync execution should preserve header observation");
let telemetry = result
.telemetry
.expect("stream sync execution should include telemetry");
let ttfb_ms = telemetry
.ttfb_ms
.expect("stream sync execution should measure the first body byte");
assert!(
observation.response_headers_observed_at_unix_ms
>= observation.request_started_at_unix_ms
);
assert!(
observation
.response_headers_observed_at_unix_ms
.saturating_sub(observation.request_started_at_unix_ms)
< ttfb_ms,
"header observation must not be derived from body-byte ttfb"
);
assert!(!observation.request_order_id.is_empty());
}
#[tokio::test]