380 lines
12 KiB
Python
380 lines
12 KiB
Python
"""Scoped Claude Code quota collection and credential-safety contracts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import stat
|
|
import sys
|
|
from pathlib import Path
|
|
from urllib.error import HTTPError
|
|
|
|
import pytest
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
SCRIPT = ROOT / "services/hermes/scripts/ai_usage_claude.py"
|
|
|
|
|
|
def load_module():
|
|
"""Load a fresh Claude usage helper for environment-independent tests."""
|
|
name = "ai_usage_claude_scoped"
|
|
spec = importlib.util.spec_from_file_location(name, SCRIPT)
|
|
module = importlib.util.module_from_spec(spec)
|
|
assert spec and spec.loader
|
|
sys.modules[name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class Response:
|
|
"""Small bounded urllib response fixture."""
|
|
|
|
def __init__(self, body, *, headers=None):
|
|
self.body = body if isinstance(body, bytes) else json.dumps(body).encode()
|
|
self.headers = headers or {}
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *_args):
|
|
return None
|
|
|
|
def read(self, maximum):
|
|
assert maximum == 1 << 20
|
|
return self.body
|
|
|
|
|
|
def write_credentials(path, **overrides):
|
|
"""Write a private synthetic Claude OAuth document."""
|
|
credentials = {
|
|
"accessToken": "scoped-access",
|
|
"refreshToken": "scoped-refresh",
|
|
"expiresAt": 2_000_000,
|
|
"scopes": [
|
|
"user:profile",
|
|
"user:inference",
|
|
"user:sessions:claude_code",
|
|
],
|
|
}
|
|
credentials.update(overrides)
|
|
path.write_text(
|
|
json.dumps({"claudeAiOauth": credentials}) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
path.chmod(0o600)
|
|
|
|
|
|
def test_scoped_payload_selects_only_overall_and_exact_fable_limits():
|
|
mod = load_module()
|
|
secret = "provider-document-detail-must-not-leak"
|
|
|
|
payload = mod._scoped_payload(
|
|
{
|
|
"five_hour": {"utilization": 58, "resets_at": "2026-08-24T17:20:00Z"},
|
|
"seven_day": {"utilization": 59, "resets_at": "2026-08-27T18:00:00Z"},
|
|
"limits": [
|
|
{
|
|
"kind": "weekly_scoped",
|
|
"percent": 81,
|
|
"resets_at": "2026-08-27T18:00:00Z",
|
|
"scope": {"model": {"display_name": "Fable"}},
|
|
"private": secret,
|
|
},
|
|
{
|
|
"kind": "weekly_scoped",
|
|
"percent": 90,
|
|
"scope": {"model": {"display_name": "Unknown"}},
|
|
},
|
|
],
|
|
"account": secret,
|
|
}
|
|
)
|
|
|
|
assert payload == {
|
|
"five_hour": {
|
|
"utilization": 58.0,
|
|
"resets_at": "2026-08-24T17:20:00Z",
|
|
},
|
|
"seven_day": {
|
|
"utilization": 59.0,
|
|
"resets_at": "2026-08-27T18:00:00Z",
|
|
},
|
|
"seven_day_fable": {
|
|
"utilization": 81.0,
|
|
"resets_at": "2026-08-27T18:00:00Z",
|
|
},
|
|
}
|
|
assert secret not in repr(payload)
|
|
|
|
|
|
def test_scoped_query_uses_current_private_access_token(tmp_path, monkeypatch):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
write_credentials(credentials)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
monkeypatch.setattr(mod.time, "time", lambda: 1_000)
|
|
requests = []
|
|
|
|
def open_request(request, timeout):
|
|
requests.append(request)
|
|
assert timeout == 30
|
|
return Response(
|
|
{
|
|
"seven_day": {"utilization": 42},
|
|
"limits": [
|
|
{
|
|
"kind": "weekly_scoped",
|
|
"percent": 73,
|
|
"scope": {"model": {"display_name": "Fable 5"}},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod, "urlopen", open_request)
|
|
|
|
assert mod._query_scoped_usage() == {
|
|
"seven_day": {"utilization": 42.0},
|
|
"seven_day_fable": {"utilization": 73.0},
|
|
}
|
|
assert len(requests) == 1
|
|
assert requests[0].full_url == mod.USAGE_URL
|
|
assert requests[0].data is None
|
|
assert requests[0].get_header("Authorization") == "Bearer scoped-access"
|
|
|
|
|
|
def test_credential_refresh_expiry_exposes_only_a_timestamp(tmp_path, monkeypatch):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
write_credentials(credentials, refreshTokenExpiresAt=2_000_000)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
|
|
assert mod.credential_refresh_expiry_timestamp() == 2_000
|
|
|
|
write_credentials(credentials, refreshTokenExpiresAt=True)
|
|
assert mod.credential_refresh_expiry_timestamp() is None
|
|
credentials.unlink()
|
|
assert mod.credential_refresh_expiry_timestamp() is None
|
|
|
|
|
|
def test_expired_scoped_token_refreshes_and_persists_rotation(
|
|
tmp_path, monkeypatch
|
|
):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
write_credentials(credentials, expiresAt=900_000)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
monkeypatch.setattr(mod.time, "time", lambda: 1_000)
|
|
requests = []
|
|
|
|
def open_request(request, timeout):
|
|
requests.append(request)
|
|
assert timeout == 30
|
|
if request.full_url == mod.TOKEN_URL:
|
|
return Response(
|
|
{
|
|
"access_token": "rotated-access",
|
|
"refresh_token": "rotated-refresh",
|
|
"expires_in": 3600,
|
|
"refresh_token_expires_in": 7200,
|
|
"scope": "user:profile user:inference",
|
|
}
|
|
)
|
|
return Response(
|
|
{
|
|
"seven_day": {"utilization": 50},
|
|
"limits": [
|
|
{
|
|
"kind": "weekly_scoped",
|
|
"percent": 75,
|
|
"scope": {"model": {"display_name": "Fable"}},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(mod, "urlopen", open_request)
|
|
|
|
assert mod._query_scoped_usage()["seven_day_fable"]["utilization"] == 75
|
|
refresh_request = requests[0]
|
|
assert refresh_request.full_url == mod.TOKEN_URL
|
|
refresh_body = json.loads(refresh_request.data)
|
|
assert refresh_body == {
|
|
"grant_type": "refresh_token",
|
|
"refresh_token": "scoped-refresh",
|
|
"client_id": mod.OAUTH_CLIENT_ID,
|
|
"scope": "user:profile user:inference user:sessions:claude_code",
|
|
}
|
|
persisted = json.loads(credentials.read_text())["claudeAiOauth"]
|
|
assert persisted["accessToken"] == "rotated-access"
|
|
assert persisted["refreshToken"] == "rotated-refresh"
|
|
assert persisted["expiresAt"] == 4_600_000
|
|
assert persisted["refreshTokenExpiresAt"] == 8_200_000
|
|
assert stat.S_IMODE(credentials.stat().st_mode) == 0o600
|
|
|
|
|
|
def test_query_falls_back_to_setup_token_headers_when_scope_is_unavailable(
|
|
tmp_path, monkeypatch
|
|
):
|
|
mod = load_module()
|
|
missing_credentials = tmp_path / "missing-credentials"
|
|
setup_token = tmp_path / "setup-token"
|
|
setup_token.write_text("setup-only-token\n", encoding="utf-8")
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", missing_credentials)
|
|
monkeypatch.setattr(mod, "OAUTH_TOKEN_FILE", setup_token)
|
|
requests = []
|
|
|
|
def open_request(request, timeout):
|
|
requests.append(request)
|
|
assert timeout == 30
|
|
return Response(
|
|
{},
|
|
headers={
|
|
"anthropic-ratelimit-unified-5h-utilization": "0.25",
|
|
"anthropic-ratelimit-unified-7d-utilization": "0.40",
|
|
},
|
|
)
|
|
|
|
monkeypatch.setattr(mod, "urlopen", open_request)
|
|
|
|
assert mod.query_claude() == {
|
|
"five_hour": {"utilization": 25.0},
|
|
"seven_day": {"utilization": 40.0},
|
|
}
|
|
assert len(requests) == 1
|
|
assert requests[0].full_url == mod.MESSAGES_URL
|
|
assert requests[0].get_header("Authorization") == "Bearer setup-only-token"
|
|
|
|
|
|
def test_scoped_validation_and_unauthorized_retry_fail_closed(
|
|
tmp_path, monkeypatch
|
|
):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
write_credentials(credentials)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
monkeypatch.setattr(mod.time, "time", lambda: 1_000)
|
|
|
|
assert mod._used_percentage(True) is None
|
|
assert mod._used_percentage(-1) is None
|
|
assert mod._used_percentage(101) is None
|
|
with pytest.raises(RuntimeError, match="supported quota windows"):
|
|
mod._scoped_payload({"limits": [{"kind": "weekly_scoped"}]})
|
|
|
|
attempts = []
|
|
|
|
def open_request(request, timeout):
|
|
attempts.append(request.full_url)
|
|
if len(attempts) == 1:
|
|
raise HTTPError(request.full_url, 401, "secret-body", {}, None)
|
|
if request.full_url == mod.TOKEN_URL:
|
|
return Response(
|
|
{
|
|
"access_token": "retried-access",
|
|
"refresh_token": "retried-refresh",
|
|
"expires_in": 3600,
|
|
}
|
|
)
|
|
return Response({"seven_day": {"utilization": 12}})
|
|
|
|
monkeypatch.setattr(mod, "urlopen", open_request)
|
|
assert mod._query_scoped_usage() == {"seven_day": {"utilization": 12.0}}
|
|
assert attempts == [mod.USAGE_URL, mod.TOKEN_URL, mod.USAGE_URL]
|
|
|
|
|
|
def test_scoped_payload_ignores_malformed_and_unrelated_limit_entries():
|
|
mod = load_module()
|
|
|
|
assert mod._scoped_payload(
|
|
{
|
|
"five_hour": {"utilization": "invalid"},
|
|
"seven_day": {"utilization": 10},
|
|
"limits": [
|
|
None,
|
|
{"kind": "session", "percent": 20},
|
|
{
|
|
"kind": "weekly_scoped",
|
|
"percent": 30,
|
|
"scope": {"model": {"display_name": "Sonnet"}},
|
|
},
|
|
{
|
|
"kind": "weekly_scoped",
|
|
"percent": "invalid",
|
|
"scope": {"model": {"display_name": "Fable"}},
|
|
},
|
|
],
|
|
}
|
|
) == {"seven_day": {"utilization": 10.0}}
|
|
with pytest.raises(RuntimeError, match="not an object"):
|
|
mod._scoped_payload(None)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("content", "mode", "message"),
|
|
[
|
|
('{"claudeAiOauth":{}}\n', 0o644, "unsafe"),
|
|
("not-json\n", 0o600, "invalid"),
|
|
("[]\n", 0o600, "invalid"),
|
|
],
|
|
)
|
|
def test_private_credential_reader_rejects_unsafe_documents(
|
|
tmp_path, monkeypatch, content, mode, message
|
|
):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
credentials.write_text(content, encoding="utf-8")
|
|
credentials.chmod(mode)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
|
|
with pytest.raises(RuntimeError, match=message):
|
|
mod._read_credentials()
|
|
|
|
|
|
def test_credential_writer_and_refresh_validation_are_bounded(tmp_path, monkeypatch):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
write_credentials(credentials)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
monkeypatch.setattr(mod, "MAX_DOCUMENT_BYTES", 1)
|
|
with pytest.raises(RuntimeError, match="too large"):
|
|
mod._write_credentials({"claudeAiOauth": {"accessToken": "large"}})
|
|
|
|
monkeypatch.setattr(mod, "MAX_DOCUMENT_BYTES", 1 << 20)
|
|
with pytest.raises(RuntimeError, match="refresh credential is unavailable"):
|
|
mod._refresh_access_token({"claudeAiOauth": {}}, {})
|
|
|
|
document = {
|
|
"claudeAiOauth": {
|
|
"accessToken": "old-access",
|
|
"refreshToken": "old-refresh",
|
|
"scopes": ["unapproved:scope"],
|
|
}
|
|
}
|
|
for response_document in ([], {"access_token": "", "expires_in": True}):
|
|
monkeypatch.setattr(
|
|
mod,
|
|
"urlopen",
|
|
lambda request, timeout, value=response_document: Response(value),
|
|
)
|
|
with pytest.raises(RuntimeError, match="refresh response is invalid"):
|
|
mod._refresh_access_token(document, document["claudeAiOauth"])
|
|
|
|
|
|
def test_non_authentication_usage_error_is_not_retried(tmp_path, monkeypatch):
|
|
mod = load_module()
|
|
credentials = tmp_path / ".credentials.json"
|
|
write_credentials(credentials)
|
|
monkeypatch.setattr(mod, "CREDENTIALS_FILE", credentials)
|
|
monkeypatch.setattr(mod.time, "time", lambda: 1_000)
|
|
attempts = []
|
|
|
|
def reject(request, timeout):
|
|
attempts.append(request.full_url)
|
|
raise HTTPError(request.full_url, 503, "provider-detail", {}, None)
|
|
|
|
monkeypatch.setattr(mod, "urlopen", reject)
|
|
with pytest.raises(HTTPError):
|
|
mod._query_scoped_usage()
|
|
assert attempts == [mod.USAGE_URL]
|