Track scheduler affinity epochs and key sorting

apps/aether-gateway/src/orchestration/attempt.rs

@@ -31,8 +31,10 @@ pub(crate) struct LocalExecutionCandidateMetadata {
     pub(crate) candidate_group_id: Option<String>,
     pub(crate) pool_key_index: Option<u32>,
     pub(crate) pool_key_lease: Option<RuntimeLockLease>,
+    pub(crate) scheduler_affinity_epoch: Option<u64>,
 }
 
+pub(crate) const SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD: &str = "scheduler_affinity_epoch";
 pub(crate) const POOL_KEY_LEASE_KEY_REPORT_FIELD: &str = "pool_key_lease_key";
 pub(crate) const POOL_KEY_LEASE_OWNER_REPORT_FIELD: &str = "pool_key_lease_owner";
 pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token";
@@ -65,6 +67,9 @@ pub(crate) fn local_execution_candidate_metadata_from_report_context(
             .and_then(Value::as_u64)
             .and_then(|value| u32::try_from(value).ok()),
         pool_key_lease: pool_key_lease_from_report_context(report_context),
+        scheduler_affinity_epoch: report_context
+            .and_then(|value| value.get(SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD))
+            .and_then(Value::as_u64),
     }
 }
 
@@ -493,6 +498,7 @@ mod tests {
                     token: "gateway-1:token-1".to_string(),
                     ttl_ms: 900000,
                 }),
+                scheduler_affinity_epoch: None,
             }
         );
     }

apps/aether-gateway/src/orchestration/effects.rs

@@ -248,12 +248,16 @@ fn remember_successful_local_scheduler_affinity(
     let Some(target) = local_scheduler_affinity_target(context.plan) else {
         return;
     };
+    let expected_epoch =
+        local_execution_candidate_metadata_from_report_context(context.report_context)
+            .scheduler_affinity_epoch;
 
-    state.remember_scheduler_affinity_target(
+    let _ = state.remember_scheduler_affinity_target_for_epoch(
         &cache_key,
         target,
         SCHEDULER_AFFINITY_TTL,
         LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES,
+        expected_epoch,
     );
 }
 

apps/aether-gateway/src/orchestration/mod.rs

@@ -20,7 +20,7 @@ pub(crate) use self::attempt::{
     attempt_identity_from_report_context, build_local_attempt_identities,
     insert_pool_key_lease_report_context_fields, local_attempt_slot_count,
     local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
-    LocalExecutionCandidateMetadata,
+    LocalExecutionCandidateMetadata, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
 };
 pub(crate) use self::classifier::{
     classify_local_failover, local_failover_error_message, LocalFailoverClassification,

apps/aether-gateway/src/orchestration/report_effects.rs

@@ -58,6 +58,22 @@ fn report_context_key_id(report_context: Option<&Value>) -> Option<String> {
         .map(ToOwned::to_owned)
 }
 
+fn report_context_provider_response_headers(
+    report_context: Option<&Value>,
+) -> Option<BTreeMap<String, String>> {
+    let headers = report_context
+        .and_then(|context| context.get("provider_response_headers"))
+        .and_then(Value::as_object)?;
+    let mut out = BTreeMap::new();
+    for (key, value) in headers {
+        let Some(value) = value.as_str() else {
+            continue;
+        };
+        out.insert(key.clone(), value.to_string());
+    }
+    (!out.is_empty()).then_some(out)
+}
+
 fn is_volatile_compare_field(key: &str) -> bool {
     key == "updated_at" || key.ends_with("_reset_seconds") || key.ends_with("_reset_after_seconds")
 }
@@ -333,7 +349,12 @@ async fn sync_codex_quota_from_response_headers(
     };
 
     let now_unix_secs = current_unix_secs();
-    let Some(parsed) = admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
+    let provider_headers = report_context_provider_response_headers(report_context);
+    let parsed_from_provider_headers = provider_headers.as_ref().and_then(|headers| {
+        admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
+    });
+    let Some(parsed) = parsed_from_provider_headers
+        .or_else(|| admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs))
     else {
         return Ok(false);
     };