gemini-cli-usage 0.1.6__tar.gz → 0.1.7__tar.gz

{gemini_cli_usage-0.1.6 → gemini_cli_usage-0.1.7}/.github/workflows/publish.yml

@@ -1,8 +1,10 @@
 name: Publish to PyPI
 
 on:
-  release:
-    types: [published]
+  workflow_dispatch:
+  push:
+    tags:
+      - "v*"
 
 permissions:
   id-token: write

{gemini_cli_usage-0.1.6 → gemini_cli_usage-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gemini-cli-usage
-Version: 0.1.6
+Version: 0.1.7
 Summary: Gemini CLI quota monitor — fetches Code Assist quota data from Google's backend
 Project-URL: Homepage, https://github.com/wakamex/gemini-cli-usage
 Project-URL: Source, https://github.com/wakamex/gemini-cli-usage

{gemini_cli_usage-0.1.6 → gemini_cli_usage-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "gemini-cli-usage"
-version = "0.1.6"
+version = "0.1.7"
 description = "Gemini CLI quota monitor — fetches Code Assist quota data from Google's backend"
 readme = "README.md"
 requires-python = ">=3.12"

{gemini_cli_usage-0.1.6 → gemini_cli_usage-0.1.7}/src/gemini_cli_usage/__init__.py

@@ -23,6 +23,7 @@ import re
 import signal
 import shutil
 import sys
+import tempfile
 import time
 import urllib.error
 import urllib.parse
@@ -181,12 +182,35 @@ def get_oauth_credentials() -> dict | None:
     return data if isinstance(data, dict) else None
 
 
-def _write_oauth_credentials(creds: dict):
+def _write_oauth_credentials(
+    creds: dict, expected_refresh_token: str | None = None
+) -> dict:
+    latest = get_oauth_credentials()
+    if expected_refresh_token and latest:
+        latest_refresh_token = latest.get("refresh_token")
+        if latest_refresh_token and latest_refresh_token != expected_refresh_token:
+            return latest
+        creds = {**latest, **creds}
+
+    tmp: Path | None = None
     try:
-        OAUTH_FILE.write_text(json.dumps(creds, indent=2) + "\n")
+        OAUTH_FILE.parent.mkdir(parents=True, exist_ok=True)
+        fd, tmp_name = tempfile.mkstemp(
+            prefix=f".{OAUTH_FILE.name}.", dir=OAUTH_FILE.parent
+        )
+        tmp = Path(tmp_name)
+        os.fchmod(fd, 0o600)
+        with os.fdopen(fd, "w") as file:
+            file.write(json.dumps(creds, indent=2) + "\n")
+        os.replace(tmp, OAUTH_FILE)
     except OSError:
         # Best effort only; the refreshed token can still be used in-memory.
-        pass
+        if tmp:
+            try:
+                tmp.unlink()
+            except OSError:
+                pass
+    return creds
 
 
 def _get_gemini_cli_oauth2_path() -> Path | None:
@@ -304,8 +328,16 @@ def refresh_access_token(creds: dict) -> dict:
         data=payload,
         headers={"Content-Type": "application/x-www-form-urlencoded"},
     )
-    with urllib.request.urlopen(req, timeout=10) as resp:
-        result = json.loads(resp.read())
+    try:
+        with urllib.request.urlopen(req, timeout=10) as resp:
+            result = json.loads(resp.read())
+    except urllib.error.HTTPError as exc:
+        latest = get_oauth_credentials()
+        if exc.code in (400, 401) and latest:
+            latest_refresh_token = latest.get("refresh_token")
+            if latest_refresh_token and latest_refresh_token != refresh_token:
+                return latest
+        raise
 
     updated = dict(creds)
     updated["access_token"] = result["access_token"]
@@ -316,17 +348,27 @@ def refresh_access_token(creds: dict) -> dict:
         updated["id_token"] = result["id_token"]
     if result.get("refresh_token"):
         updated["refresh_token"] = result["refresh_token"]
-    _write_oauth_credentials(updated)
-    return updated
+    return _write_oauth_credentials(updated, expected_refresh_token=refresh_token)
 
 
-def get_access_token() -> str:
+def get_access_token(
+    force_refresh: bool = False, rejected_token: str | None = None
+) -> str:
     creds = get_oauth_credentials()
     if not creds:
         raise RuntimeError("No OAuth credentials at ~/.gemini/oauth_creds.json")
 
+    token = creds.get("access_token")
+    token_was_replaced = (
+        rejected_token is not None
+        and isinstance(token, str)
+        and bool(token)
+        and token != rejected_token
+    )
     expiry_date = int(creds.get("expiry_date", 0) or 0)
-    if time.time() * 1000 >= expiry_date - 60_000:
+    if not token_was_replaced and (
+        force_refresh or time.time() * 1000 >= expiry_date - 60_000
+    ):
         creds = refresh_access_token(creds)
 
     token = creds.get("access_token")
@@ -426,32 +468,41 @@ def fetch_quota(project_root: Path | None = None) -> dict:
         )
 
     access_token = get_access_token()
