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
+641
-14
No files matched your search
+234
-14
@@ -679,6 +679,7 @@ else:
|
|||||||
print(json.dumps(_fetch_snapshot(provider, api_key)), flush=True)
|
print(json.dumps(_fetch_snapshot(provider, api_key)), flush=True)
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
|
||||||
# SECTION: Provider ↔ env var mapping
|
# SECTION: Provider ↔ env var mapping
|
||||||
|
|
||||||
# Maps canonical provider slug → env var name for API key.
|
# Maps canonical provider slug → env var name for API key.
|
||||||
@@ -751,7 +752,201 @@ _OAUTH_PROVIDERS = frozenset({
|
|||||||
"xai-oauth",
|
"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:
|
def _get_hermes_home() -> Path:
|
||||||
@@ -1721,18 +1916,7 @@ def get_provider_quota(provider_id: str | None = None, *, refresh: bool = False)
|
|||||||
if provider in _ACCOUNT_USAGE_PROVIDERS:
|
if provider in _ACCOUNT_USAGE_PROVIDERS:
|
||||||
return _provider_account_usage_status(provider, display_name, refresh=refresh)
|
return _provider_account_usage_status(provider, display_name, refresh=refresh)
|
||||||
|
|
||||||
if provider != "openrouter":
|
if provider == "openrouter":
|
||||||
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}",
|
|
||||||
}
|
|
||||||
|
|
||||||
api_key = _get_provider_api_key("openrouter")
|
api_key = _get_provider_api_key("openrouter")
|
||||||
if not api_key:
|
if not api_key:
|
||||||
return {
|
return {
|
||||||
@@ -1744,7 +1928,6 @@ def get_provider_quota(provider_id: str | None = None, *, refresh: bool = False)
|
|||||||
"quota": None,
|
"quota": None,
|
||||||
"message": "OpenRouter quota status needs an OPENROUTER_API_KEY configured on the server.",
|
"message": "OpenRouter quota status needs an OPENROUTER_API_KEY configured on the server.",
|
||||||
}
|
}
|
||||||
|
|
||||||
req = urllib.request.Request(
|
req = urllib.request.Request(
|
||||||
_OPENROUTER_KEY_URL,
|
_OPENROUTER_KEY_URL,
|
||||||
headers={
|
headers={
|
||||||
@@ -1794,6 +1977,43 @@ def get_provider_quota(provider_id: str | None = None, *, refresh: bool = False)
|
|||||||
"message": "OpenRouter quota status is temporarily unavailable.",
|
"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": True,
|
||||||
|
"status": "unavailable",
|
||||||
|
"quota": None,
|
||||||
|
"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:
|
def _provider_is_oauth(provider_id: str) -> bool:
|
||||||
"""Check whether a provider uses OAuth/token flows (managed by CLI)."""
|
"""Check whether a provider uses OAuth/token flows (managed by CLI)."""
|
||||||
|
|||||||
@@ -7412,6 +7412,14 @@ def handle_get(handler, parsed) -> bool:
|
|||||||
query = parse_qs(parsed.query)
|
query = parse_qs(parsed.query)
|
||||||
provider_id = (query.get("provider", [""])[0] or None)
|
provider_id = (query.get("provider", [""])[0] or None)
|
||||||
refresh = (query.get("refresh", [""])[0] or "").strip().lower() in {"1", "true", "yes", "on"}
|
refresh = (query.get("refresh", [""])[0] or "").strip().lower() in {"1", "true", "yes", "on"}
|
||||||
|
# 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))
|
return j(handler, get_provider_quota(provider_id, refresh=refresh))
|
||||||
|
|
||||||
if parsed.path == "/api/provider/cost-history":
|
if parsed.path == "/api/provider/cost-history":
|
||||||
|
|||||||
@@ -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