Unify quota snapshots and oauth refresh handling

crates/aether-admin/src/provider/pool.rs

@@ -121,10 +121,72 @@ fn admin_pool_json_f64(value: Option<&Value>) -> Option<f64> {
     .filter(|value| value.is_finite())
 }
 
+fn admin_pool_quota_snapshot_matches_provider(
+    quota_snapshot: &serde_json::Map<String, Value>,
+    provider_type: &str,
+) -> bool {
+    let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
+    match quota_snapshot
+        .get("provider_type")
+        .and_then(Value::as_str)
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+    {
+        Some(provider_type) => provider_type.eq_ignore_ascii_case(&normalized_provider_type),
+        None => {
+            admin_pool_json_bool(quota_snapshot.get("exhausted")) == Some(true)
+                || quota_snapshot
+                    .get("code")
+                    .and_then(Value::as_str)
+                    .is_some_and(|code| !code.trim().eq_ignore_ascii_case("unknown"))
+                || quota_snapshot
+                    .get("updated_at")
+                    .is_some_and(|value| !value.is_null())
+                || quota_snapshot
+                    .get("observed_at")
+                    .is_some_and(|value| !value.is_null())
+                || quota_snapshot
+                    .get("usage_ratio")
+                    .is_some_and(|value| !value.is_null())
+                || quota_snapshot
+                    .get("reset_seconds")
+                    .is_some_and(|value| !value.is_null())
+                || quota_snapshot
+                    .get("windows")
+                    .and_then(Value::as_array)
+                    .is_some_and(|windows| !windows.is_empty())
+                || quota_snapshot
+                    .get("credits")
+                    .and_then(Value::as_object)
+                    .is_some_and(|credits| !credits.is_empty())
+        }
+    }
+}
+
+fn admin_pool_key_quota_snapshot<'a>(
+    key: &'a StoredProviderCatalogKey,
+    provider_type: &str,
+) -> Option<&'a serde_json::Map<String, Value>> {
+    let quota_snapshot = key
+        .status_snapshot
+        .as_ref()
+        .and_then(Value::as_object)
+        .and_then(|snapshot| snapshot.get("quota"))
+        .and_then(Value::as_object)?;
+    admin_pool_quota_snapshot_matches_provider(quota_snapshot, provider_type)
+        .then_some(quota_snapshot)
+}
+
 pub fn admin_pool_key_account_quota_exhausted(
     key: &StoredProviderCatalogKey,
     provider_type: &str,
 ) -> bool {
+    if let Some(exhausted) = admin_pool_key_quota_snapshot(key, provider_type)
+        .and_then(|quota_snapshot| admin_pool_json_bool(quota_snapshot.get("exhausted")))
+    {
+        return exhausted;
+    }
+
     let provider_type = provider_type.trim().to_ascii_lowercase();
     let Some(bucket) = admin_pool_metadata_bucket(key.upstream_metadata.as_ref(), &provider_type)
     else {
@@ -663,6 +725,39 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn prefers_quota_snapshot_over_metadata_for_codex_exhaustion() {
+        let mut key = sample_key(Some(json!({
+            "codex": {
+                "secondary_used_percent": 100.0
+            }
+        })));
+        key.status_snapshot = Some(json!({
+            "quota": {
+                "version": 2,
+                "provider_type": "codex",
+                "code": "ok",
+                "exhausted": false,
+                "usage_ratio": 0.0,
+                "updated_at": 1_776_395_200u64,
+                "windows": [
+                    {
+                        "code": "weekly",
+                        "used_ratio": 0.0,
+                        "remaining_ratio": 1.0
+                    },
+                    {
+                        "code": "5h",
+                        "used_ratio": 0.0,
+                        "remaining_ratio": 1.0
+                    }
+                ]
+            }
+        }));
+
+        assert!(!admin_pool_key_account_quota_exhausted(&key, "codex"));
+    }
+
     #[test]
     fn detects_kiro_exhaustion_from_metadata() {
         assert!(admin_pool_key_account_quota_exhausted(

crates/aether-admin/src/provider/quota.rs

@@ -145,14 +145,11 @@ pub fn parse_antigravity_usage_response(
         let remaining_fraction = quota_info
             .and_then(|object| object.get("remainingFraction"))
             .and_then(coerce_json_f64);
-        let used_percent = remaining_fraction
-            .map(|value| ((1.0 - value).max(0.0) * 100.0).min(100.0))
-            .unwrap_or(100.0);
-        payload.insert(
-            "remaining_fraction".to_string(),
-            json!(remaining_fraction.unwrap_or(0.0)),
-        );
-        payload.insert("used_percent".to_string(), json!(used_percent));
+        if let Some(remaining_fraction) = remaining_fraction {
+            let used_percent = ((1.0 - remaining_fraction).max(0.0) * 100.0).min(100.0);
+            payload.insert("remaining_fraction".to_string(), json!(remaining_fraction));
+            payload.insert("used_percent".to_string(), json!(used_percent));
+        }
         if let Some(reset_time) = quota_info
             .and_then(|object| object.get("resetTime"))
             .cloned()

crates/aether-data/src/migrate.rs

@@ -417,7 +417,16 @@ mod tests {
                 return Ok(None);
             }
 
-            Ok(Some(Self::start(initdb_bin, postgres_bin).await?))
+            match Self::start(initdb_bin, postgres_bin).await {
+                Ok(server) => Ok(Some(server)),
+                Err(err) if postgres_shared_memory_unavailable(err.to_string().as_str()) => {
+                    eprintln!(
+                        "skipping postgres integration test because local postgres could not allocate shared memory: {err}"
+                    );
+                    Ok(None)
+                }
+                Err(err) => Err(err),
+            }
         }
 
         async fn start(
@@ -468,6 +477,14 @@ mod tests {
                 .arg("synchronous_commit=off")
                 .arg("-c")
                 .arg("full_page_writes=off")
+                .arg("-c")
+                .arg("shared_buffers=8MB")
+                .arg("-c")
+                .arg("max_connections=8")
+                .arg("-c")
+                .arg("dynamic_shared_memory_type=none")
+                .arg("-c")
+                .arg("autovacuum=off")
                 .stdout(Stdio::from(stdout))
                 .stderr(Stdio::from(stderr))
                 .spawn()?;
@@ -519,6 +536,14 @@ mod tests {
         Ok(port)
     }
 
+    fn postgres_shared_memory_unavailable(message: &str) -> bool {
+        let message = message.to_ascii_lowercase();
+        message.contains("shared memory")
+            && (message.contains("could not create shared memory segment")
+                || message.contains("shmget")
+                || message.contains("no space left on device"))
+    }
+
     async fn wait_for_postgres(database_url: &str) -> Result<(), Box<dyn std::error::Error>> {
         let deadline = Instant::now() + Duration::from_secs(10);
         loop {