-    load_res = _load_code_assist(access_token)
-
-    env_project = (
-        os.environ.get("GOOGLE_CLOUD_PROJECT")
-        or os.environ.get("GOOGLE_CLOUD_PROJECT_ID")
-        or None
-    )
-    project_id = load_res.get("cloudaicompanionProject") or env_project
-    if not project_id:
-        raise RuntimeError(
-            "No Code Assist project ID available. Set GOOGLE_CLOUD_PROJECT if your account requires it."
-        )
-
-    quota_res = _code_assist_post(
-        "retrieveUserQuota", {"project": project_id}, access_token
-    )
-    current_tier = load_res.get("currentTier") or {}
-    paid_tier = load_res.get("paidTier") or {}
-    result = {
-        "project_id": project_id,
-        "user_tier": paid_tier.get("id") or current_tier.get("id"),
-        "user_tier_name": paid_tier.get("name") or current_tier.get("name"),
-        "buckets": _parse_quota_buckets(quota_res.get("buckets") or []),
-    }
-    result["summary_bucket"] = _select_summary_bucket(result)
-    return result
+    for attempt in range(2):
+        try:
+            load_res = _load_code_assist(access_token)
+            env_project = (
+                os.environ.get("GOOGLE_CLOUD_PROJECT")
+                or os.environ.get("GOOGLE_CLOUD_PROJECT_ID")
+                or None
+            )
+            project_id = load_res.get("cloudaicompanionProject") or env_project
+            if not project_id:
+                raise RuntimeError(
+                    "No Code Assist project ID available. Set GOOGLE_CLOUD_PROJECT if your account requires it."
+                )
+
+            quota_res = _code_assist_post(
+                "retrieveUserQuota", {"project": project_id}, access_token
+            )
+            current_tier = load_res.get("currentTier") or {}
+            paid_tier = load_res.get("paidTier") or {}
+            result = {
+                "project_id": project_id,
+                "user_tier": paid_tier.get("id") or current_tier.get("id"),
+                "user_tier_name": paid_tier.get("name") or current_tier.get("name"),
+                "buckets": _parse_quota_buckets(quota_res.get("buckets") or []),
+            }
+            result["summary_bucket"] = _select_summary_bucket(result)
+            return result
+        except urllib.error.HTTPError as exc:
+            if exc.code in (401, 403) and attempt == 0:
+                access_token = get_access_token(
+                    force_refresh=True, rejected_token=access_token
+                )
+                continue
+            raise
+    raise RuntimeError("Failed to fetch quota after token refresh")
 
 
 def build_usage_json(project_root: Path | None = None) -> dict:

{gemini_cli_usage-0.1.6 → gemini_cli_usage-0.1.7}/tests/test_gemini_cli_usage.py

@@ -4,6 +4,7 @@ import json
 import os
 import tempfile
 import unittest
+import urllib.error
 from datetime import UTC, datetime, timedelta
 from pathlib import Path
 from unittest import mock
@@ -266,6 +267,131 @@ class GeminiUsageTests(unittest.TestCase):
         self.assertIn("GEMINI_OAUTH_CLIENT_ID", str(exc.exception))
         self.assertIn("run `gemini`", str(exc.exception))
 
