gemini-cli-usage 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- gemini_cli_usage-0.1.0/.github/workflows/publish.yml +18 -0
- gemini_cli_usage-0.1.0/.gitignore +2 -0
- gemini_cli_usage-0.1.0/PKG-INFO +116 -0
- gemini_cli_usage-0.1.0/README.md +109 -0
- gemini_cli_usage-0.1.0/pyproject.toml +20 -0
- gemini_cli_usage-0.1.0/src/gemini_cli_usage/__init__.py +624 -0
- gemini_cli_usage-0.1.0/tests/test_gemini_cli_usage.py +271 -0
- gemini_cli_usage-0.1.0/uv.lock +8 -0
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
name: Publish to PyPI
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
release:
|
|
5
|
+
types: [published]
|
|
6
|
+
|
|
7
|
+
permissions:
|
|
8
|
+
id-token: write
|
|
9
|
+
|
|
10
|
+
jobs:
|
|
11
|
+
publish:
|
|
12
|
+
runs-on: ubuntu-latest
|
|
13
|
+
environment: pypi
|
|
14
|
+
steps:
|
|
15
|
+
- uses: actions/checkout@v4
|
|
16
|
+
- uses: astral-sh/setup-uv@v4
|
|
17
|
+
- run: uv build
|
|
18
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: gemini-cli-usage
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Gemini CLI quota monitor — fetches Code Assist quota data from Google's backend
|
|
5
|
+
Requires-Python: >=3.12
|
|
6
|
+
Description-Content-Type: text/markdown
|
|
7
|
+
|
|
8
|
+
# gemini-cli-usage
|
|
9
|
+
|
|
10
|
+
Gemini CLI quota monitor. Fetches live Code Assist quota data from Google's
|
|
11
|
+
backend when Gemini is using Google login.
|
|
12
|
+
|
|
13
|
+
## Example output
|
|
14
|
+
|
|
15
|
+
`gemini-cli-usage` command:
|
|
16
|
+
|
|
17
|
+
```text
|
|
18
|
+
Project: gemini-cli-usage
|
|
19
|
+
Auth: Google login
|
|
20
|
+
gemini-2.5-pro 3.5% used resets 19h38m
|
|
21
|
+
gemini-2.5-flash-lite 0.07% used resets 19h37m
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
`gemini-cli-usage statusline` command:
|
|
25
|
+
|
|
26
|
+
```text
|
|
27
|
+
q:3.5% reset:19h38m
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Install
|
|
31
|
+
|
|
32
|
+
```bash
|
|
33
|
+
uv tool install gemini-cli-usage
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
For local development from a checkout:
|
|
37
|
+
|
|
38
|
+
```bash
|
|
39
|
+
uv tool install .
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
Then run:
|
|
43
|
+
|
|
44
|
+
```bash
|
|
45
|
+
# Check usage once
|
|
46
|
+
gemini-cli-usage
|
|
47
|
+
|
|
48
|
+
# Raw JSON
|
|
49
|
+
gemini-cli-usage json
|
|
50
|
+
|
|
51
|
+
# Compact shell/statusline output
|
|
52
|
+
gemini-cli-usage statusline
|
|
53
|
+
|
|
54
|
+
# Force a fresh cache rebuild and print full status
|
|
55
|
+
gemini-cli-usage refresh
|
|
56
|
+
|
|
57
|
+
# Keep ~/.gemini/usage-limits.json fresh
|
|
58
|
+
gemini-cli-usage daemon
|
|
59
|
+
```
|
|
60
|
+
|
|
61
|
+
## Commands
|
|
62
|
+
|
|
63
|
+
| Command | Description |
|
|
64
|
+
|---------|-------------|
|
|
65
|
+
| `gemini-cli-usage` | Show current usage (colored terminal output) |
|
|
66
|
+
| `gemini-cli-usage status` | Same as above |
|
|
67
|
+
| `gemini-cli-usage json` | Print raw JSON |
|
|
68
|
+
| `gemini-cli-usage daemon [-i SECS]` | Run in foreground, refresh every 5 min |
|
|
69
|
+
| `gemini-cli-usage statusline` | Compact statusline (reads cache, refreshes if stale) |
|
|
70
|
+
| `gemini-cli-usage refresh` | Force a fresh fetch, rewrite cache, and print status |
|
|
71
|
+
| `gemini-cli-usage install` | Print setup instructions |
|
|
72
|
+
|
|
73
|
+
## Data source
|
|
74
|
+
|
|
75
|
+
### `account_quota`
|
|
76
|
+
|
|
77
|
+
When Gemini CLI is configured for Google login (`oauth-personal`), it calls
|
|
78
|
+
Google's internal Code Assist API:
|
|
79
|
+
|
|
80
|
+
- `loadCodeAssist`
|
|
81
|
+
- `retrieveUserQuota`
|
|
82
|
+
|
|
83
|
+
This tool mirrors that flow using the OAuth credentials in
|
|
84
|
+
`~/.gemini/oauth_creds.json`.
|
|
85
|
+
|
|
86
|
+
## Notes
|
|
87
|
+
|
|
88
|
+
- Quota fetches are best-effort. If auth is not Google login, or quota lookup
|
|
89
|
+
fails, the tool reports the auth state plus the quota error.
|
|
90
|
+
- If the Google OAuth access token expires, the tool reuses Gemini CLI's
|
|
91
|
+
installed OAuth client metadata when available. If Gemini is installed in a
|
|
92
|
+
nonstandard location, set `GEMINI_OAUTH_CLIENT_ID` and
|
|
93
|
+
`GEMINI_OAUTH_CLIENT_SECRET`, or rerun `gemini` and retry.
|
|
94
|
+
- `status` and `json` always build fresh data.
|
|
95
|
+
- `statusline` reads the cache by default; use `--refresh` or `--max-age 0` to
|
|
96
|
+
force a live refresh.
|
|
97
|
+
- `refresh` is a convenience command that rebuilds the cache and prints the full
|
|
98
|
+
status output.
|
|
99
|
+
- Absolute quota counts are only shown when Google's response includes both
|
|
100
|
+
`remainingAmount` and a usable fraction. Otherwise the tool reports `% used`
|
|
101
|
+
plus reset time.
|
|
102
|
+
- Auth detection follows Gemini CLI precedence: environment variables first,
|
|
103
|
+
then workspace `.gemini/settings.json`, then global `~/.gemini/settings.json`.
|
|
104
|
+
|
|
105
|
+
## Options
|
|
106
|
+
|
|
107
|
+
```text
|
|
108
|
+
usage: gemini-cli-usage [-h] [--root ROOT] [--interval INTERVAL]
|
|
109
|
+
[--max-age MAX_AGE] [--refresh]
|
|
110
|
+
{status,json,daemon,statusline,refresh,install}
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
- `--root ROOT`: inspect a different project root instead of the current
|
|
114
|
+
directory
|
|
115
|
+
- `--max-age MAX_AGE`: cache TTL for `statusline`
|
|
116
|
+
- `--refresh`: ignore the cache and rebuild fresh data where applicable
|
|
@@ -0,0 +1,109 @@
|
|
|
1
|
+
# gemini-cli-usage
|
|
2
|
+
|
|
3
|
+
Gemini CLI quota monitor. Fetches live Code Assist quota data from Google's
|
|
4
|
+
backend when Gemini is using Google login.
|
|
5
|
+
|
|
6
|
+
## Example output
|
|
7
|
+
|
|
8
|
+
`gemini-cli-usage` command:
|
|
9
|
+
|
|
10
|
+
```text
|
|
11
|
+
Project: gemini-cli-usage
|
|
12
|
+
Auth: Google login
|
|
13
|
+
gemini-2.5-pro 3.5% used resets 19h38m
|
|
14
|
+
gemini-2.5-flash-lite 0.07% used resets 19h37m
|
|
15
|
+
```
|
|
16
|
+
|
|
17
|
+
`gemini-cli-usage statusline` command:
|
|
18
|
+
|
|
19
|
+
```text
|
|
20
|
+
q:3.5% reset:19h38m
|
|
21
|
+
```
|
|
22
|
+
|
|
23
|
+
## Install
|
|
24
|
+
|
|
25
|
+
```bash
|
|
26
|
+
uv tool install gemini-cli-usage
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
For local development from a checkout:
|
|
30
|
+
|
|
31
|
+
```bash
|
|
32
|
+
uv tool install .
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
Then run:
|
|
36
|
+
|
|
37
|
+
```bash
|
|
38
|
+
# Check usage once
|
|
39
|
+
gemini-cli-usage
|
|
40
|
+
|
|
41
|
+
# Raw JSON
|
|
42
|
+
gemini-cli-usage json
|
|
43
|
+
|
|
44
|
+
# Compact shell/statusline output
|
|
45
|
+
gemini-cli-usage statusline
|
|
46
|
+
|
|
47
|
+
# Force a fresh cache rebuild and print full status
|
|
48
|
+
gemini-cli-usage refresh
|
|
49
|
+
|
|
50
|
+
# Keep ~/.gemini/usage-limits.json fresh
|
|
51
|
+
gemini-cli-usage daemon
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
## Commands
|
|
55
|
+
|
|
56
|
+
| Command | Description |
|
|
57
|
+
|---------|-------------|
|
|
58
|
+
| `gemini-cli-usage` | Show current usage (colored terminal output) |
|
|
59
|
+
| `gemini-cli-usage status` | Same as above |
|
|
60
|
+
| `gemini-cli-usage json` | Print raw JSON |
|
|
61
|
+
| `gemini-cli-usage daemon [-i SECS]` | Run in foreground, refresh every 5 min |
|
|
62
|
+
| `gemini-cli-usage statusline` | Compact statusline (reads cache, refreshes if stale) |
|
|
63
|
+
| `gemini-cli-usage refresh` | Force a fresh fetch, rewrite cache, and print status |
|
|
64
|
+
| `gemini-cli-usage install` | Print setup instructions |
|
|
65
|
+
|
|
66
|
+
## Data source
|
|
67
|
+
|
|
68
|
+
### `account_quota`
|
|
69
|
+
|
|
70
|
+
When Gemini CLI is configured for Google login (`oauth-personal`), it calls
|
|
71
|
+
Google's internal Code Assist API:
|
|
72
|
+
|
|
73
|
+
- `loadCodeAssist`
|
|
74
|
+
- `retrieveUserQuota`
|
|
75
|
+
|
|
76
|
+
This tool mirrors that flow using the OAuth credentials in
|
|
77
|
+
`~/.gemini/oauth_creds.json`.
|
|
78
|
+
|
|
79
|
+
## Notes
|
|
80
|
+
|
|
81
|
+
- Quota fetches are best-effort. If auth is not Google login, or quota lookup
|
|
82
|
+
fails, the tool reports the auth state plus the quota error.
|
|
83
|
+
- If the Google OAuth access token expires, the tool reuses Gemini CLI's
|
|
84
|
+
installed OAuth client metadata when available. If Gemini is installed in a
|
|
85
|
+
nonstandard location, set `GEMINI_OAUTH_CLIENT_ID` and
|
|
86
|
+
`GEMINI_OAUTH_CLIENT_SECRET`, or rerun `gemini` and retry.
|
|
87
|
+
- `status` and `json` always build fresh data.
|
|
88
|
+
- `statusline` reads the cache by default; use `--refresh` or `--max-age 0` to
|
|
89
|
+
force a live refresh.
|
|
90
|
+
- `refresh` is a convenience command that rebuilds the cache and prints the full
|
|
91
|
+
status output.
|
|
92
|
+
- Absolute quota counts are only shown when Google's response includes both
|
|
93
|
+
`remainingAmount` and a usable fraction. Otherwise the tool reports `% used`
|
|
94
|
+
plus reset time.
|
|
95
|
+
- Auth detection follows Gemini CLI precedence: environment variables first,
|
|
96
|
+
then workspace `.gemini/settings.json`, then global `~/.gemini/settings.json`.
|
|
97
|
+
|
|
98
|
+
## Options
|
|
99
|
+
|
|
100
|
+
```text
|
|
101
|
+
usage: gemini-cli-usage [-h] [--root ROOT] [--interval INTERVAL]
|
|
102
|
+
[--max-age MAX_AGE] [--refresh]
|
|
103
|
+
{status,json,daemon,statusline,refresh,install}
|
|
104
|
+
```
|
|
105
|
+
|
|
106
|
+
- `--root ROOT`: inspect a different project root instead of the current
|
|
107
|
+
directory
|
|
108
|
+
- `--max-age MAX_AGE`: cache TTL for `statusline`
|
|
109
|
+
- `--refresh`: ignore the cache and rebuild fresh data where applicable
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "gemini-cli-usage"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Gemini CLI quota monitor — fetches Code Assist quota data from Google's backend"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.12"
|
|
7
|
+
dependencies = []
|
|
8
|
+
|
|
9
|
+
[tool.uv]
|
|
10
|
+
package = true
|
|
11
|
+
|
|
12
|
+
[build-system]
|
|
13
|
+
requires = ["hatchling"]
|
|
14
|
+
build-backend = "hatchling.build"
|
|
15
|
+
|
|
16
|
+
[tool.hatch.build.targets.wheel]
|
|
17
|
+
packages = ["src/gemini_cli_usage"]
|
|
18
|
+
|
|
19
|
+
[project.scripts]
|
|
20
|
+
gemini-cli-usage = "gemini_cli_usage:main"
|
|
@@ -0,0 +1,624 @@
|
|
|
1
|
+
#!/usr/bin/env python3
|
|
2
|
+
"""gemini-cli-usage - Gemini CLI quota monitor.
|
|
3
|
+
|
|
4
|
+
Fetches Gemini Code Assist quota data using the OAuth credentials
|
|
5
|
+
stored in ~/.gemini/oauth_creds.json.
|
|
6
|
+
|
|
7
|
+
Usage:
|
|
8
|
+
gemini-cli-usage
|
|
9
|
+
gemini-cli-usage status
|
|
10
|
+
gemini-cli-usage json
|
|
11
|
+
gemini-cli-usage daemon
|
|
12
|
+
gemini-cli-usage statusline
|
|
13
|
+
gemini-cli-usage refresh
|
|
14
|
+
gemini-cli-usage install
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import argparse
|
|
20
|
+
import json
|
|
21
|
+
import os
|
|
22
|
+
import re
|
|
23
|
+
import signal
|
|
24
|
+
import shutil
|
|
25
|
+
import sys
|
|
26
|
+
import time
|
|
27
|
+
import urllib.error
|
|
28
|
+
import urllib.parse
|
|
29
|
+
import urllib.request
|
|
30
|
+
from datetime import UTC, datetime
|
|
31
|
+
from pathlib import Path
|
|
32
|
+
|
|
33
|
+
GEMINI_DIR = Path.home() / ".gemini"
|
|
34
|
+
OAUTH_FILE = GEMINI_DIR / "oauth_creds.json"
|
|
35
|
+
SETTINGS_FILE = GEMINI_DIR / "settings.json"
|
|
36
|
+
DEFAULT_USAGE_FILE = GEMINI_DIR / "usage-limits.json"
|
|
37
|
+
|
|
38
|
+
DAEMON_INTERVAL = 300 # 5 minutes
|
|
39
|
+
CACHE_MAX_AGE = 300
|
|
40
|
+
|
|
41
|
+
CODE_ASSIST_BASE_URL = "https://cloudcode-pa.googleapis.com/v1internal"
|
|
42
|
+
TOKEN_URL = "https://oauth2.googleapis.com/token"
|
|
43
|
+
|
|
44
|
+
AUTH_LABELS = {
|
|
45
|
+
"oauth-personal": "Google login",
|
|
46
|
+
"gemini-api-key": "Gemini API key",
|
|
47
|
+
"vertex-ai": "Vertex AI",
|
|
48
|
+
"cloud-shell": "Cloud Shell",
|
|
49
|
+
"compute-default-credentials": "Compute ADC",
|
|
50
|
+
"gateway": "Gateway",
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
OAUTH_CLIENT_ID_PATTERN = re.compile(r"const OAUTH_CLIENT_ID = '([^']+)';")
|
|
54
|
+
OAUTH_CLIENT_SECRET_PATTERN = re.compile(r"const OAUTH_CLIENT_SECRET = '([^']+)';")
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _read_json(path: Path) -> dict | list | None:
|
|
58
|
+
try:
|
|
59
|
+
return json.loads(path.read_text())
|
|
60
|
+
except (FileNotFoundError, json.JSONDecodeError):
|
|
61
|
+
return None
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def _iso_now() -> str:
|
|
65
|
+
return datetime.now(UTC).isoformat()
|
|
66
|
+
|
|
67
|
+
|
|
68
|
+
def _parse_iso(timestamp: str | None) -> datetime | None:
|
|
69
|
+
if not timestamp:
|
|
70
|
+
return None
|
|
71
|
+
try:
|
|
72
|
+
return datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
|
73
|
+
except ValueError:
|
|
74
|
+
return None
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def _read_auth_type_from_settings(path: Path) -> str | None:
|
|
78
|
+
settings = _read_json(path)
|
|
79
|
+
if not isinstance(settings, dict):
|
|
80
|
+
return None
|
|
81
|
+
security = settings.get("security")
|
|
82
|
+
if not isinstance(security, dict):
|
|
83
|
+
return None
|
|
84
|
+
auth = security.get("auth")
|
|
85
|
+
if not isinstance(auth, dict):
|
|
86
|
+
return None
|
|
87
|
+
selected = auth.get("selectedType")
|
|
88
|
+
return selected if isinstance(selected, str) else None
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def _get_env_auth_type() -> str | None:
|
|
92
|
+
if os.environ.get("GOOGLE_GENAI_USE_GCA") == "true":
|
|
93
|
+
return "oauth-personal"
|
|
94
|
+
if os.environ.get("GOOGLE_GENAI_USE_VERTEXAI") == "true":
|
|
95
|
+
return "vertex-ai"
|
|
96
|
+
if os.environ.get("GEMINI_API_KEY"):
|
|
97
|
+
return "gemini-api-key"
|
|
98
|
+
if (
|
|
99
|
+
os.environ.get("CLOUD_SHELL") == "true"
|
|
100
|
+
or os.environ.get("GEMINI_CLI_USE_COMPUTE_ADC") == "true"
|
|
101
|
+
):
|
|
102
|
+
return "compute-default-credentials"
|
|
103
|
+
return None
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def get_auth_type(project_root: Path | None = None) -> str | None:
|
|
107
|
+
env_auth = _get_env_auth_type()
|
|
108
|
+
if env_auth:
|
|
109
|
+
return env_auth
|
|
110
|
+
|
|
111
|
+
if project_root:
|
|
112
|
+
workspace_auth = _read_auth_type_from_settings(
|
|
113
|
+
project_root.resolve() / ".gemini" / "settings.json"
|
|
114
|
+
)
|
|
115
|
+
if workspace_auth:
|
|
116
|
+
return workspace_auth
|
|
117
|
+
|
|
118
|
+
return _read_auth_type_from_settings(SETTINGS_FILE)
|
|
119
|
+
|
|
120
|
+
|
|
121
|
+
def get_auth_label(auth_type: str | None) -> str:
|
|
122
|
+
return AUTH_LABELS.get(auth_type or "", auth_type or "unknown")
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
def get_oauth_credentials() -> dict | None:
|
|
126
|
+
data = _read_json(OAUTH_FILE)
|
|
127
|
+
return data if isinstance(data, dict) else None
|
|
128
|
+
|
|
129
|
+
|
|
130
|
+
def _write_oauth_credentials(creds: dict):
|
|
131
|
+
try:
|
|
132
|
+
OAUTH_FILE.write_text(json.dumps(creds, indent=2) + "\n")
|
|
133
|
+
except OSError:
|
|
134
|
+
# Best effort only; the refreshed token can still be used in-memory.
|
|
135
|
+
pass
|
|
136
|
+
|
|
137
|
+
|
|
138
|
+
def _get_gemini_cli_oauth2_path() -> Path | None:
|
|
139
|
+
gemini_bin = shutil.which("gemini")
|
|
140
|
+
if not gemini_bin:
|
|
141
|
+
return None
|
|
142
|
+
|
|
143
|
+
resolved = Path(gemini_bin).resolve()
|
|
144
|
+
package_root = resolved.parent.parent
|
|
145
|
+
oauth2_path = package_root / "node_modules" / "@google" / "gemini-cli-core" / "dist" / "src" / "code_assist" / "oauth2.js"
|
|
146
|
+
return oauth2_path if oauth2_path.exists() else None
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
def _get_gemini_cli_oauth_client_credentials() -> tuple[str, str] | None:
|
|
150
|
+
oauth2_path = _get_gemini_cli_oauth2_path()
|
|
151
|
+
if not oauth2_path:
|
|
152
|
+
return None
|
|
153
|
+
|
|
154
|
+
try:
|
|
155
|
+
source = oauth2_path.read_text()
|
|
156
|
+
except OSError:
|
|
157
|
+
return None
|
|
158
|
+
|
|
159
|
+
client_id_match = OAUTH_CLIENT_ID_PATTERN.search(source)
|
|
160
|
+
client_secret_match = OAUTH_CLIENT_SECRET_PATTERN.search(source)
|
|
161
|
+
if not client_id_match or not client_secret_match:
|
|
162
|
+
return None
|
|
163
|
+
|
|
164
|
+
return client_id_match.group(1), client_secret_match.group(1)
|
|
165
|
+
|
|
166
|
+
|
|
167
|
+
def _get_oauth_client_credentials(creds: dict | None = None) -> tuple[str, str]:
|
|
168
|
+
client_id = os.environ.get("GEMINI_OAUTH_CLIENT_ID")
|
|
169
|
+
client_secret = os.environ.get("GEMINI_OAUTH_CLIENT_SECRET")
|
|
170
|
+
if client_id and client_secret:
|
|
171
|
+
return client_id, client_secret
|
|
172
|
+
|
|
173
|
+
creds = creds or get_oauth_credentials() or {}
|
|
174
|
+
client_id = creds.get("client_id")
|
|
175
|
+
client_secret = creds.get("client_secret")
|
|
176
|
+
if isinstance(client_id, str) and isinstance(client_secret, str):
|
|
177
|
+
if client_id and client_secret:
|
|
178
|
+
return client_id, client_secret
|
|
179
|
+
|
|
180
|
+
live_credentials = _get_gemini_cli_oauth_client_credentials()
|
|
181
|
+
if live_credentials:
|
|
182
|
+
return live_credentials
|
|
183
|
+
|
|
184
|
+
raise RuntimeError(
|
|
185
|
+
"OAuth access token expired and no Gemini CLI OAuth client metadata "
|
|
186
|
+
"was found. Set GEMINI_OAUTH_CLIENT_ID and "
|
|
187
|
+
"GEMINI_OAUTH_CLIENT_SECRET, or run `gemini` and retry."
|
|
188
|
+
)
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
def refresh_access_token(creds: dict) -> dict:
|
|
192
|
+
refresh_token = creds.get("refresh_token")
|
|
193
|
+
if not refresh_token:
|
|
194
|
+
raise RuntimeError("No refresh token in ~/.gemini/oauth_creds.json")
|
|
195
|
+
|
|
196
|
+
client_id, client_secret = _get_oauth_client_credentials(creds)
|
|
197
|
+
payload = urllib.parse.urlencode(
|
|
198
|
+
{
|
|
199
|
+
"grant_type": "refresh_token",
|
|
200
|
+
"refresh_token": refresh_token,
|
|
201
|
+
"client_id": client_id,
|
|
202
|
+
"client_secret": client_secret,
|
|
203
|
+
}
|
|
204
|
+
).encode()
|
|
205
|
+
|
|
206
|
+
req = urllib.request.Request(
|
|
207
|
+
TOKEN_URL,
|
|
208
|
+
data=payload,
|
|
209
|
+
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
210
|
+
)
|
|
211
|
+
with urllib.request.urlopen(req, timeout=10) as resp:
|
|
212
|
+
result = json.loads(resp.read())
|
|
213
|
+
|
|
214
|
+
updated = dict(creds)
|
|
215
|
+
updated["access_token"] = result["access_token"]
|
|
216
|
+
updated["token_type"] = result.get("token_type", updated.get("token_type", "Bearer"))
|
|
217
|
+
updated["scope"] = result.get("scope", updated.get("scope"))
|
|
218
|
+
updated["expiry_date"] = int(time.time() * 1000 + int(result.get("expires_in", 3600)) * 1000)
|
|
219
|
+
if result.get("id_token"):
|
|
220
|
+
updated["id_token"] = result["id_token"]
|
|
221
|
+
if result.get("refresh_token"):
|
|
222
|
+
updated["refresh_token"] = result["refresh_token"]
|
|
223
|
+
_write_oauth_credentials(updated)
|
|
224
|
+
return updated
|
|
225
|
+
|
|
226
|
+
|
|
227
|
+
def get_access_token() -> str:
|
|
228
|
+
creds = get_oauth_credentials()
|
|
229
|
+
if not creds:
|
|
230
|
+
raise RuntimeError("No OAuth credentials at ~/.gemini/oauth_creds.json")
|
|
231
|
+
|
|
232
|
+
expiry_date = int(creds.get("expiry_date", 0) or 0)
|
|
233
|
+
if time.time() * 1000 >= expiry_date - 60_000:
|
|
234
|
+
creds = refresh_access_token(creds)
|
|
235
|
+
|
|
236
|
+
token = creds.get("access_token")
|
|
237
|
+
if not token:
|
|
238
|
+
raise RuntimeError("No access token in ~/.gemini/oauth_creds.json")
|
|
239
|
+
return token
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
def _code_assist_post(method: str, payload: dict, access_token: str) -> dict:
|
|
243
|
+
req = urllib.request.Request(
|
|
244
|
+
f"{CODE_ASSIST_BASE_URL}:{method}",
|
|
245
|
+
data=json.dumps(payload).encode(),
|
|
246
|
+
headers={
|
|
247
|
+
"Authorization": f"Bearer {access_token}",
|
|
248
|
+
"Content-Type": "application/json",
|
|
249
|
+
"Accept": "application/json",
|
|
250
|
+
"User-Agent": "gemini-cli-usage/0.1.0",
|
|
251
|
+
},
|
|
252
|
+
)
|
|
253
|
+
with urllib.request.urlopen(req, timeout=15) as resp:
|
|
254
|
+
return json.loads(resp.read())
|
|
255
|
+
|
|
256
|
+
|
|
257
|
+
def _load_code_assist(access_token: str) -> dict:
|
|
258
|
+
project_id = (
|
|
259
|
+
os.environ.get("GOOGLE_CLOUD_PROJECT")
|
|
260
|
+
or os.environ.get("GOOGLE_CLOUD_PROJECT_ID")
|
|
261
|
+
or None
|
|
262
|
+
)
|
|
263
|
+
metadata = {
|
|
264
|
+
"ideType": "IDE_UNSPECIFIED",
|
|
265
|
+
"platform": "PLATFORM_UNSPECIFIED",
|
|
266
|
+
"pluginType": "GEMINI",
|
|
267
|
+
}
|
|
268
|
+
if project_id:
|
|
269
|
+
metadata["duetProject"] = project_id
|
|
270
|
+
return _code_assist_post(
|
|
271
|
+
"loadCodeAssist",
|
|
272
|
+
{
|
|
273
|
+
"cloudaicompanionProject": project_id,
|
|
274
|
+
"metadata": metadata,
|
|
275
|
+
},
|
|
276
|
+
access_token,
|
|
277
|
+
)
|
|
278
|
+
|
|
279
|
+
|
|
280
|
+
def _parse_quota_buckets(buckets: list[dict]) -> list[dict]:
|
|
281
|
+
parsed = []
|
|
282
|
+
for bucket in buckets:
|
|
283
|
+
remaining = None
|
|
284
|
+
limit = None
|
|
285
|
+
used_pct = None
|
|
286
|
+
remaining_fraction = bucket.get("remainingFraction")
|
|
287
|
+
|
|
288
|
+
try:
|
|
289
|
+
if bucket.get("remainingAmount") is not None:
|
|
290
|
+
remaining = int(bucket["remainingAmount"])
|
|
291
|
+
except (TypeError, ValueError):
|
|
292
|
+
remaining = None
|
|
293
|
+
|
|
294
|
+
if isinstance(remaining_fraction, int | float):
|
|
295
|
+
used_pct = (1 - float(remaining_fraction)) * 100
|
|
296
|
+
if remaining is not None and isinstance(remaining_fraction, int | float):
|
|
297
|
+
if remaining_fraction > 0:
|
|
298
|
+
limit = round(remaining / float(remaining_fraction))
|
|
299
|
+
|
|
300
|
+
parsed.append(
|
|
301
|
+
{
|
|
302
|
+
"model": bucket.get("modelId"),
|
|
303
|
+
"remaining": remaining,
|
|
304
|
+
"limit": limit,
|
|
305
|
+
"used_pct": used_pct,
|
|
306
|
+
"remaining_fraction": remaining_fraction,
|
|
307
|
+
"reset_time": bucket.get("resetTime"),
|
|
308
|
+
"token_type": bucket.get("tokenType"),
|
|
309
|
+
}
|
|
310
|
+
)
|
|
311
|
+
return parsed
|
|
312
|
+
|
|
313
|
+
|
|
314
|
+
def _select_summary_bucket(quota: dict) -> dict | None:
|
|
315
|
+
buckets = quota.get("buckets") or []
|
|
316
|
+
if not isinstance(buckets, list):
|
|
317
|
+
return None
|
|
318
|
+
scored = [bucket for bucket in buckets if bucket.get("used_pct") is not None]
|
|
319
|
+
if scored:
|
|
320
|
+
return max(scored, key=lambda bucket: bucket["used_pct"])
|
|
321
|
+
return buckets[0] if buckets else None
|
|
322
|
+
|
|
323
|
+
|
|
324
|
+
def fetch_quota(project_root: Path | None = None) -> dict:
|
|
325
|
+
auth_type = get_auth_type(project_root)
|
|
326
|
+
if auth_type != "oauth-personal":
|
|
327
|
+
raise RuntimeError(
|
|
328
|
+
"Quota lookup requires Google login; current auth is "
|
|
329
|
+
f"{get_auth_label(auth_type)}"
|
|
330
|
+
)
|
|
331
|
+
|
|
332
|
+
access_token = get_access_token()
|
|
333
|
+
load_res = _load_code_assist(access_token)
|
|
334
|
+
|
|
335
|
+
env_project = (
|
|
336
|
+
os.environ.get("GOOGLE_CLOUD_PROJECT")
|
|
337
|
+
or os.environ.get("GOOGLE_CLOUD_PROJECT_ID")
|
|
338
|
+
or None
|
|
339
|
+
)
|
|
340
|
+
project_id = load_res.get("cloudaicompanionProject") or env_project
|
|
341
|
+
if not project_id:
|
|
342
|
+
raise RuntimeError(
|
|
343
|
+
"No Code Assist project ID available. Set GOOGLE_CLOUD_PROJECT if your account requires it."
|
|
344
|
+
)
|
|
345
|
+
|
|
346
|
+
quota_res = _code_assist_post(
|
|
347
|
+
"retrieveUserQuota", {"project": project_id}, access_token
|
|
348
|
+
)
|
|
349
|
+
current_tier = load_res.get("currentTier") or {}
|
|
350
|
+
paid_tier = load_res.get("paidTier") or {}
|
|
351
|
+
result = {
|
|
352
|
+
"project_id": project_id,
|
|
353
|
+
"user_tier": paid_tier.get("id") or current_tier.get("id"),
|
|
354
|
+
"user_tier_name": paid_tier.get("name") or current_tier.get("name"),
|
|
355
|
+
"buckets": _parse_quota_buckets(quota_res.get("buckets") or []),
|
|
356
|
+
}
|
|
357
|
+
result["summary_bucket"] = _select_summary_bucket(result)
|
|
358
|
+
return result
|
|
359
|
+
|
|
360
|
+
|
|
361
|
+
def build_usage_json(project_root: Path | None = None) -> dict:
|
|
362
|
+
root = (project_root or Path.cwd()).resolve()
|
|
363
|
+
auth_type = get_auth_type(root)
|
|
364
|
+
|
|
365
|
+
result = {
|
|
366
|
+
"project_root": str(root),
|
|
367
|
+
"auth_type": auth_type,
|
|
368
|
+
"auth_label": get_auth_label(auth_type),
|
|
369
|
+
"source": [],
|
|
370
|
+
"updated_at": _iso_now(),
|
|
371
|
+
}
|
|
372
|
+
|
|
373
|
+
try:
|
|
374
|
+
result["account_quota"] = fetch_quota(root)
|
|
375
|
+
result["source"].append("quota_api")
|
|
376
|
+
except Exception as exc:
|
|
377
|
+
result["quota_error"] = str(exc)
|
|
378
|
+
|
|
379
|
+
return result
|
|
380
|
+
|
|
381
|
+
|
|
382
|
+
def get_usage_file() -> Path:
|
|
383
|
+
override = os.environ.get("GEMINI_CLI_USAGE_FILE") or os.environ.get(
|
|
384
|
+
"GEMINI_USAGE_FILE"
|
|
385
|
+
)
|
|
386
|
+
return Path(override).expanduser() if override else DEFAULT_USAGE_FILE
|
|
387
|
+
|
|
388
|
+
|
|
389
|
+
def write_usage_file(data: dict):
|
|
390
|
+
usage_file = get_usage_file()
|
|
391
|
+
usage_file.parent.mkdir(parents=True, exist_ok=True)
|
|
392
|
+
usage_file.write_text(json.dumps(data, indent=2) + "\n")
|
|
393
|
+
|
|
394
|
+
|
|
395
|
+
def _format_duration_until(iso_timestamp: str | None) -> str:
|
|
396
|
+
reset = _parse_iso(iso_timestamp)
|
|
397
|
+
if not reset:
|
|
398
|
+
return ""
|
|
399
|
+
seconds = int((reset - datetime.now(UTC)).total_seconds())
|
|
400
|
+
if seconds <= 0:
|
|
401
|
+
return ""
|
|
402
|
+
minutes = seconds // 60
|
|
403
|
+
if minutes >= 60:
|
|
404
|
+
return f"{minutes // 60}h{minutes % 60}m"
|
|
405
|
+
return f"{minutes}m"
|
|
406
|
+
|
|
407
|
+
|
|
408
|
+
def _color_pct(pct: float | int | None) -> str:
|
|
409
|
+
if pct is None:
|
|
410
|
+
return "?"
|
|
411
|
+
p = float(pct)
|
|
412
|
+
red = "\033[0;31m"
|
|
413
|
+
yellow = "\033[0;33m"
|
|
414
|
+
green = "\033[0;32m"
|
|
415
|
+
reset = "\033[0m"
|
|
416
|
+
color = red if p >= 70 else yellow if p >= 40 else green
|
|
417
|
+
return f"{color}{_format_pct(p)}{reset}"
|
|
418
|
+
|
|
419
|
+
|
|
420
|
+
def _format_pct(pct: float | int | None) -> str:
|
|
421
|
+
if pct is None:
|
|
422
|
+
return "?"
|
|
423
|
+
p = float(pct)
|
|
424
|
+
if p >= 1:
|
|
425
|
+
return f"{p:.1f}%"
|
|
426
|
+
return f"{p:.2f}%"
|
|
427
|
+
|
|
428
|
+
|
|
429
|
+
def _print_status(data: dict):
|
|
430
|
+
dim = "\033[0;90m"
|
|
431
|
+
reset = "\033[0m"
|
|
432
|
+
|
|
433
|
+
print(f"Project: {Path(data['project_root']).name}")
|
|
434
|
+
print(f"Auth: {data.get('auth_label', 'unknown')}")
|
|
435
|
+
|
|
436
|
+
quota = data.get("account_quota")
|
|
437
|
+
if quota:
|
|
438
|
+
bucket_names = [
|
|
439
|
+
(bucket.get("model") or "unknown")
|
|
440
|
+
for bucket in quota.get("buckets", [])
|
|
441
|
+
if isinstance(bucket, dict)
|
|
442
|
+
]
|
|
443
|
+
name_width = max(map(len, bucket_names), default=len("Quota"))
|
|
444
|
+
|
|
445
|
+
def print_bucket_line(label: str, bucket: dict):
|
|
446
|
+
reset_time = _format_duration_until(bucket.get("reset_time"))
|
|
447
|
+
reset_part = f" resets {reset_time}" if reset_time else ""
|
|
448
|
+
remaining = bucket.get("remaining")
|
|
449
|
+
limit = bucket.get("limit")
|
|
450
|
+
remain_part = (
|
|
451
|
+
f" {remaining} / {limit} remaining"
|
|
452
|
+
if remaining is not None and limit is not None
|
|
453
|
+
else ""
|
|
454
|
+
)
|
|
455
|
+
print(
|
|
456
|
+
f" {label:{name_width}s} {_color_pct(bucket.get('used_pct'))} used"
|
|
457
|
+
f"{remain_part}{dim}{reset_part}{reset}"
|
|
458
|
+
)
|
|
459
|
+
|
|
460
|
+
for bucket in quota.get("buckets", []):
|
|
461
|
+
model = bucket.get("model") or "unknown"
|
|
462
|
+
print_bucket_line(model, bucket)
|
|
463
|
+
elif data.get("quota_error"):
|
|
464
|
+
print(f" {'Quota':20s} {dim}{data['quota_error']}{reset}")
|
|
465
|
+
|
|
466
|
+
|
|
467
|
+
def _statusline_text(data: dict) -> str:
|
|
468
|
+
parts = []
|
|
469
|
+
quota = data.get("account_quota")
|
|
470
|
+
if quota and quota.get("summary_bucket"):
|
|
471
|
+
summary_bucket = quota["summary_bucket"]
|
|
472
|
+
if summary_bucket.get("used_pct") is not None:
|
|
473
|
+
parts.append(f"q:{_format_pct(summary_bucket['used_pct'])}")
|
|
474
|
+
reset_time = _format_duration_until(summary_bucket.get("reset_time"))
|
|
475
|
+
if reset_time:
|
|
476
|
+
parts.append(f"reset:{reset_time}")
|
|
477
|
+
elif data.get("quota_error"):
|
|
478
|
+
parts.append("q:err")
|
|
479
|
+
return " ".join(parts)
|
|
480
|
+
|
|
481
|
+
|
|
482
|
+
def _get_cached_usage(
|
|
483
|
+
project_root: Path | None = None,
|
|
484
|
+
max_age: int = CACHE_MAX_AGE,
|
|
485
|
+
force_refresh: bool = False,
|
|
486
|
+
) -> dict:
|
|
487
|
+
usage_file = get_usage_file()
|
|
488
|
+
if not force_refresh:
|
|
489
|
+
try:
|
|
490
|
+
cached = json.loads(usage_file.read_text())
|
|
491
|
+
updated = _parse_iso(cached.get("updated_at"))
|
|
492
|
+
root = str((project_root or Path.cwd()).resolve())
|
|
493
|
+
if updated and cached.get("project_root") == root:
|
|
494
|
+
age = (datetime.now(UTC) - updated).total_seconds()
|
|
495
|
+
if age < max_age and "quota_api" in cached.get("source", []):
|
|
496
|
+
return cached
|
|
497
|
+
except Exception:
|
|
498
|
+
pass
|
|
499
|
+
|
|
500
|
+
try:
|
|
501
|
+
fresh = build_usage_json(project_root)
|
|
502
|
+
write_usage_file(fresh)
|
|
503
|
+
return fresh
|
|
504
|
+
except Exception:
|
|
505
|
+
try:
|
|
506
|
+
return json.loads(usage_file.read_text())
|
|
507
|
+
except Exception:
|
|
508
|
+
return build_usage_json(project_root)
|
|
509
|
+
|
|
510
|
+
|
|
511
|
+
def cmd_status(args):
|
|
512
|
+
data = build_usage_json(project_root=Path(args.root).resolve() if args.root else None)
|
|
513
|
+
_print_status(data)
|
|
514
|
+
|
|
515
|
+
|
|
516
|
+
def cmd_json(args):
|
|
517
|
+
data = build_usage_json(project_root=Path(args.root).resolve() if args.root else None)
|
|
518
|
+
print(json.dumps(data, indent=2))
|
|
519
|
+
|
|
520
|
+
|
|
521
|
+
def cmd_daemon(args):
|
|
522
|
+
signal.signal(signal.SIGINT, lambda *_: sys.exit(0))
|
|
523
|
+
signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
|
|
524
|
+
|
|
525
|
+
root = Path(args.root).resolve() if args.root else None
|
|
526
|
+
usage_file = get_usage_file()
|
|
527
|
+
|
|
528
|
+
print(f"gemini-cli-usage daemon started (refreshing every {args.interval}s)")
|
|
529
|
+
print(f"Writing to {usage_file}")
|
|
530
|
+
|
|
531
|
+
while True:
|
|
532
|
+
try:
|
|
533
|
+
data = build_usage_json(project_root=root)
|
|
534
|
+
write_usage_file(data)
|
|
535
|
+
print(
|
|
536
|
+
f"[{datetime.now().strftime('%H:%M:%S')}] "
|
|
537
|
+
f"{_statusline_text(data)}"
|
|
538
|
+
)
|
|
539
|
+
except Exception as exc:
|
|
540
|
+
print(
|
|
541
|
+
f"[{datetime.now().strftime('%H:%M:%S')}] Error: {exc}",
|
|
542
|
+
file=sys.stderr,
|
|
543
|
+
)
|
|
544
|
+
time.sleep(args.interval)
|
|
545
|
+
|
|
546
|
+
|
|
547
|
+
def cmd_statusline(args):
|
|
548
|
+
data = _get_cached_usage(
|
|
549
|
+
project_root=Path(args.root).resolve() if args.root else None,
|
|
550
|
+
max_age=args.max_age,
|
|
551
|
+
force_refresh=args.refresh,
|
|
552
|
+
)
|
|
553
|
+
print(_statusline_text(data))
|
|
554
|
+
|
|
555
|
+
|
|
556
|
+
def cmd_refresh(args):
|
|
557
|
+
data = build_usage_json(project_root=Path(args.root).resolve() if args.root else None)
|
|
558
|
+
write_usage_file(data)
|
|
559
|
+
_print_status(data)
|
|
560
|
+
|
|
561
|
+
|
|
562
|
+
def cmd_install(_args):
|
|
563
|
+
print(
|
|
564
|
+
"Install with:\n"
|
|
565
|
+
" uv tool install gemini-cli-usage\n\n"
|
|
566
|
+
"For local development:\n"
|
|
567
|
+
" uv tool install .\n\n"
|
|
568
|
+
"Then run:\n"
|
|
569
|
+
" gemini-cli-usage\n"
|
|
570
|
+
" gemini-cli-usage statusline\n"
|
|
571
|
+
" gemini-cli-usage refresh\n"
|
|
572
|
+
)
|
|
573
|
+
|
|
574
|
+
|
|
575
|
+
def _build_parser() -> argparse.ArgumentParser:
|
|
576
|
+
parser = argparse.ArgumentParser(description="Gemini CLI quota monitor")
|
|
577
|
+
parser.add_argument(
|
|
578
|
+
"command",
|
|
579
|
+
nargs="?",
|
|
580
|
+
default="status",
|
|
581
|
+
choices=["status", "json", "daemon", "statusline", "refresh", "install"],
|
|
582
|
+
)
|
|
583
|
+
parser.add_argument(
|
|
584
|
+
"--root",
|
|
585
|
+
help="Project root to inspect (default: current working directory)",
|
|
586
|
+
)
|
|
587
|
+
parser.add_argument(
|
|
588
|
+
"-i",
|
|
589
|
+
"--interval",
|
|
590
|
+
type=int,
|
|
591
|
+
default=DAEMON_INTERVAL,
|
|
592
|
+
help="Daemon refresh interval in seconds",
|
|
593
|
+
)
|
|
594
|
+
parser.add_argument(
|
|
595
|
+
"--max-age",
|
|
596
|
+
type=int,
|
|
597
|
+
default=CACHE_MAX_AGE,
|
|
598
|
+
help="Maximum cache age in seconds for statusline",
|
|
599
|
+
)
|
|
600
|
+
parser.add_argument(
|
|
601
|
+
"--refresh",
|
|
602
|
+
action="store_true",
|
|
603
|
+
help="Ignore cache and force a fresh fetch where applicable",
|
|
604
|
+
)
|
|
605
|
+
return parser
|
|
606
|
+
|
|
607
|
+
|
|
608
|
+
def main():
|
|
609
|
+
parser = _build_parser()
|
|
610
|
+
args = parser.parse_args()
|
|
611
|
+
|
|
612
|
+
commands = {
|
|
613
|
+
"status": cmd_status,
|
|
614
|
+
"json": cmd_json,
|
|
615
|
+
"daemon": cmd_daemon,
|
|
616
|
+
"statusline": cmd_statusline,
|
|
617
|
+
"refresh": cmd_refresh,
|
|
618
|
+
"install": cmd_install,
|
|
619
|
+
}
|
|
620
|
+
commands[args.command](args)
|
|
621
|
+
|
|
622
|
+
|
|
623
|
+
if __name__ == "__main__":
|
|
624
|
+
main()
|
|
@@ -0,0 +1,271 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import tempfile
|
|
6
|
+
import unittest
|
|
7
|
+
from datetime import UTC, datetime, timedelta
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
from unittest import mock
|
|
10
|
+
from urllib.parse import parse_qs
|
|
11
|
+
|
|
12
|
+
import gemini_cli_usage
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _write_json(path: Path, payload: dict):
|
|
16
|
+
path.parent.mkdir(parents=True, exist_ok=True)
|
|
17
|
+
path.write_text(json.dumps(payload) + "\n")
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _settings_payload(selected_type: str) -> dict:
|
|
21
|
+
return {"security": {"auth": {"selectedType": selected_type}}}
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def _usage_payload(project_root: Path, updated_at: str) -> dict:
|
|
25
|
+
return {
|
|
26
|
+
"project_root": str(project_root.resolve()),
|
|
27
|
+
"auth_type": "oauth-personal",
|
|
28
|
+
"auth_label": "Google login",
|
|
29
|
+
"source": ["quota_api"],
|
|
30
|
+
"updated_at": updated_at,
|
|
31
|
+
"account_quota": {
|
|
32
|
+
"buckets": [],
|
|
33
|
+
"summary_bucket": None,
|
|
34
|
+
},
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
class GeminiUsageTests(unittest.TestCase):
|
|
39
|
+
def test_env_auth_overrides_workspace_and_global_settings(self):
|
|
40
|
+
with tempfile.TemporaryDirectory() as tmp:
|
|
41
|
+
tmp_path = Path(tmp)
|
|
42
|
+
project_root = tmp_path / "project"
|
|
43
|
+
global_settings = tmp_path / "global-settings.json"
|
|
44
|
+
workspace_settings = project_root / ".gemini" / "settings.json"
|
|
45
|
+
|
|
46
|
+
_write_json(global_settings, _settings_payload("oauth-personal"))
|
|
47
|
+
_write_json(workspace_settings, _settings_payload("vertex-ai"))
|
|
48
|
+
|
|
49
|
+
with (
|
|
50
|
+
mock.patch.object(gemini_cli_usage, "SETTINGS_FILE", global_settings),
|
|
51
|
+
mock.patch.dict(os.environ, {"GEMINI_API_KEY": "secret"}, clear=True),
|
|
52
|
+
):
|
|
53
|
+
self.assertEqual(
|
|
54
|
+
gemini_cli_usage.get_auth_type(project_root), "gemini-api-key"
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
def test_workspace_settings_override_global_settings(self):
|
|
58
|
+
with tempfile.TemporaryDirectory() as tmp:
|
|
59
|
+
tmp_path = Path(tmp)
|
|
60
|
+
project_root = tmp_path / "project"
|
|
61
|
+
global_settings = tmp_path / "global-settings.json"
|
|
62
|
+
workspace_settings = project_root / ".gemini" / "settings.json"
|
|
63
|
+
|
|
64
|
+
_write_json(global_settings, _settings_payload("oauth-personal"))
|
|
65
|
+
_write_json(workspace_settings, _settings_payload("vertex-ai"))
|
|
66
|
+
|
|
67
|
+
with (
|
|
68
|
+
mock.patch.object(gemini_cli_usage, "SETTINGS_FILE", global_settings),
|
|
69
|
+
mock.patch.dict(os.environ, {}, clear=True),
|
|
70
|
+
):
|
|
71
|
+
self.assertEqual(gemini_cli_usage.get_auth_type(project_root), "vertex-ai")
|
|
72
|
+
|
|
73
|
+
def test_build_usage_json_is_quota_only(self):
|
|
74
|
+
with tempfile.TemporaryDirectory() as tmp:
|
|
75
|
+
tmp_path = Path(tmp)
|
|
76
|
+
project_root = tmp_path / "project"
|
|
77
|
+
project_root.mkdir()
|
|
78
|
+
quota = {"buckets": [], "summary_bucket": None}
|
|
79
|
+
|
|
80
|
+
with mock.patch.object(gemini_cli_usage, "fetch_quota", return_value=quota):
|
|
81
|
+
usage = gemini_cli_usage.build_usage_json(project_root)
|
|
82
|
+
|
|
83
|
+
self.assertEqual(usage["project_root"], str(project_root.resolve()))
|
|
84
|
+
self.assertEqual(usage["source"], ["quota_api"])
|
|
85
|
+
self.assertNotIn("local_usage", usage)
|
|
86
|
+
self.assertEqual(usage["account_quota"], quota)
|
|
87
|
+
|
|
88
|
+
def test_summary_bucket_and_statusline_use_highest_used_bucket(self):
|
|
89
|
+
future = (datetime.now(UTC) + timedelta(hours=2)).isoformat()
|
|
90
|
+
quota = {
|
|
91
|
+
"buckets": [
|
|
92
|
+
{
|
|
93
|
+
"model": "gemini-2.5-flash-lite",
|
|
94
|
+
"used_pct": 0.07,
|
|
95
|
+
"reset_time": future,
|
|
96
|
+
},
|
|
97
|
+
{
|
|
98
|
+
"model": "gemini-2.5-pro",
|
|
99
|
+
"used_pct": 3.5,
|
|
100
|
+
"reset_time": future,
|
|
101
|
+
},
|
|
102
|
+
]
|
|
103
|
+
}
|
|
104
|
+
|
|
105
|
+
summary = gemini_cli_usage._select_summary_bucket(quota)
|
|
106
|
+
|
|
107
|
+
data = {
|
|
108
|
+
"account_quota": {
|
|
109
|
+
"summary_bucket": summary,
|
|
110
|
+
},
|
|
111
|
+
}
|
|
112
|
+
|
|
113
|
+
statusline = gemini_cli_usage._statusline_text(data)
|
|
114
|
+
|
|
115
|
+
self.assertEqual(summary["model"], "gemini-2.5-pro")
|
|
116
|
+
self.assertIn("q:3.5%", statusline)
|
|
117
|
+
self.assertNotIn("q:0.07%", statusline)
|
|
118
|
+
|
|
119
|
+
def test_force_refresh_bypasses_cache(self):
|
|
120
|
+
with tempfile.TemporaryDirectory() as tmp:
|
|
121
|
+
tmp_path = Path(tmp)
|
|
122
|
+
project_root = tmp_path / "project"
|
|
123
|
+
project_root.mkdir()
|
|
124
|
+
usage_file = tmp_path / "usage-limits.json"
|
|
125
|
+
now = datetime.now(UTC).isoformat()
|
|
126
|
+
cached = _usage_payload(project_root, now)
|
|
127
|
+
fresh = _usage_payload(project_root, "2026-03-12T12:00:00+00:00")
|
|
128
|
+
fresh["auth_type"] = "gemini-api-key"
|
|
129
|
+
fresh["auth_label"] = "Gemini API key"
|
|
130
|
+
|
|
131
|
+
usage_file.write_text(json.dumps(cached) + "\n")
|
|
132
|
+
|
|
133
|
+
with (
|
|
134
|
+
mock.patch.object(gemini_cli_usage, "DEFAULT_USAGE_FILE", usage_file),
|
|
135
|
+
mock.patch.dict(os.environ, {}, clear=True),
|
|
136
|
+
mock.patch.object(gemini_cli_usage, "build_usage_json", return_value=fresh) as build_mock,
|
|
137
|
+
):
|
|
138
|
+
result = gemini_cli_usage._get_cached_usage(project_root=project_root)
|
|
139
|
+
self.assertEqual(result["updated_at"], now)
|
|
140
|
+
build_mock.assert_not_called()
|
|
141
|
+
|
|
142
|
+
result = gemini_cli_usage._get_cached_usage(
|
|
143
|
+
project_root=project_root, force_refresh=True
|
|
144
|
+
)
|
|
145
|
+
|
|
146
|
+
self.assertEqual(result, fresh)
|
|
147
|
+
self.assertEqual(json.loads(usage_file.read_text()), fresh)
|
|
148
|
+
|
|
149
|
+
def test_refresh_access_token_uses_env_client_credentials(self):
|
|
150
|
+
creds = {
|
|
151
|
+
"refresh_token": "refresh-token",
|
|
152
|
+
"access_token": "old-access-token",
|
|
153
|
+
"expiry_date": 0,
|
|
154
|
+
}
|
|
155
|
+
|
|
156
|
+
class FakeResponse:
|
|
157
|
+
def __enter__(self):
|
|
158
|
+
return self
|
|
159
|
+
|
|
160
|
+
def __exit__(self, exc_type, exc, tb):
|
|
161
|
+
return False
|
|
162
|
+
|
|
163
|
+
def read(self):
|
|
164
|
+
return json.dumps(
|
|
165
|
+
{
|
|
166
|
+
"access_token": "new-access-token",
|
|
167
|
+
"expires_in": 3600,
|
|
168
|
+
"token_type": "Bearer",
|
|
169
|
+
}
|
|
170
|
+
).encode()
|
|
171
|
+
|
|
172
|
+
def fake_urlopen(request, timeout=10):
|
|
173
|
+
payload = parse_qs(request.data.decode())
|
|
174
|
+
self.assertEqual(payload["client_id"], ["test-client-id"])
|
|
175
|
+
self.assertEqual(payload["client_secret"], ["test-client-secret"])
|
|
176
|
+
self.assertEqual(payload["refresh_token"], ["refresh-token"])
|
|
177
|
+
self.assertEqual(payload["grant_type"], ["refresh_token"])
|
|
178
|
+
return FakeResponse()
|
|
179
|
+
|
|
180
|
+
with (
|
|
181
|
+
mock.patch.dict(
|
|
182
|
+
os.environ,
|
|
183
|
+
{
|
|
184
|
+
"GEMINI_OAUTH_CLIENT_ID": "test-client-id",
|
|
185
|
+
"GEMINI_OAUTH_CLIENT_SECRET": "test-client-secret",
|
|
186
|
+
},
|
|
187
|
+
clear=True,
|
|
188
|
+
),
|
|
189
|
+
mock.patch.object(gemini_cli_usage.urllib.request, "urlopen", side_effect=fake_urlopen),
|
|
190
|
+
):
|
|
191
|
+
updated = gemini_cli_usage.refresh_access_token(creds)
|
|
192
|
+
|
|
193
|
+
self.assertEqual(updated["access_token"], "new-access-token")
|
|
194
|
+
self.assertGreater(updated["expiry_date"], 0)
|
|
195
|
+
|
|
196
|
+
def test_refresh_access_token_reads_client_credentials_from_installed_gemini(self):
|
|
197
|
+
with tempfile.TemporaryDirectory() as tmp:
|
|
198
|
+
tmp_path = Path(tmp)
|
|
199
|
+
package_root = tmp_path / "lib" / "node_modules" / "@google" / "gemini-cli"
|
|
200
|
+
oauth2_path = (
|
|
201
|
+
package_root
|
|
202
|
+
/ "node_modules"
|
|
203
|
+
/ "@google"
|
|
204
|
+
/ "gemini-cli-core"
|
|
205
|
+
/ "dist"
|
|
206
|
+
/ "src"
|
|
207
|
+
/ "code_assist"
|
|
208
|
+
/ "oauth2.js"
|
|
209
|
+
)
|
|
210
|
+
oauth2_path.parent.mkdir(parents=True, exist_ok=True)
|
|
211
|
+
oauth2_path.write_text(
|
|
212
|
+
"const OAUTH_CLIENT_ID = 'live-client-id';\n"
|
|
213
|
+
"const OAUTH_CLIENT_SECRET = 'live-client-secret';\n"
|
|
214
|
+
)
|
|
215
|
+
gemini_path = package_root / "dist" / "index.js"
|
|
216
|
+
gemini_path.parent.mkdir(parents=True, exist_ok=True)
|
|
217
|
+
gemini_path.write_text("// stub\n")
|
|
218
|
+
|
|
219
|
+
creds = {"refresh_token": "refresh-token"}
|
|
220
|
+
|
|
221
|
+
class FakeResponse:
|
|
222
|
+
def __enter__(self):
|
|
223
|
+
return self
|
|
224
|
+
|
|
225
|
+
def __exit__(self, exc_type, exc, tb):
|
|
226
|
+
return False
|
|
227
|
+
|
|
228
|
+
def read(self):
|
|
229
|
+
return json.dumps({"access_token": "live-access-token"}).encode()
|
|
230
|
+
|
|
231
|
+
def fake_urlopen(request, timeout=10):
|
|
232
|
+
payload = parse_qs(request.data.decode())
|
|
233
|
+
self.assertEqual(payload["client_id"], ["live-client-id"])
|
|
234
|
+
self.assertEqual(payload["client_secret"], ["live-client-secret"])
|
|
235
|
+
return FakeResponse()
|
|
236
|
+
|
|
237
|
+
with (
|
|
238
|
+
mock.patch.dict(os.environ, {}, clear=True),
|
|
239
|
+
mock.patch.object(gemini_cli_usage.shutil, "which", return_value=str(gemini_path)),
|
|
240
|
+
mock.patch.object(gemini_cli_usage.urllib.request, "urlopen", side_effect=fake_urlopen),
|
|
241
|
+
):
|
|
242
|
+
updated = gemini_cli_usage.refresh_access_token(creds)
|
|
243
|
+
|
|
244
|
+
self.assertEqual(updated["access_token"], "live-access-token")
|
|
245
|
+
|
|
246
|
+
def test_get_access_token_requires_local_client_metadata_when_expired(self):
|
|
247
|
+
with tempfile.TemporaryDirectory() as tmp:
|
|
248
|
+
oauth_file = Path(tmp) / "oauth_creds.json"
|
|
249
|
+
_write_json(
|
|
250
|
+
oauth_file,
|
|
251
|
+
{
|
|
252
|
+
"access_token": "expired-access-token",
|
|
253
|
+
"refresh_token": "refresh-token",
|
|
254
|
+
"expiry_date": 0,
|
|
255
|
+
},
|
|
256
|
+
)
|
|
257
|
+
|
|
258
|
+
with (
|
|
259
|
+
mock.patch.object(gemini_cli_usage, "OAUTH_FILE", oauth_file),
|
|
260
|
+
mock.patch.dict(os.environ, {}, clear=True),
|
|
261
|
+
mock.patch.object(gemini_cli_usage.shutil, "which", return_value=None),
|
|
262
|
+
):
|
|
263
|
+
with self.assertRaises(RuntimeError) as exc:
|
|
264
|
+
gemini_cli_usage.get_access_token()
|
|
265
|
+
|
|
266
|
+
self.assertIn("GEMINI_OAUTH_CLIENT_ID", str(exc.exception))
|
|
267
|
+
self.assertIn("run `gemini`", str(exc.exception))
|
|
268
|
+
|
|
269
|
+
|
|
270
|
+
if __name__ == "__main__":
|
|
271
|
+
unittest.main()
|