i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-12 22:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(codex): fence concurrent quota updates

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/execution_runtime/chatgpt_web_image.rs
+3
View File
@@ -2573,6 +2573,7 @@ fn json_execution_result(
candidate_id: plan.candidate_id.clone(),
status_code,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(body),
body_bytes_b64: None,
@@ -2615,6 +2616,7 @@ fn bytes_execution_result(
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
response_observation: None,
body: Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)),
@@ -2637,6 +2639,7 @@ fn execution_result_frame_stream(
payload: StreamFramePayload::Headers {
status_code: result.status_code,
headers: result.headers.clone(),
response_observation: result.response_observation.clone(),
},
},
StreamFrame {
apps/aether-gateway/src/execution_runtime/fallback.rs
+12
View File
@@ -480,6 +480,7 @@ mod tests {
candidate_id: None,
status_code: 502,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -519,6 +520,7 @@ mod tests {
candidate_id: None,
status_code: 502,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -583,6 +585,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -614,6 +617,7 @@ mod tests {
candidate_id: None,
status_code: 401,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -645,6 +649,7 @@ mod tests {
candidate_id: None,
status_code: 502,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: Some(ExecutionError {
@@ -708,6 +713,7 @@ mod tests {
candidate_id: None,
status_code: 404,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -900,6 +906,7 @@ mod tests {
candidate_id: None,
status_code: 200,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1022,6 +1029,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1068,6 +1076,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1174,6 +1183,7 @@ mod tests {
candidate_id: None,
status_code: 200,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1211,6 +1221,7 @@ mod tests {
candidate_id: None,
status_code: 400,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -1259,6 +1270,7 @@ mod tests {
candidate_id: None,
status_code: 429,
headers: Default::default(),
response_observation: None,
body: None,
telemetry: None,
error: None,
apps/aether-gateway/src/execution_runtime/grok.rs
+3
View File
@@ -841,6 +841,7 @@ fn encode_grok_headers_frame(
payload: StreamFramePayload::Headers {
status_code,
headers,
response_observation: None,
},
})
}
@@ -2157,6 +2158,7 @@ fn grok_execution_result(
candidate_id: plan.candidate_id.clone(),
status_code,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
@@ -2220,6 +2222,7 @@ fn grok_collected_frame_stream(
"application/json".to_string()
},
)]),
response_observation: None,
},
},
StreamFrame {
apps/aether-gateway/src/execution_runtime/kiro_web_search.rs
+2
View File
@@ -279,6 +279,7 @@ fn raw_response_frame_stream(
payload: StreamFramePayload::Headers {
status_code,
headers,
response_observation: None,
},
},
StreamFrame {
@@ -1449,6 +1450,7 @@ mod tests {
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"jsonrpc": "2.0",
apps/aether-gateway/src/execution_runtime/oauth_retry.rs
+105 -18
View File
@@ -2,10 +2,10 @@ use aether_contracts::ExecutionPlan;
use tracing::warn;
use crate::orchestration::{
oauth_status_may_be_invalid as status_may_be_oauth_invalid,
local_failover_error_message, oauth_status_may_be_invalid as status_may_be_oauth_invalid,
oauth_status_proves_access_token_invalid as status_proves_access_token_invalid,
};
use crate::state::AgentIdentityAuthConfigFence;
use crate::state::{AgentIdentityAuthConfigFence, CodexRuntimeOAuthObservation};
use crate::{provider_transport::LocalOAuthRefreshError, AppState};
pub(crate) async fn refresh_oauth_plan_auth_for_retry(
@@ -14,6 +14,9 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
status_code: u16,
response_text: Option<&str>,
trace_id: &str,
report_context: Option<&serde_json::Value>,
request_started_at_unix_ms: Option<u64>,
request_order_id: Option<&str>,
) -> bool {
if !status_may_be_oauth_invalid(status_code, response_text) {
return false;
@@ -109,15 +112,49 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
body_excerpt,
..
}) if matches!(refresh_status_code, 400 | 401 | 403) => {
if let Err(err) = state
.persist_local_oauth_refresh_failure_state(
&transport,
refresh_status_code,
body_excerpt.as_str(),
access_token_invalid_proven,
)
.await
{
let observed_credential_generation =
report_context_string(report_context, "codex_credential_generation");
let runtime_invalid_message = local_failover_error_message(response_text);
let runtime_invalid_reason =
aether_admin::provider::quota::codex_runtime_invalid_reason(
status_code,
runtime_invalid_message.as_deref(),
);
let persist_result = match (request_started_at_unix_ms, request_order_id) {
(Some(request_started_at_unix_ms), Some(request_order_id))
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex") =>
{
state
.persist_local_oauth_refresh_failure_state_observed(
&transport,
refresh_status_code,
body_excerpt.as_str(),
access_token_invalid_proven,
CodexRuntimeOAuthObservation {
request_started_at_unix_ms,
request_order_id,
observed_credential_generation,
runtime_invalid_reason: runtime_invalid_reason.as_deref(),
},
)
.await
}
_ => {
state
.persist_local_oauth_refresh_failure_state(
&transport,
refresh_status_code,
body_excerpt.as_str(),
access_token_invalid_proven,
)
.await
}
};
if let Err(err) = persist_result {
warn!(
event_name = "local_oauth_retry_refresh_failure_persist_failed",
log_type = "ops",
@@ -161,6 +198,17 @@ pub(crate) async fn refresh_oauth_plan_auth_for_retry(
}
}
fn report_context_string<'a>(
report_context: Option<&'a serde_json::Value>,
field: &str,
) -> Option<&'a str> {
report_context
.and_then(|context| context.get(field))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn execution_plan_authorization(plan: &ExecutionPlan) -> Option<&str> {
plan.headers
.iter()
@@ -209,6 +257,7 @@ mod tests {
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -312,7 +361,7 @@ mod tests {
}
#[tokio::test]
async fn auto_removes_codex_key_after_request_proven_terminal_refresh_failure() {
async fn retains_codex_key_after_request_proven_terminal_refresh_failure() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
@@ -458,16 +507,34 @@ mod tests {
401,
Some(r#"{"error":"oauth_token_invalid"}"#),
"trace-oauth-retry",
None,
Some(1_000),
Some("01900000-0000-7000-8000-000000000010"),
)
.await;
assert!(!retried);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let keys = provider_catalog_repository
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-retry".to_string()])
.await
.expect("keys should read");
assert!(keys.is_empty());
.expect("keys should read")
.into_iter()
.next()
.expect("request-scoped refresh failure should retain the key");
let invalid_reason = stored_key
.oauth_invalid_reason
.as_deref()
.expect("combined invalid reason should persist");
assert!(invalid_reason.contains("[OAUTH_EXPIRED]"));
assert!(invalid_reason.contains("[REFRESH_FAILED]"));
assert_eq!(
stored_key
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/oauth_state_request_id")),
Some(&json!("01900000-0000-7000-8000-000000000010"))
);
token_handle.abort();
}
@@ -619,6 +686,9 @@ mod tests {
401,
Some(r#"{"error":"invalid_token"}"#),
"trace-claude-oauth-fence-first",
None,
None,
None,
)
.await
);
@@ -647,6 +717,9 @@ mod tests {
401,
Some(r#"{"error":"invalid_token"}"#),
"trace-claude-oauth-fence-stale",
None,
None,
None,
)
.await
);
@@ -665,6 +738,7 @@ mod tests {
.expect("Claude key should load")
.pop()
.expect("Claude key should exist");
let expected_admin_replacement = admin_replacement.clone();
admin_replacement.encrypted_api_key = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
@@ -673,10 +747,23 @@ mod tests {
.expect("admin access token should encrypt"),
);
admin_replacement.expires_at_unix_secs = Some(4_102_444_800);
provider_catalog_repository
.update_key(&admin_replacement)
assert!(provider_catalog_repository
.compare_and_update_key_admin_state(&ProviderCatalogKeyAdminCasUpdate {
expected_encrypted_auth_config: expected_admin_replacement
.encrypted_auth_config
.clone(),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: expected_admin_replacement.encrypted_api_key.clone(),
auth_type: expected_admin_replacement.auth_type.clone(),
provider_id: expected_admin_replacement.provider_id.clone(),
provider_type: "claude_code".to_string(),
},
key: admin_replacement,
codex_rotation: None,
reset_oauth_runtime: true,
})
.await
.expect("admin replacement should persist");
.expect("admin replacement CAS should run"));
let admin_result = state
.force_local_oauth_refresh_entry(&stale_transport)
apps/aether-gateway/src/execution_runtime/response_header_rules.rs
+101 -4
View File
@@ -10,6 +10,10 @@ use crate::{AppState, GatewayError};
const RESPONSE_HEADER_RULES_KEY: &str = "response_header_rules";
const RESPONSE_HEADER_RULES_CAMEL_KEY: &str = "responseHeaderRules";
const PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY: &str = "provider_response_headers";
const PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY: &str = "provider_request_started_at_unix_ms";
const PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY: &str = "provider_request_order_id";
const PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY: &str =
"provider_response_headers_observed_at_unix_ms";
const RESPONSE_HEADER_RULE_PROTECTED_KEYS: &[&str] = &["content-length"];
const RESPONSE_HEADER_RULES_CACHE_TTL: Duration = Duration::from_secs(5);
@@ -98,6 +102,9 @@ pub(crate) async fn apply_endpoint_response_header_rules(
pub(crate) fn attach_provider_response_headers_to_report_context(
report_context: Option<Value>,
provider_headers: &BTreeMap<String, String>,
provider_request_started_at_unix_ms: u64,
provider_response_headers_observed_at_unix_ms: u64,
provider_request_order_id: &str,
) -> Option<Value> {
let provider_headers = serde_json::to_value(provider_headers).ok()?;
let mut object = match report_context {
@@ -105,9 +112,99 @@ pub(crate) fn attach_provider_response_headers_to_report_context(
Some(other) => Map::from_iter([("seed".to_string(), other)]),
None => Map::new(),
};
object.insert(
PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(),
provider_headers,
);
let observation_is_absent = !object.contains_key(PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY)
&& !object.contains_key(PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY)
&& !object.contains_key(PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY)
&& !object.contains_key(PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY);
if observation_is_absent {
object.insert(
PROVIDER_RESPONSE_HEADERS_CONTEXT_KEY.to_string(),
provider_headers,
);
object.insert(
PROVIDER_REQUEST_STARTED_AT_UNIX_MS_CONTEXT_KEY.to_string(),
Value::from(provider_request_started_at_unix_ms),
);
object.insert(
PROVIDER_RESPONSE_HEADERS_OBSERVED_AT_UNIX_MS_CONTEXT_KEY.to_string(),
Value::from(provider_response_headers_observed_at_unix_ms),
);
object.insert(
PROVIDER_REQUEST_ORDER_ID_CONTEXT_KEY.to_string(),
Value::from(provider_request_order_id),
);
}
Some(Value::Object(object))
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn provider_response_observation_is_first_write_wins() {
let first_headers =
BTreeMap::from([("x-codex-primary-used-percent".to_string(), "10".to_string())]);
let second_headers =
BTreeMap::from([("x-codex-primary-used-percent".to_string(), "20".to_string())]);
let report_context = attach_provider_response_headers_to_report_context(
Some(json!("seed-value")),
&first_headers,
100,
200,
"observation-1",
);
let report_context = attach_provider_response_headers_to_report_context(
report_context,
&second_headers,
300,
400,
"observation-2",
)
.expect("report context should exist");
assert_eq!(report_context["seed"], json!("seed-value"));
assert_eq!(
report_context["provider_response_headers"]["x-codex-primary-used-percent"],
json!("10")
);
assert_eq!(
report_context["provider_request_started_at_unix_ms"],
json!(100)
);
assert_eq!(
report_context["provider_response_headers_observed_at_unix_ms"],
json!(200)
);
assert_eq!(
report_context["provider_request_order_id"],
json!("observation-1")
);
}
#[test]
fn provider_response_observation_does_not_complete_a_partial_triplet() {
let report_context = attach_provider_response_headers_to_report_context(
Some(json!({"provider_response_headers": {"x-existing": "1"}})),
&BTreeMap::from([("x-new".to_string(), "2".to_string())]),
300,
400,
"observation-2",
)
.expect("report context should exist");
assert_eq!(
report_context["provider_response_headers"]["x-existing"],
json!("1")
);
assert!(report_context
.get("provider_request_started_at_unix_ms")
.is_none());
assert!(report_context
.get("provider_response_headers_observed_at_unix_ms")
.is_none());
assert!(report_context.get("provider_request_order_id").is_none());
}
}
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+99 -8
View File
@@ -11,8 +11,8 @@ use std::time::{Duration, Instant};
use aether_ai_serving::{AiAttemptExecutionOutcome, AiAttemptRetryScope};
use aether_contracts::{
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StandardizedUsage,
StreamFrame, StreamFramePayload,
ExecutionPlan, ExecutionResponseObservation, ExecutionStreamTerminalSummary,
ExecutionTelemetry, StandardizedUsage, StreamFrame, StreamFramePayload,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, UpsertRequestCandidateRecord,
@@ -112,12 +112,13 @@ use crate::execution_runtime::{
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, build_local_error_flow_metadata, classify_failure_disposition,
cyber_continue_failover_enabled, trace_upstream_response_body, with_error_flow_report_context,
cyber_continue_failover_enabled, spawn_local_oauth_success_effect,
trace_upstream_response_body, with_error_flow_report_context,
with_upstream_response_report_context, FailureDisposition, FailureTokenAction,
LocalAdaptiveRateLimitEffect, LocalAdaptiveSuccessEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalFailoverAnalysis,
LocalHealthFailureEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect,
LocalPoolErrorEffect,
LocalOAuthSuccessEffect, LocalPoolErrorEffect,
};
use crate::provider_pool_demand::{
acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
@@ -1249,6 +1250,9 @@ async fn execute_in_process_stream_with_oauth_retry(
retry_status_code,
response_text.as_deref(),
trace_id,
report_context,
Some(execution.response_observation.request_started_at_unix_ms),
Some(&execution.response_observation.request_order_id),
)
.await
{
@@ -2818,6 +2822,7 @@ async fn execute_stream_from_direct_passthrough(
stream_precommit_committed: _,
response,
started_at: upstream_started_at,
response_observation,
stream_first_byte_timeout,
upstream_target_permit,
} = execution;
@@ -2834,8 +2839,23 @@ async fn execute_stream_from_direct_passthrough(
let request_id = plan.request_id.clone();
let candidate_id = plan.candidate_id.clone();
let request_id_for_log = short_request_id(request_id.as_str());
let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers);
let mut report_context = attach_provider_response_headers_to_report_context(
report_context,
&headers,
response_observation.request_started_at_unix_ms,
response_observation.response_headers_observed_at_unix_ms,
&response_observation.request_order_id,
);
spawn_local_oauth_success_effect(
state.clone(),
&plan,
report_context.as_ref(),
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),
},
);
if status_code == 200 {
seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
@@ -3819,6 +3839,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -3891,6 +3912,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -3963,6 +3985,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -4035,6 +4058,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
@@ -4192,6 +4216,15 @@ async fn execute_execution_runtime_stream_inner(
record_stream_pending_lifecycle(state, seed, &mut stage_trace).await;
lifecycle_pending_recorded = true;
}
let report_context = attach_provider_response_headers_to_report_context(
report_context,
&execution.headers,
execution.response_observation.request_started_at_unix_ms,
execution
.response_observation
.response_headers_observed_at_unix_ms,
&execution.response_observation.request_order_id,
);
let stream_precommit_committed = execution.stream_precommit_committed;
let frame_stream = build_direct_execution_frame_stream(execution).boxed();
return execute_stream_from_frame_stream_with_retry_scope(
@@ -4211,6 +4244,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out,
retry_fallback_out,
None,
)
.await;
}
@@ -4327,6 +4361,15 @@ async fn execute_execution_runtime_stream_inner(
record_stream_pending_lifecycle(state, seed, &mut stage_trace).await;
lifecycle_pending_recorded = true;
}
let report_context = attach_provider_response_headers_to_report_context(
report_context,
&execution.headers,
execution.response_observation.request_started_at_unix_ms,
execution
.response_observation
.response_headers_observed_at_unix_ms,
&execution.response_observation.request_order_id,
);
let stream_precommit_committed = execution.stream_precommit_committed;
let frame_stream = build_direct_execution_frame_stream(execution).boxed();
return execute_stream_from_frame_stream_with_retry_scope(
@@ -4346,10 +4389,13 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
None,
)
.await;
}
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_stream_plan_to_remote_execution_runtime(
state,
remote_execution_runtime_base_url,
@@ -4431,6 +4477,12 @@ async fn execute_execution_runtime_stream_inner(
)?));
}
let remote_response_observed_at_unix_ms = current_request_candidate_unix_ms();
let remote_fallback_observation = 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,
};
let frame_stream = response
.bytes_stream()
.map_err(|err| IoError::other(err.to_string()))
@@ -4452,6 +4504,7 @@ async fn execute_execution_runtime_stream_inner(
provider_pool_in_flight_guard.take(),
retry_scope_out.as_deref_mut(),
retry_fallback_out.as_deref_mut(),
Some(remote_fallback_observation),
)
.await;
}
@@ -5481,6 +5534,7 @@ async fn execute_stream_from_frame_stream(
in_flight_guard,
None,
None,
None,
)
.await
}
@@ -5503,6 +5557,7 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
in_flight_guard: Option<ProviderPoolInFlightGuard>,
mut retry_scope_out: Option<&mut AiAttemptRetryScope>,
mut retry_fallback_out: Option<&mut Option<Response<Body>>>,
fallback_response_observation: Option<ExecutionResponseObservation>,
) -> Result<Option<Response<Body>>, GatewayError> {
let request_id = plan.request_id.as_str();
let request_id_for_log = short_request_id(request_id);
@@ -5535,14 +5590,37 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let StreamFramePayload::Headers {
status_code,
mut headers,
response_observation,
} = first_frame.payload
else {
return Err(GatewayError::Internal(
"execution runtime stream must start with headers frame".to_string(),
));
};
let mut report_context =
attach_provider_response_headers_to_report_context(report_context, &headers);
let response_observation = response_observation
.or(fallback_response_observation)
.unwrap_or(ExecutionResponseObservation {
request_started_at_unix_ms: candidate_started_unix_secs,
response_headers_observed_at_unix_ms: current_request_candidate_unix_ms(),
request_order_id: uuid::Uuid::now_v7().to_string(),
});
let mut report_context = attach_provider_response_headers_to_report_context(
report_context,
&headers,
response_observation.request_started_at_unix_ms,
response_observation.response_headers_observed_at_unix_ms,
&response_observation.request_order_id,
);
spawn_local_oauth_success_effect(
state.clone(),
&plan,
report_context.as_ref(),
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),
},
);
if status_code == 200 {
seed_kiro_simulated_cache_enabled(state, &plan, &mut report_context).await;
if kiro_simulated_cache_enabled_from_report_context(report_context.as_ref()) {
@@ -8310,6 +8388,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -8389,6 +8468,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -8431,6 +8511,7 @@ mod tests {
None,
Some(&mut retry_scope),
None,
None,
)
.await
.expect("prefetch transport execution should resolve");
@@ -8480,6 +8561,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -8522,6 +8604,7 @@ mod tests {
None,
Some(&mut retry_scope),
None,
None,
)
.await
.expect("prefetch HTTP status execution should resolve");
@@ -8680,6 +8763,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
for chunk in chunks {
@@ -8725,6 +8809,7 @@ mod tests {
None,
Some(&mut retry_scope),
Some(&mut fallback_response),
None,
)
.await
.expect("native Anthropic stream execution should succeed");
@@ -9364,6 +9449,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -9413,6 +9499,7 @@ mod tests {
None,
None,
None,
None,
),
)
.await
@@ -9850,6 +9937,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
}
@@ -11532,6 +11620,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -11660,6 +11749,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
@@ -12384,6 +12474,7 @@ mod tests {
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
response_observation: None,
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
apps/aether-gateway/src/execution_runtime/stream_pump.rs
+71 -41
View File
@@ -4,8 +4,9 @@ use std::io::Error as IoError;
use std::time::{Duration, Instant};
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionStreamTerminalSummary,
ExecutionTelemetry, StreamFrame, StreamFramePayload, StreamFrameType,
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionResponseObservation,
ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame, StreamFramePayload,
StreamFrameType,
};
use async_stream::stream;
use axum::body::Bytes;
@@ -44,6 +45,7 @@ pub(crate) fn build_direct_execution_frame_stream(
stream_precommit_committed: _,
response,
started_at,
response_observation,
stream_first_byte_timeout,
upstream_target_permit,
} = execution;
@@ -108,7 +110,11 @@ pub(crate) fn build_direct_execution_frame_stream(
}
}
match encode_headers_frame(status_code, response_headers) {
match encode_headers_frame(
status_code,
response_headers,
&response_observation,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -153,7 +159,11 @@ pub(crate) fn build_direct_execution_frame_stream(
upstream_bytes,
first_byte_timeout,
}) => {
match encode_headers_frame(status_code, original_headers) {
match encode_headers_frame(
status_code,
original_headers,
&response_observation,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -192,7 +202,11 @@ pub(crate) fn build_direct_execution_frame_stream(
return;
}
match encode_headers_frame(status_code, headers) {
match encode_headers_frame(
status_code,
headers,
&response_observation,
) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -611,12 +625,14 @@ pub(crate) fn build_direct_execution_frame_stream(
fn encode_headers_frame(
status_code: u16,
headers: BTreeMap<String, String>,
response_observation: &ExecutionResponseObservation,
) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code,
headers,
response_observation: Some(response_observation.clone()),
},
})
}
@@ -1606,43 +1622,47 @@ mod tests {
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let server = tokio::spawn(async move {
let app = Router::new().route(
"/responses",
post(|| async {
let body = serde_json::json!({
"id": "resp_sync_bridge_123",
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [{
"type": "message",
"id": "msg_sync_bridge_123",
"role": "assistant",
"content": [{
"type": "output_text",
"text": "Hello from buffered JSON stream",
"annotations": []
}]
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
});
let mut response = axum::http::Response::new(Body::from(
serde_json::to_vec(&body).expect("json should encode"),
));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
response
}),
);
axum::serve(listener, app)
let (mut socket, _) = listener.accept().await.expect("client should connect");
let mut request = [0_u8; 4096];
let _ = socket
.read(&mut request)
.await
.expect("server should start");
.expect("request should read");
let body = serde_json::to_vec(&serde_json::json!({
"id": "resp_sync_bridge_123",
"object": "response",
"model": "gpt-5.4",
"status": "completed",
"output": [{
"type": "message",
"id": "msg_sync_bridge_123",
"role": "assistant",
"content": [{
"type": "output_text",
"text": "Hello from buffered JSON stream",
"annotations": []
}]
}],
"usage": {
"input_tokens": 1,
"output_tokens": 2,
"total_tokens": 3
}
}))
.expect("json should encode");
socket
.write_all(
format!(
"HTTP/1.1 200 OK\r\ncontent-type: application/json\r\ncontent-length: {}\r\n\r\n",
body.len()
)
.as_bytes(),
)
.await
.expect("headers should write");
socket.flush().await.expect("headers should flush");
tokio::time::sleep(Duration::from_millis(75)).await;
socket.write_all(&body).await.expect("body should write");
});
let runtime = DirectSyncExecutionRuntime::new();
@@ -1678,6 +1698,12 @@ mod tests {
})
.await
.expect("stream execution should succeed");
let expected_observation = execution.response_observation.clone();
assert!(
expected_observation.response_headers_observed_at_unix_ms
>= expected_observation.request_started_at_unix_ms
);
assert!(!expected_observation.request_order_id.is_empty());
let frames = build_direct_execution_frame_stream(execution)
.map(|item| item.expect("frame should encode"))
@@ -1691,6 +1717,10 @@ mod tests {
let header_frame: Value =
serde_json::from_str(&frames[0]).expect("headers frame should parse");
let encoded_observation: aether_contracts::ExecutionResponseObservation =
serde_json::from_value(header_frame["payload"]["response_observation"].clone())
.expect("headers frame should retain the response observation");
assert_eq!(encoded_observation, expected_observation);
assert_eq!(
header_frame
.get("payload")
apps/aether-gateway/src/execution_runtime/sync/execution.rs
+97 -14
View File
@@ -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)]
apps/aether-gateway/src/execution_runtime/transport.rs
+103 -9
View File
@@ -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]
apps/aether-gateway/src/execution_runtime/windsurf.rs
+2
View File
@@ -266,6 +266,7 @@ pub(crate) async fn maybe_execute_windsurf_sync(
candidate_id: prepared.candidate_id,
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
@@ -527,6 +528,7 @@ fn build_windsurf_stream_frame_stream(
("cache-control".to_string(), "no-cache".to_string()),
("content-type".to_string(), "text/event-stream".to_string()),
]),
response_observation: None,
},
});