+    def test_write_oauth_credentials_is_atomic_and_private(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            oauth_file = Path(tmp) / "oauth_creds.json"
+            with mock.patch.object(gemini_cli_usage, "OAUTH_FILE", oauth_file):
+                written = gemini_cli_usage._write_oauth_credentials(
+                    {"access_token": "token", "refresh_token": "refresh"}
+                )
+
+            self.assertEqual(json.loads(oauth_file.read_text()), written)
+            self.assertEqual(oauth_file.stat().st_mode & 0o777, 0o600)
+            self.assertEqual(list(oauth_file.parent.glob(".oauth_creds.json.*")), [])
+
+    def test_refresh_does_not_overwrite_concurrently_rotated_credentials(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            oauth_file = Path(tmp) / "oauth_creds.json"
+            original = {
+                "access_token": "old-access",
+                "refresh_token": "old-refresh",
+                "expiry_date": 0,
+            }
+            rotated = {
+                "access_token": "other-access",
+                "refresh_token": "other-refresh",
+                "expiry_date": 9_999_999_999_999,
+            }
+            _write_json(oauth_file, original)
+
+            class FakeResponse:
+                def __enter__(self):
+                    return self
+
+                def __exit__(self, *_args):
+                    return False
+
+                def read(self):
+                    _write_json(oauth_file, rotated)
+                    return json.dumps(
+                        {"access_token": "our-access", "refresh_token": "our-refresh"}
+                    ).encode()
+
+            with (
+                mock.patch.object(gemini_cli_usage, "OAUTH_FILE", oauth_file),
+                mock.patch.object(
+                    gemini_cli_usage,
+                    "_get_oauth_client_credentials",
+                    return_value=("client", "secret"),
+                ),
+                mock.patch.object(
+                    gemini_cli_usage.urllib.request,
+                    "urlopen",
+                    return_value=FakeResponse(),
+                ),
+            ):
+                result = gemini_cli_usage.refresh_access_token(original)
+
+            self.assertEqual(result, rotated)
+            self.assertEqual(json.loads(oauth_file.read_text()), rotated)
+
+    def test_rejected_token_uses_concurrently_replaced_access_token(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            oauth_file = Path(tmp) / "oauth_creds.json"
+            _write_json(
+                oauth_file,
+                {
+                    "access_token": "replacement",
+                    "refresh_token": "refresh",
+                    "expiry_date": 9_999_999_999_999,
+                },
+            )
+            with (
+                mock.patch.object(gemini_cli_usage, "OAUTH_FILE", oauth_file),
+                mock.patch.object(
+                    gemini_cli_usage, "refresh_access_token"
+                ) as refresh_mock,
+            ):
+                token = gemini_cli_usage.get_access_token(
+                    force_refresh=True, rejected_token="rejected"
+                )
+
+            self.assertEqual(token, "replacement")
+            refresh_mock.assert_not_called()
+
+    def test_fetch_quota_recovers_once_from_auth_error(self):
+        auth_error = urllib.error.HTTPError(
+            url="https://cloudcode-pa.googleapis.com",
+            code=401,
+            msg="Unauthorized",
+            hdrs={},
+            fp=None,
+        )
+        with (
+            mock.patch.object(
+                gemini_cli_usage,
+                "get_auth_type",
+                return_value="oauth-personal",
+            ),
+            mock.patch.object(
+                gemini_cli_usage,
+                "get_access_token",
+                side_effect=["stale", "fresh"],
+            ) as token_mock,
+            mock.patch.object(
+                gemini_cli_usage,
+                "_load_code_assist",
+                side_effect=[
+                    auth_error,
+                    {"cloudaicompanionProject": "project"},
+                ],
+            ),
+            mock.patch.object(
+                gemini_cli_usage,
+                "_code_assist_post",
+                return_value={"buckets": []},
+            ),
+        ):
+            quota = gemini_cli_usage.fetch_quota()
+
+        self.assertEqual(quota["project_id"], "project")
+        token_mock.assert_has_calls(
+            [
+                mock.call(),
+                mock.call(force_refresh=True, rejected_token="stale"),
+            ]
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

{gemini_cli_usage-0.1.6 → gemini_cli_usage-0.1.7}/uv.lock

@@ -4,5 +4,5 @@ requires-python = ">=3.12"
 
 [[package]]
 name = "gemini-cli-usage"
-version = "0.1.4"
+version = "0.1.7"
 source = { editable = "." }