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:
nesquena-hermes committed 2026-06-19 05:27:42 +00:00
1 parent 16a94ba0b3
commit 83a479d7a5
3 files changed
+641 -14

No files matched your search

+234 -14
View File
@@ -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)."""
+8
View File
@@ -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":
+399
View File
@@ -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()