fix: read current Codex quota from desktop logs

Author: ericweichun

codex_loader.py

@@ -94,6 +94,9 @@ def _session_total_tokens(entries: list[UsageEntry]) -> int:
 
 
 def load_rate_limits() -> CodexRateLimits | None:
+    sqlite_limits = _load_sqlite_rate_limits()
+    if sqlite_limits is not None:
+        return sqlite_limits
     if not SESSIONS_DIR.is_dir():
         return None
     models = _load_thread_models()
@@ -105,6 +108,132 @@ def load_rate_limits() -> CodexRateLimits | None:
     return None
 
 
+def _load_sqlite_rate_limits() -> CodexRateLimits | None:
+    if not LOGS_DB.exists():
+        return None
+    query = (
+        "SELECT ts, feedback_log_body FROM logs "
+        "WHERE target = 'codex_api::endpoint::responses_websocket' "
+        "AND feedback_log_body LIKE '%websocket event:%' "
+        "AND (feedback_log_body LIKE '%\"type\":\"codex.rate_limits\"%' "
+        "OR feedback_log_body LIKE '%\"type\":\"error\"%usage_limit_reached%') "
+        "ORDER BY ts DESC, ts_nanos DESC, id DESC LIMIT 50"
+    )
+    try:
+        with sqlite3.connect(f"file:{LOGS_DB}?mode=ro", uri=True) as conn:
+            rows = conn.execute(query).fetchall()
+    except (OSError, sqlite3.Error):
+        if os.environ.get("USAGE_DEBUG") == "1":
+            logger.warning("codex sqlite rate limits load failed", exc_info=True)
+        return None
+
+    for ts, body in rows:
+        parsed = _parse_sqlite_rate_limits_row(ts, body)
+        if parsed is not None:
+            return parsed
+    return None
+
+
+def _parse_sqlite_rate_limits_row(ts: Any, body: Any) -> CodexRateLimits | None:
+    if not isinstance(body, str):
+        return None
+    event = _websocket_event_payload(body)
+    if not event:
+        return None
+    if event.get("type") == "codex.rate_limits":
+        return _rate_limits_from_websocket_event(event, body, ts)
+    if event.get("type") == "error":
+        return _rate_limits_from_websocket_error(event, body, ts)
+    return None
+
+
+def _websocket_event_payload(body: str) -> dict[str, Any]:
+    marker = "websocket event: "
+    index = body.find(marker)
+    if index < 0:
+        return {}
+    try:
+        data = json.loads(body[index + len(marker):])
+    except json.JSONDecodeError:
+        return {}
+    return data if isinstance(data, dict) else {}
+
+
+def _rate_limits_from_websocket_event(
+    event: dict[str, Any],
+    body: str,
+    ts: Any,
+) -> CodexRateLimits | None:
+    rate_limits = _as_dict(event.get("rate_limits"))
+    primary = _as_dict(rate_limits.get("primary"))
+    secondary = _as_dict(rate_limits.get("secondary"))
+    return _build_rate_limits(
+        primary_pct=_as_optional_float(primary.get("used_percent")),
+        primary_reset=_as_optional_float(primary.get("reset_at")),
+        secondary_pct=_as_optional_float(secondary.get("used_percent")),
+        secondary_reset=_as_optional_float(secondary.get("reset_at")),
+        model=_event_value(body, "model") or "unknown",
+        updated_at=_timestamp_from_log_ts(ts),
+    )
+
+
+def _rate_limits_from_websocket_error(
+    event: dict[str, Any],
+    body: str,
+    ts: Any,
+) -> CodexRateLimits | None:
+    headers = _as_dict(event.get("headers"))
+    primary_reset = _as_optional_float(headers.get("X-Codex-Primary-Reset-At"))
+    secondary_reset = _as_optional_float(headers.get("X-Codex-Secondary-Reset-At"))
+    now_ts = datetime.now(UTC).timestamp()
+    if primary_reset is None:
+        primary_reset_after = _as_optional_float(headers.get("X-Codex-Primary-Reset-After-Seconds"))
+        primary_reset = now_ts + primary_reset_after if primary_reset_after is not None else None
+    if secondary_reset is None:
+        secondary_reset_after = _as_optional_float(
+            headers.get("X-Codex-Secondary-Reset-After-Seconds")
+        )
+        secondary_reset = (
+            now_ts + secondary_reset_after if secondary_reset_after is not None else None
+        )
+    return _build_rate_limits(
+        primary_pct=_as_optional_float(headers.get("X-Codex-Primary-Used-Percent")),
+        primary_reset=primary_reset,
+        secondary_pct=_as_optional_float(headers.get("X-Codex-Secondary-Used-Percent")),
+        secondary_reset=secondary_reset,
+        model=_event_value(body, "model") or "unknown",
+        updated_at=_timestamp_from_log_ts(ts),
+    )
+
+
+def _build_rate_limits(
+    *,
+    primary_pct: float | None,
+    primary_reset: float | None,
+    secondary_pct: float | None,
+    secondary_reset: float | None,
+    model: str,
+    updated_at: datetime | None,
+) -> CodexRateLimits | None:
+    now_ts = datetime.now(UTC).timestamp()
+    if primary_reset is not None and primary_reset < now_ts:
+        primary_pct = None
+        primary_reset = None
+    if secondary_reset is not None and secondary_reset < now_ts:
+        secondary_pct = None
+        secondary_reset = None
+    if primary_pct is None and secondary_pct is None:
+        return None
+    return CodexRateLimits(
+        five_hour_pct=primary_pct,
+        five_hour_resets_at=primary_reset,
+        seven_day_pct=secondary_pct,
+        seven_day_resets_at=secondary_reset,
+        model=model,
+        updated_at=updated_at.isoformat() if updated_at is not None else "",
+    )
+
+
 def _load_thread_models() -> dict[str, str]:
     return {
         thread_id: metadata.model
@@ -225,7 +354,7 @@ def _parse_sqlite_log_row(
 def _event_value(body: str, key: str) -> str:
     pattern = _EVENT_VALUE_RE_CACHE.get(key)
     if pattern is None:
-        pattern = re.compile(rf'(?:^|\s){re.escape(key)}=(?:"([^"]*)"|([^\s]+))')
+        pattern = re.compile(rf'(?:^|[\s{{]){re.escape(key)}=(?:"([^"]*)"|([^\s}}]+))')
         _EVENT_VALUE_RE_CACHE[key] = pattern
     match = pattern.search(body)
     if match is None:

tests/test_codex_loader.py

@@ -228,7 +228,12 @@ def _create_state_db(path: Path, rows: list[tuple[str, str, str]]) -> None:
         conn.executemany("INSERT INTO threads (id, model, cwd) VALUES (?, ?, ?)", rows)
 
 
-def _create_logs_db(path: Path, rows: list[tuple[int, int, str]]) -> None:
+def _create_logs_db(
+    path: Path,
+    rows: list[tuple[int, int, str]],
+    *,
+    target: str = "codex_otel.trace_safe",
+) -> None:
     with sqlite3.connect(path) as conn:
         conn.execute(
             "CREATE TABLE logs ("
@@ -249,7 +254,7 @@ def _create_logs_db(path: Path, rows: list[tuple[int, int, str]]) -> None:
             conn.execute(
                 "INSERT INTO logs (id, ts, ts_nanos, target, feedback_log_body) "
                 "VALUES (?, ?, ?, ?, ?)",
-                (row_id, ts, row_id * 10, "codex_otel.trace_safe", body),
+                (row_id, ts, row_id * 10, target, body),
             )
 
 
@@ -466,6 +471,81 @@ def test_load_rate_limits_reads_primary_and_secondary_windows(
     )
 
 
+def test_load_rate_limits_prefers_sqlite_websocket_rate_limits(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    sessions_dir = tmp_path / "sessions"
+    logs_db = tmp_path / "logs.sqlite"
+    now = datetime(2026, 1, 1, 12, 0, tzinfo=UTC)
+    stale_limits = _rate_limits()
+    stale_limits["primary"]["used_percent"] = 9
+    _write_rate_limit_session(
+        sessions_dir / "rate.jsonl",
+        now.isoformat(),
+        stale_limits,
+        now.timestamp(),
+    )
+    body = (
+        "session_loop{thread_id=session-sqlite}:turn{model=gpt-5.5}: "
+        'websocket event: {"type":"codex.rate_limits","plan_type":"plus",'
+        '"rate_limits":{"allowed":true,"limit_reached":false,'
+        '"primary":{"used_percent":40,"window_minutes":300,"reset_at":9999999999},'
+        '"secondary":{"used_percent":6,"window_minutes":10080,"reset_at":9999999998}},'
+        '"code_review_rate_limits":null}'
+    )
+    _create_logs_db(
+        logs_db,
+        [(1, int(now.timestamp()), body)],
+        target="codex_api::endpoint::responses_websocket",
+    )
+    monkeypatch.setattr(codex_loader, "SESSIONS_DIR", sessions_dir)
+    monkeypatch.setattr(codex_loader, "LOGS_DB", logs_db)
+
+    result = codex_loader.load_rate_limits()
+
+    assert result == codex_loader.CodexRateLimits(
+        five_hour_pct=40.0,
+        five_hour_resets_at=9999999999.0,
+        seven_day_pct=6.0,
+        seven_day_resets_at=9999999998.0,
+        model="gpt-5.5",
+        updated_at=now.isoformat(),
+    )
+
+
+def test_load_rate_limits_reads_sqlite_usage_limit_error_headers(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    logs_db = tmp_path / "logs.sqlite"
+    now = datetime(2026, 1, 1, 12, 0, tzinfo=UTC)
+    body = (
+        "session_loop{thread_id=session-error}:turn{model=gpt-5.4}: "
+        'websocket event: {"type":"error","error":{"type":"usage_limit_reached"},'
+        '"headers":{"X-Codex-Primary-Used-Percent":"100",'
+        '"X-Codex-Secondary-Used-Percent":"47",'
+        '"X-Codex-Primary-Reset-At":"9999999999",'
+        '"X-Codex-Secondary-Reset-At":"9999999998"}}'
+    )
+    _create_logs_db(
+        logs_db,
+        [(1, int(now.timestamp()), body)],
+        target="codex_api::endpoint::responses_websocket",
+    )
+    monkeypatch.setattr(codex_loader, "SESSIONS_DIR", tmp_path / "missing-sessions")
+    monkeypatch.setattr(codex_loader, "LOGS_DB", logs_db)
+
+    result = codex_loader.load_rate_limits()
+
+    assert result == codex_loader.CodexRateLimits(
+        five_hour_pct=100.0,
+        five_hour_resets_at=9999999999.0,
+        seven_day_pct=47.0,
+        seven_day_resets_at=9999999998.0,
+        model="gpt-5.4",
+        updated_at=now.isoformat(),
+    )
+
+
 def test_load_rate_limits_clears_expired_primary_window(
     monkeypatch: pytest.MonkeyPatch, tmp_path: Path
 ) -> None: