Prefer codex rate limit bucket over spark

Author: yieldthought

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "codexapi"
-version = "0.7.2"
+version = "0.7.3"
 description = "Minimal Python API for running the Codex CLI."
 readme = "README.md"
 requires-python = ">=3.8"

src/codexapi/__init__.py

@@ -27,4 +27,4 @@
     "task_result",
     "lead",
 ]
-__version__ = "0.7.2"
+__version__ = "0.7.3"

src/codexapi/rate_limits.py

@@ -36,6 +36,7 @@ def rate_limits():
 
 def _extract_rate_limits(path):
     last = None
+    preferred = None
     try:
         with open(path, "r", encoding="utf-8", errors="replace") as handle:
             for line in handle:
@@ -49,9 +50,23 @@ def _extract_rate_limits(path):
                 rate_data = payload.get("rate_limits")
                 if isinstance(rate_data, dict):
                     last = rate_data
+                    if _is_primary_limit(rate_data):
+                        preferred = rate_data
     except OSError:
         return None
-    return last
+    return preferred or last
+
+
+def _is_primary_limit(rate_data):
+    limit_id = rate_data.get("limit_id")
+    if isinstance(limit_id, str):
+        return limit_id == "codex"
+    limit_name = rate_data.get("limit_name")
+    if limit_name is None:
+        return True
+    if isinstance(limit_name, str):
+        return not limit_name.strip()
+    return False
 
 
 def quota_line():

tests/test_rate_limits.py

@@ -0,0 +1,71 @@
+import json
+import os
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
+
+from codexapi.rate_limits import _extract_rate_limits
+
+
+def _write_events(path, rates):
+    with open(path, "w", encoding="utf-8") as handle:
+        for rate in rates:
+            event = {"payload": {"rate_limits": rate}}
+            handle.write(json.dumps(event))
+            handle.write("\n")
+
+
+class RateLimitsTests(unittest.TestCase):
+    def test_extract_prefers_codex_limit_id_when_spark_is_last(self):
+        with tempfile.NamedTemporaryFile(delete=False) as handle:
+            path = handle.name
+        try:
+            _write_events(
+                path,
+                [
+                    {"limit_id": "codex", "primary": {"used_percent": 3}},
+                    {
+                        "limit_id": "codex_bengalfox",
+                        "limit_name": "GPT-5.3-Codex-Spark",
+                        "primary": {"used_percent": 0},
+                    },
+                ],
+            )
+            found = _extract_rate_limits(path)
+        finally:
+            os.unlink(path)
+        self.assertEqual(found["limit_id"], "codex")
+        self.assertEqual(found["primary"]["used_percent"], 3)
+
+    def test_extract_falls_back_to_last_when_codex_missing(self):
+        with tempfile.NamedTemporaryFile(delete=False) as handle:
+            path = handle.name
+        try:
+            _write_events(
+                path,
+                [
+                    {"limit_id": "codex_bengalfox", "primary": {"used_percent": 1}},
+                    {"limit_id": "codex_lynx", "primary": {"used_percent": 2}},
+                ],
+            )
+            found = _extract_rate_limits(path)
+        finally:
+            os.unlink(path)
+        self.assertEqual(found["limit_id"], "codex_lynx")
+
+    def test_extract_keeps_legacy_single_limit_without_limit_id(self):
+        with tempfile.NamedTemporaryFile(delete=False) as handle:
+            path = handle.name
+        try:
+            _write_events(path, [{"primary": {"used_percent": 22}}])
+            found = _extract_rate_limits(path)
+        finally:
+            os.unlink(path)
+        self.assertEqual(found["primary"]["used_percent"], 22)
+
+
+if __name__ == "__main__":
+    unittest.main()