fix(codex): fence concurrent quota updates

2026-10-07 01:47:47 +08:00 · 2026-08-14 09:28:07 +08:00
parent f3a12c1008
commit 5b0c763086
65 changed files with 13009 additions and 738 deletions
@@ -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());
+    }
+}