fix(#4365): wrap /api/provider/quota in profile_env_for_active_request
Codex gate caught that the new load_pool() path runs under the process-default profile (sibling endpoints /api/providers + /api/models/live already wrap). Without it a multi-profile client sees+seeds the default profile's pool (#4247/#4067 profile-isolation class).
This commit is contained in:
1 parent
16a94ba0b3
commit
83a479d7a5
3 files changed
+693
-66
No files matched your search
+285
-65
@@ -679,6 +679,7 @@ else:
|
||||
print(json.dumps(_fetch_snapshot(provider, api_key)), flush=True)
|
||||
"""
|
||||
|
||||
|
||||
# SECTION: Provider ↔ env var mapping
|
||||
|
||||
# Maps canonical provider slug → env var name for API key.
|
||||
@@ -751,7 +752,201 @@ _OAUTH_PROVIDERS = frozenset({
|
||||
"xai-oauth",
|
||||
})
|
||||
|
||||
# SECTION: Helper functions
|
||||
|
||||
def _entry_value(entry, *names):
|
||||
for name in names:
|
||||
try:
|
||||
value = getattr(entry, name)
|
||||
except Exception:
|
||||
value = None
|
||||
if value in (None, ""):
|
||||
continue
|
||||
text = str(value).strip()
|
||||
if text:
|
||||
return text
|
||||
return None
|
||||
|
||||
|
||||
def _parse_dt(value):
|
||||
if value in (None, ""):
|
||||
return None
|
||||
if isinstance(value, (int, float)):
|
||||
try:
|
||||
return datetime.fromtimestamp(float(value), tz=timezone.utc)
|
||||
except Exception:
|
||||
return None
|
||||
text = str(value).strip()
|
||||
if not text:
|
||||
return None
|
||||
if text.endswith("Z"):
|
||||
text = text[:-1] + "+00:00"
|
||||
try:
|
||||
dt = datetime.fromisoformat(text)
|
||||
except ValueError:
|
||||
return None
|
||||
return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _iso(value):
|
||||
if value in (None, ""):
|
||||
return None
|
||||
if hasattr(value, "isoformat"):
|
||||
text = value.isoformat()
|
||||
return text.replace("+00:00", "Z")
|
||||
text = str(value).strip()
|
||||
return text or None
|
||||
|
||||
|
||||
def _entry_exhausted_ttl_seconds(error_code):
|
||||
code = str(error_code or "").strip()
|
||||
if code == "401":
|
||||
return 5 * 60
|
||||
return 60 * 60
|
||||
|
||||
|
||||
def _entry_pool_exhausted_until(entry):
|
||||
if str(_entry_value(entry, "last_status") or "").strip().lower() != "exhausted":
|
||||
return None
|
||||
reset_at = _parse_dt(getattr(entry, "last_error_reset_at", None))
|
||||
if reset_at is not None:
|
||||
return reset_at
|
||||
status_at = _parse_dt(getattr(entry, "last_status_at", None))
|
||||
if status_at is None:
|
||||
return None
|
||||
return status_at + timedelta(seconds=_entry_exhausted_ttl_seconds(_entry_value(entry, "last_error_code")))
|
||||
|
||||
|
||||
def _entry_is_pool_exhausted(entry):
|
||||
exhausted_until = _entry_pool_exhausted_until(entry)
|
||||
return exhausted_until is not None and datetime.now(timezone.utc) < exhausted_until
|
||||
|
||||
|
||||
def _safe_entry_label(entry, index):
|
||||
label = _entry_value(entry, "label", "source") or ""
|
||||
if not label:
|
||||
label = "Credential " + str(index)
|
||||
label = " ".join(str(label).split())
|
||||
if len(label) > 64:
|
||||
label = label[:61].rstrip() + "..."
|
||||
return label
|
||||
|
||||
|
||||
def _entry_pool_retry_after(entry):
|
||||
return _iso(_entry_pool_exhausted_until(entry))
|
||||
|
||||
|
||||
def _entry_pool_exhausted_reason(entry):
|
||||
code = _entry_value(entry, "last_error_code")
|
||||
reset_at = _entry_pool_retry_after(entry)
|
||||
reason = "Credential pool marked this credential exhausted"
|
||||
if code:
|
||||
reason += " after provider status " + code
|
||||
if reset_at:
|
||||
reason += "; retry after " + reset_at
|
||||
return reason + "."
|
||||
|
||||
|
||||
def _local_pool_snapshot(provider):
|
||||
"""Probe-free pool snapshot from local auth.json entries.
|
||||
|
||||
Returns a SimpleNamespace compatible with _serialize_account_usage_snapshot,
|
||||
or None if the provider has no pool or no entries.
|
||||
"""
|
||||
try:
|
||||
from agent.credential_pool import load_pool
|
||||
from api.config import _is_ambient_gh_cli_entry
|
||||
|
||||
pool = load_pool(provider)
|
||||
entries = list(pool.entries()) if pool is not None and hasattr(pool, "entries") else []
|
||||
except Exception:
|
||||
return None
|
||||
if not entries:
|
||||
return None
|
||||
|
||||
rows = []
|
||||
available_count = 0
|
||||
exhausted_count = 0
|
||||
dead_count = 0
|
||||
for index, entry in enumerate(entries, start=1):
|
||||
source = str(_entry_value(entry, "source") or "")
|
||||
label_val = str(_entry_value(entry, "label", "source") or "")
|
||||
key_source = str(_entry_value(entry, "key_source") or "")
|
||||
if _is_ambient_gh_cli_entry(source, label_val, key_source):
|
||||
continue
|
||||
label = _safe_entry_label(entry, index)
|
||||
entry_status = str(_entry_value(entry, "last_status") or "").strip().lower()
|
||||
if entry_status == "dead":
|
||||
dead_count += 1
|
||||
rows.append({
|
||||
"label": label,
|
||||
"status": "dead",
|
||||
"plan": None,
|
||||
"windows": [],
|
||||
"details": [],
|
||||
"unavailable_reason": "Credential permanently revoked or invalid.",
|
||||
"retry_after": None,
|
||||
"fetched_at": None,
|
||||
})
|
||||
elif _entry_is_pool_exhausted(entry):
|
||||
exhausted_count += 1
|
||||
rows.append({
|
||||
"label": label,
|
||||
"status": "exhausted",
|
||||
"plan": None,
|
||||
"windows": [],
|
||||
"details": [],
|
||||
"unavailable_reason": _entry_pool_exhausted_reason(entry),
|
||||
"retry_after": _entry_pool_retry_after(entry),
|
||||
"fetched_at": None,
|
||||
})
|
||||
else:
|
||||
available_count += 1
|
||||
rows.append({
|
||||
"label": label,
|
||||
"status": "available",
|
||||
"plan": None,
|
||||
"windows": [],
|
||||
"details": [],
|
||||
"unavailable_reason": None,
|
||||
"retry_after": None,
|
||||
"fetched_at": None,
|
||||
})
|
||||
|
||||
if not rows:
|
||||
return None
|
||||
|
||||
total = available_count + exhausted_count + dead_count
|
||||
pool_dict = {
|
||||
"total_credentials": total,
|
||||
"queried_credentials": 0,
|
||||
"available_credentials": available_count,
|
||||
"exhausted_credentials": exhausted_count,
|
||||
"dead_credentials": dead_count,
|
||||
"failed_credentials": 0,
|
||||
"plans": [],
|
||||
"next_reset_at": None,
|
||||
"best_remaining_by_window": [],
|
||||
"credentials": rows,
|
||||
}
|
||||
|
||||
details = [str(available_count) + "/" + str(total) + " credentials available"]
|
||||
if exhausted_count:
|
||||
details.append(str(exhausted_count) + " exhausted")
|
||||
if dead_count:
|
||||
details.append(str(dead_count) + " dead")
|
||||
|
||||
return SimpleNamespace(
|
||||
provider=provider,
|
||||
source="local_pool",
|
||||
title="Credential pool",
|
||||
plan=None,
|
||||
windows=(),
|
||||
details=tuple(details),
|
||||
available=available_count > 0,
|
||||
unavailable_reason=None if available_count > 0 else "All pool credentials are unavailable.",
|
||||
fetched_at=datetime.now(timezone.utc),
|
||||
pool=pool_dict,
|
||||
)
|
||||
|
||||
|
||||
def _get_hermes_home() -> Path:
|
||||
@@ -1721,79 +1916,104 @@ def get_provider_quota(provider_id: str | None = None, *, refresh: bool = False)
|
||||
if provider in _ACCOUNT_USAGE_PROVIDERS:
|
||||
return _provider_account_usage_status(provider, display_name, refresh=refresh)
|
||||
|
||||
if provider != "openrouter":
|
||||
detail = "OpenAI/Anthropic rate-limit headers are a follow-up once WebUI captures provider response metadata."
|
||||
if provider == "openrouter":
|
||||
api_key = _get_provider_api_key("openrouter")
|
||||
if not api_key:
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "no_key",
|
||||
"quota": None,
|
||||
"message": "OpenRouter quota status needs an OPENROUTER_API_KEY configured on the server.",
|
||||
}
|
||||
req = urllib.request.Request(
|
||||
_OPENROUTER_KEY_URL,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=_PROVIDER_QUOTA_TIMEOUT_SECONDS) as resp:
|
||||
raw = resp.read()
|
||||
payload = json.loads(raw.decode("utf-8")) if isinstance(raw, (bytes, bytearray)) else json.loads(raw)
|
||||
quota = _sanitize_openrouter_quota(payload)
|
||||
return {
|
||||
"ok": True,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "available",
|
||||
"label": "OpenRouter credits",
|
||||
"quota": quota,
|
||||
"message": "OpenRouter quota status loaded.",
|
||||
}
|
||||
except urllib.error.HTTPError as exc:
|
||||
status = "invalid_key" if exc.code in (401, 403) else "unavailable"
|
||||
message = (
|
||||
"OpenRouter rejected the configured API key."
|
||||
if status == "invalid_key"
|
||||
else "OpenRouter quota status is temporarily unavailable."
|
||||
)
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": status,
|
||||
"quota": None,
|
||||
"message": message,
|
||||
}
|
||||
except (TimeoutError, urllib.error.URLError, json.JSONDecodeError, OSError, ValueError):
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "unavailable",
|
||||
"quota": None,
|
||||
"message": "OpenRouter quota status is temporarily unavailable.",
|
||||
}
|
||||
|
||||
local_snapshot = _local_pool_snapshot(provider)
|
||||
if local_snapshot is not None:
|
||||
account_limits = _serialize_account_usage_snapshot(local_snapshot)
|
||||
if account_limits and account_limits.get("available"):
|
||||
return {
|
||||
"ok": True,
|
||||
"provider": provider,
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "available",
|
||||
"label": account_limits.get("title") or "Credential pool",
|
||||
"quota": None,
|
||||
"account_limits": account_limits,
|
||||
"message": f"{display_name} credential pool status loaded.",
|
||||
}
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": provider,
|
||||
"display_name": display_name,
|
||||
"supported": False,
|
||||
"status": "unsupported",
|
||||
"quota": None,
|
||||
"message": f"Quota status is not available for {display_name}. {detail}",
|
||||
}
|
||||
|
||||
api_key = _get_provider_api_key("openrouter")
|
||||
if not api_key:
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "no_key",
|
||||
"quota": None,
|
||||
"message": "OpenRouter quota status needs an OPENROUTER_API_KEY configured on the server.",
|
||||
}
|
||||
|
||||
req = urllib.request.Request(
|
||||
_OPENROUTER_KEY_URL,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Accept": "application/json",
|
||||
},
|
||||
)
|
||||
try:
|
||||
with urllib.request.urlopen(req, timeout=_PROVIDER_QUOTA_TIMEOUT_SECONDS) as resp:
|
||||
raw = resp.read()
|
||||
payload = json.loads(raw.decode("utf-8")) if isinstance(raw, (bytes, bytearray)) else json.loads(raw)
|
||||
quota = _sanitize_openrouter_quota(payload)
|
||||
return {
|
||||
"ok": True,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "available",
|
||||
"label": "OpenRouter credits",
|
||||
"quota": quota,
|
||||
"message": "OpenRouter quota status loaded.",
|
||||
}
|
||||
except urllib.error.HTTPError as exc:
|
||||
status = "invalid_key" if exc.code in (401, 403) else "unavailable"
|
||||
message = (
|
||||
"OpenRouter rejected the configured API key."
|
||||
if status == "invalid_key"
|
||||
else "OpenRouter quota status is temporarily unavailable."
|
||||
)
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": status,
|
||||
"quota": None,
|
||||
"message": message,
|
||||
}
|
||||
except (TimeoutError, urllib.error.URLError, json.JSONDecodeError, OSError, ValueError):
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": "openrouter",
|
||||
"display_name": display_name,
|
||||
"supported": True,
|
||||
"status": "unavailable",
|
||||
"quota": None,
|
||||
"message": "OpenRouter quota status is temporarily unavailable.",
|
||||
"account_limits": account_limits,
|
||||
"message": f"{display_name} credential pool: all credentials are unavailable.",
|
||||
}
|
||||
|
||||
detail = "OpenAI/Anthropic rate-limit headers are a follow-up once WebUI captures provider response metadata."
|
||||
return {
|
||||
"ok": False,
|
||||
"provider": provider,
|
||||
"display_name": display_name,
|
||||
"supported": False,
|
||||
"status": "unsupported",
|
||||
"quota": None,
|
||||
"message": f"Quota status is not available for {display_name}. {detail}",
|
||||
}
|
||||
|
||||
|
||||
def _provider_is_oauth(provider_id: str) -> bool:
|
||||
"""Check whether a provider uses OAuth/token flows (managed by CLI)."""
|
||||
|
||||
+9
-1
@@ -7412,7 +7412,15 @@ def handle_get(handler, parsed) -> bool:
|
||||
query = parse_qs(parsed.query)
|
||||
provider_id = (query.get("provider", [""])[0] or None)
|
||||
refresh = (query.get("refresh", [""])[0] or "").strip().lower() in {"1", "true", "yes", "on"}
|
||||
return j(handler, get_provider_quota(provider_id, refresh=refresh))
|
||||
# Bind the active request's profile env (matches /api/providers and
|
||||
# /api/models/live). #4365 added a credential_pool.load_pool() path in
|
||||
# get_provider_quota for all pooled providers; without this wrapper that
|
||||
# read/write runs under the process-default profile, so a multi-profile
|
||||
# client would see (and seed) the default profile's pool instead of its
|
||||
# own (#4247/#4067 profile-isolation class).
|
||||
from api.profiles import profile_env_for_active_request
|
||||
with profile_env_for_active_request("/api/provider/quota", logger_override=logger):
|
||||
return j(handler, get_provider_quota(provider_id, refresh=refresh))
|
||||
|
||||
if parsed.path == "/api/provider/cost-history":
|
||||
query = parse_qs(parsed.query)
|
||||
|
||||
@@ -0,0 +1,399 @@
|
||||
"""Unit tests for PR #4360: Extend credential pool quota status to non-Codex providers.
|
||||
|
||||
Tests cover:
|
||||
- _local_pool_snapshot(provider) function with various pool states
|
||||
- Integration with get_provider_quota for non-allowlist providers
|
||||
- Preservation of allowlist and openrouter paths
|
||||
"""
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
from contextlib import contextmanager
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from api.providers import _local_pool_snapshot, get_provider_quota
|
||||
|
||||
|
||||
def _is_ambient_gh_cli_entry_real(source, label, key_source):
|
||||
markers = frozenset({"gh_cli", "gh auth token"})
|
||||
env_sources = frozenset({"env:github_token", "env:gh_token"})
|
||||
source_lower = source.strip().lower()
|
||||
return (
|
||||
source_lower in markers
|
||||
or source_lower in env_sources
|
||||
or label.strip().lower() == "gh auth token"
|
||||
or key_source.strip().lower() == "gh auth token"
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _fake_credential_pool(**kwargs):
|
||||
"""Context manager that ensures agent.credential_pool and api.config exist and patches load_pool."""
|
||||
had_agent = "agent" in sys.modules
|
||||
had_cp = "agent.credential_pool" in sys.modules
|
||||
had_api = "api" in sys.modules
|
||||
had_api_config = "api.config" in sys.modules
|
||||
old_agent = sys.modules.get("agent")
|
||||
|
||||
if not had_agent:
|
||||
_agent = types.ModuleType("agent")
|
||||
_agent.__path__ = []
|
||||
sys.modules["agent"] = _agent
|
||||
if not had_cp:
|
||||
_cp = types.ModuleType("agent.credential_pool")
|
||||
sys.modules["agent.credential_pool"] = _cp
|
||||
sys.modules["agent"].credential_pool = _cp
|
||||
|
||||
if not had_api:
|
||||
_api = types.ModuleType("api")
|
||||
_api.__path__ = []
|
||||
sys.modules["api"] = _api
|
||||
if not had_api_config:
|
||||
_api_config = types.ModuleType("api.config")
|
||||
_api_config._is_ambient_gh_cli_entry = _is_ambient_gh_cli_entry_real
|
||||
sys.modules["api.config"] = _api_config
|
||||
sys.modules["api"].config = _api_config
|
||||
elif not hasattr(sys.modules["api.config"], "_is_ambient_gh_cli_entry"):
|
||||
sys.modules["api.config"]._is_ambient_gh_cli_entry = _is_ambient_gh_cli_entry_real
|
||||
|
||||
cp = sys.modules["agent.credential_pool"]
|
||||
if not hasattr(cp, "load_pool"):
|
||||
cp.load_pool = None
|
||||
|
||||
with patch("agent.credential_pool.load_pool", **kwargs) as m:
|
||||
yield m
|
||||
|
||||
if not had_cp:
|
||||
sys.modules.pop("agent.credential_pool", None)
|
||||
if had_agent and old_agent is not None and hasattr(old_agent, "credential_pool"):
|
||||
delattr(old_agent, "credential_pool")
|
||||
if not had_agent:
|
||||
sys.modules.pop("agent", None)
|
||||
if not had_api_config:
|
||||
sys.modules.pop("api.config", None)
|
||||
if not had_api:
|
||||
sys.modules.pop("api", None)
|
||||
|
||||
|
||||
class TestLocalPoolSnapshot(unittest.TestCase):
|
||||
"""Tests for _local_pool_snapshot(provider) function."""
|
||||
|
||||
def test_local_pool_snapshot_returns_none_for_empty_pool(self):
|
||||
"""_local_pool_snapshot returns None when load_pool returns None."""
|
||||
with _fake_credential_pool( return_value=None):
|
||||
result = _local_pool_snapshot("some-provider")
|
||||
assert result is None
|
||||
|
||||
def test_local_pool_snapshot_returns_none_for_pool_with_no_entries(self):
|
||||
"""_local_pool_snapshot returns None when pool has no entries."""
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = []
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
result = _local_pool_snapshot("some-provider")
|
||||
assert result is None
|
||||
|
||||
def test_local_pool_snapshot_with_available_entries(self):
|
||||
"""_local_pool_snapshot returns correct snapshot with available entries."""
|
||||
entry1 = SimpleNamespace(
|
||||
label="Cred 1",
|
||||
source="test",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
entry2 = SimpleNamespace(
|
||||
label="Cred 2",
|
||||
source="test",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry1, entry2]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
result = _local_pool_snapshot("google")
|
||||
assert result is not None
|
||||
assert result.provider == "google"
|
||||
assert result.source == "local_pool"
|
||||
assert result.available is True
|
||||
assert result.pool["total_credentials"] == 2
|
||||
assert result.pool["available_credentials"] == 2
|
||||
assert result.pool["exhausted_credentials"] == 0
|
||||
assert len(result.pool["credentials"]) == 2
|
||||
assert result.pool["credentials"][0]["status"] == "available"
|
||||
assert "2/2 credentials available" in result.details
|
||||
|
||||
def test_local_pool_snapshot_with_exhausted_entries(self):
|
||||
"""_local_pool_snapshot returns correct snapshot with all exhausted entries."""
|
||||
exhausted_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
||||
entry1 = SimpleNamespace(
|
||||
label="Exhausted Cred",
|
||||
source="test",
|
||||
last_status="exhausted",
|
||||
last_status_at=exhausted_until - timedelta(hours=1),
|
||||
last_error_code="401",
|
||||
last_error_reset_at=exhausted_until,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry1]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
result = _local_pool_snapshot("openai")
|
||||
assert result is not None
|
||||
assert result.available is False
|
||||
assert result.pool["total_credentials"] == 1
|
||||
assert result.pool["available_credentials"] == 0
|
||||
assert result.pool["exhausted_credentials"] == 1
|
||||
assert result.pool["credentials"][0]["status"] == "exhausted"
|
||||
assert "All pool credentials are unavailable" in result.unavailable_reason
|
||||
assert "1 exhausted" in result.details
|
||||
|
||||
def test_local_pool_snapshot_with_mixed_entries(self):
|
||||
"""_local_pool_snapshot returns correct counts with mixed available/exhausted."""
|
||||
entry_avail = SimpleNamespace(
|
||||
label="Available",
|
||||
source="test",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
exhausted_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
||||
entry_exhausted = SimpleNamespace(
|
||||
label="Exhausted",
|
||||
source="test",
|
||||
last_status="exhausted",
|
||||
last_status_at=exhausted_until - timedelta(hours=1),
|
||||
last_error_code="429",
|
||||
last_error_reset_at=exhausted_until,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry_avail, entry_exhausted]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
result = _local_pool_snapshot("vertex")
|
||||
assert result is not None
|
||||
assert result.available is True
|
||||
assert result.pool["total_credentials"] == 2
|
||||
assert result.pool["available_credentials"] == 1
|
||||
assert result.pool["exhausted_credentials"] == 1
|
||||
assert "1/2 credentials available" in result.details
|
||||
assert "1 exhausted" in result.details
|
||||
|
||||
def test_local_pool_snapshot_with_dead_entry(self):
|
||||
"""_local_pool_snapshot treats dead credentials as unavailable, not available."""
|
||||
entry = SimpleNamespace(
|
||||
label="Dead Cred",
|
||||
source="test",
|
||||
last_status="dead",
|
||||
last_status_at=None,
|
||||
last_error_code="401",
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
result = _local_pool_snapshot("xai-oauth")
|
||||
assert result is not None
|
||||
assert result.available is False
|
||||
assert result.pool["available_credentials"] == 0
|
||||
assert result.pool["dead_credentials"] == 1
|
||||
assert result.pool["credentials"][0]["status"] == "dead"
|
||||
assert "All pool credentials are unavailable" in result.unavailable_reason
|
||||
|
||||
def test_local_pool_snapshot_handles_load_pool_exception(self):
|
||||
"""_local_pool_snapshot returns None when load_pool raises exception."""
|
||||
with _fake_credential_pool( side_effect=RuntimeError("Pool error")):
|
||||
result = _local_pool_snapshot("custom-provider")
|
||||
assert result is None
|
||||
|
||||
def test_local_pool_snapshot_snapshot_shape(self):
|
||||
"""_local_pool_snapshot returns properly shaped SimpleNamespace."""
|
||||
entry = SimpleNamespace(
|
||||
label="Test",
|
||||
source="test",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
result = _local_pool_snapshot("test-provider")
|
||||
assert hasattr(result, "provider")
|
||||
assert hasattr(result, "source")
|
||||
assert hasattr(result, "title")
|
||||
assert hasattr(result, "plan")
|
||||
assert hasattr(result, "windows")
|
||||
assert hasattr(result, "details")
|
||||
assert hasattr(result, "available")
|
||||
assert hasattr(result, "unavailable_reason")
|
||||
assert hasattr(result, "fetched_at")
|
||||
assert hasattr(result, "pool")
|
||||
assert result.title == "Credential pool"
|
||||
assert result.plan is None
|
||||
assert result.windows == ()
|
||||
assert isinstance(result.fetched_at, datetime)
|
||||
|
||||
|
||||
def test_local_pool_snapshot_filters_ambient_gh_cli_entries(self):
|
||||
"""_local_pool_snapshot skips ambient gh-cli entries so providers like
|
||||
copilot don't show a phantom "1 credential available" (#4247 class)."""
|
||||
entry = SimpleNamespace(
|
||||
label="gh auth token",
|
||||
source="gh_cli",
|
||||
key_source="gh auth token",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry]
|
||||
|
||||
with _fake_credential_pool(return_value=pool):
|
||||
result = _local_pool_snapshot("copilot")
|
||||
assert result is None
|
||||
|
||||
def test_local_pool_snapshot_keeps_non_ambient_alongside_ambient(self):
|
||||
"""Non-ambient entries are kept even when ambient ones are present."""
|
||||
ambient = SimpleNamespace(
|
||||
label="gh auth token",
|
||||
source="gh_cli",
|
||||
key_source="gh auth token",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
real = SimpleNamespace(
|
||||
label="My API Key",
|
||||
source="user",
|
||||
key_source="",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [ambient, real]
|
||||
|
||||
with _fake_credential_pool(return_value=pool):
|
||||
result = _local_pool_snapshot("copilot")
|
||||
assert result is not None
|
||||
assert result.pool["total_credentials"] == 1
|
||||
assert result.pool["available_credentials"] == 1
|
||||
assert len(result.pool["credentials"]) == 1
|
||||
|
||||
|
||||
class TestGetProviderQuotaLocalPool(unittest.TestCase):
|
||||
"""Tests for get_provider_quota integration with local pool snapshots."""
|
||||
|
||||
def test_get_provider_quota_uses_local_pool_for_non_allowlist_provider(self):
|
||||
"""get_provider_quota returns available status for non-allowlist provider with pool."""
|
||||
entry = SimpleNamespace(
|
||||
label="Test Cred",
|
||||
source="test",
|
||||
last_status="available",
|
||||
last_status_at=None,
|
||||
last_error_code=None,
|
||||
last_error_reset_at=None,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
with patch("api.providers._active_provider_id", return_value="google"):
|
||||
result = get_provider_quota("google")
|
||||
assert result["ok"] is True
|
||||
assert result["provider"] == "google"
|
||||
assert result["status"] == "available"
|
||||
assert result["supported"] is True
|
||||
assert "account_limits" in result
|
||||
assert result["account_limits"]["source"] == "local_pool"
|
||||
assert result["account_limits"]["available"] is True
|
||||
|
||||
def test_get_provider_quota_unavailable_when_all_pool_exhausted(self):
|
||||
"""get_provider_quota returns unavailable when all pool credentials exhausted."""
|
||||
exhausted_until = datetime.now(timezone.utc) + timedelta(hours=1)
|
||||
entry = SimpleNamespace(
|
||||
label="Exhausted",
|
||||
source="test",
|
||||
last_status="exhausted",
|
||||
last_status_at=exhausted_until - timedelta(hours=1),
|
||||
last_error_code="401",
|
||||
last_error_reset_at=exhausted_until,
|
||||
)
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = [entry]
|
||||
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
with patch("api.providers._active_provider_id", return_value="openai"):
|
||||
result = get_provider_quota("openai")
|
||||
assert result["ok"] is False
|
||||
assert result["provider"] == "openai"
|
||||
assert result["status"] == "unavailable"
|
||||
assert result["supported"] is True
|
||||
assert "all credentials are unavailable" in result["message"]
|
||||
assert result["account_limits"]["available"] is False
|
||||
|
||||
def test_get_provider_quota_falls_through_to_unsupported_when_no_pool(self):
|
||||
"""get_provider_quota returns unsupported when no pool exists for provider."""
|
||||
with _fake_credential_pool( return_value=None):
|
||||
with patch("api.providers._active_provider_id", return_value="some-other-provider"):
|
||||
result = get_provider_quota("some-other-provider")
|
||||
assert result["ok"] is False
|
||||
assert result["status"] == "unsupported"
|
||||
assert result["supported"] is False
|
||||
|
||||
def test_get_provider_quota_falls_through_when_pool_has_no_entries(self):
|
||||
"""get_provider_quota returns unsupported when pool exists but has no entries."""
|
||||
pool = MagicMock()
|
||||
pool.entries.return_value = []
|
||||
with _fake_credential_pool( return_value=pool):
|
||||
with patch("api.providers._active_provider_id", return_value="custom-provider"):
|
||||
result = get_provider_quota("custom-provider")
|
||||
assert result["ok"] is False
|
||||
assert result["status"] == "unsupported"
|
||||
|
||||
def test_allowlist_providers_bypass_local_pool(self):
|
||||
"""get_provider_quota uses probe path for allowlist providers, not local pool."""
|
||||
with patch("api.providers._provider_account_usage_status") as mock_probe:
|
||||
with _fake_credential_pool(return_value=None) as mock_load_pool:
|
||||
mock_probe.return_value = {
|
||||
"ok": True,
|
||||
"provider": "openai-codex",
|
||||
"status": "available",
|
||||
}
|
||||
get_provider_quota("openai-codex")
|
||||
mock_probe.assert_called_once()
|
||||
mock_load_pool.assert_not_called()
|
||||
|
||||
def test_allowlist_anthropic_uses_probe_path(self):
|
||||
"""get_provider_quota uses probe path for anthropic, not local pool."""
|
||||
with patch("api.providers._provider_account_usage_status") as mock_probe:
|
||||
with _fake_credential_pool(return_value=None) as mock_load_pool:
|
||||
mock_probe.return_value = {
|
||||
"ok": True,
|
||||
"provider": "anthropic",
|
||||
"status": "available",
|
||||
}
|
||||
get_provider_quota("anthropic")
|
||||
mock_probe.assert_called_once()
|
||||
mock_load_pool.assert_not_called()
|
||||
|
||||
def test_openrouter_uses_api_key_path_not_local_pool(self):
|
||||
"""get_provider_quota uses openrouter API-key path, not local pool."""
|
||||
with _fake_credential_pool(return_value=None) as mock_load_pool:
|
||||
with patch("api.providers._get_provider_api_key", return_value=None):
|
||||
result = get_provider_quota("openrouter")
|
||||
assert result["status"] == "no_key"
|
||||
mock_load_pool.assert_not_called()
|
||||
Reference in new issue
Block a user