Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions .fork/features/umans.md
Original file line number Diff line number Diff line change
Expand Up @@ -76,3 +76,24 @@ Verification:
Notes:
- Rebased onto current `origin/dev` (fe3ee86) to fix the stale base that
caused PR #62 to show 16 commits instead of 1.

## 2026-06-22 — Fix Umans quota fetch 404 and Web UI visibility (post #62)

Target: `fix(umans): normalize API base and show quota without proxy requests`
Files:
- `src/rotator_library/providers/utilities/umans_quota_tracker.py`
- `src/rotator_library/providers/umans_provider.py`
- `src/rotator_library/client/quota.py`
- `tests/test_umans_quota_tracker.py`

Changes:
- `_normalize_umans_api_base()` strips trailing `/v1` so `UMANS_API_BASE=https://api.code.umans.ai/v1` does not call `/v1/v1/usage` (404).
- `_resolve_umans_api_key()` resolves `env://umans/N` to `UMANS_API_KEY_N` for Bearer auth.
- `QuotaService.get_quota_stats` keeps providers with quota baselines when `total_requests == 0`.
- Tests for URL normalization, env key resolution, and ISO timestamp assertion.

Verification:
- `uv run pytest tests/test_umans_quota_tracker.py -q` — 36 passed
- Live: `GET https://api.code.umans.ai/v1/usage` returns 200; double `/v1` returns 404.

Follow-up PR after merge of #62.
19 changes: 18 additions & 1 deletion src/rotator_library/client/quota.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,23 @@
lib_logger = logging.getLogger("rotator_library")


def _stats_has_quota_data(stats: Dict[str, Any]) -> bool:
"""True if provider has quota baselines or per-credential group windows (e.g. Umans API-only)."""
quota_groups = stats.get("quota_groups") or {}
if quota_groups:
for group in quota_groups.values():
windows = group.get("windows") or {}
for win in windows.values():
if (win.get("total_max") or 0) > 0:
return True
for cred in (stats.get("credentials") or {}).values():
for group in (cred.get("group_usage") or {}).values():
for win in (group.get("windows") or {}).values():
if win.get("limit") is not None:

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SUGGESTION: Inconsistent threshold vs. the provider-level branch above.

Line 21 checks (win.get("total_max") or 0) > 0, but line 26 checks win.get("limit") is not None. A credential window with limit == 0 (typically meaning "no quota" or "untracked") would still pass and cause the provider to be included in the Web UI quota view with no displayable bar.

Consider using the same > 0 semantics for consistency:

Suggested change
if win.get("limit") is not None:
if (win.get("limit") or 0) > 0:

Reply with @kilocode-bot fix it to have Kilo Code address this issue.

return True
return False


class QuotaService:
"""Aggregate usage stats and force-refresh provider quota baselines."""

Expand Down Expand Up @@ -56,7 +73,7 @@ async def get_quota_stats(
if classifier is not None:
stats["classifier"] = classifier

if stats.get("total_requests", 0) == 0:
if stats.get("total_requests", 0) == 0 and not _stats_has_quota_data(stats):
continue

providers[manager_key if classifier is not None else provider_name] = stats
Expand Down
6 changes: 4 additions & 2 deletions src/rotator_library/providers/umans_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
import httpx

from .provider_interface import ProviderInterface, UsageResetConfigDef
from .utilities.umans_quota_tracker import UmansQuotaTracker
from .utilities.umans_quota_tracker import UmansQuotaTracker, _normalize_umans_api_base

lib_logger = logging.getLogger("rotator_library")

Expand Down Expand Up @@ -139,7 +139,9 @@ async def get_models(self, api_key: str, client: httpx.AsyncClient) -> List[str]
List of model names prefixed with 'umans/'
"""
try:
base = os.getenv("UMANS_API_BASE", "https://api.code.umans.ai").rstrip("/")
base = _normalize_umans_api_base(
os.getenv("UMANS_API_BASE", "https://api.code.umans.ai")
)
response = await client.get(
f"{base}/v1/models",
headers={"Authorization": f"Bearer {api_key}"},
Expand Down
40 changes: 38 additions & 2 deletions src/rotator_library/providers/utilities/umans_quota_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,39 @@ def _get_credential_identifier(credential: str) -> str:
return f"{credential[:4]}...{credential[-4:]}"


def _normalize_umans_api_base(raw: str) -> str:
"""
Host root for Umans API paths (/v1/usage, /v1/models).

Docs often set UMANS_API_BASE to https://api.code.umans.ai/v1 for LiteLLM;
appending /v1/usage again would hit /v1/v1/usage (404).
"""
base = (raw or UMANS_API_BASE_DEFAULT).strip().rstrip("/")
if base.endswith("/v1"):
base = base[:-3].rstrip("/")
return base or UMANS_API_BASE_DEFAULT


def _resolve_umans_api_key(credential_path: str) -> str:
"""Raw Bearer token for env://umans/N virtual paths."""
if not credential_path.startswith("env://"):
return credential_path
parts = credential_path[6:].split("/")
if len(parts) < 2:
return credential_path
provider, index_s = parts[0], parts[1]
if provider != "umans":
return credential_path
idx = _safe_int(index_s, 0)
if idx <= 0:
key = os.getenv("UMANS_API_KEY", "").strip()
if key:
return key
return credential_path
key = os.getenv(f"UMANS_API_KEY_{idx}", "").strip()
return key or credential_path

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

SUGGESTION: Silent fallback to the raw credential_path will produce a confusing 401 Unauthorized from the Umans API when the corresponding UMANS_API_KEY_<idx> env var is unset.

A lib_logger.warning(...) here would surface the misconfiguration directly in logs ("Umans env://umans/N has no UMANS_API_KEY_N set") instead of forcing users to debug an opaque 401 in _fetch_usage_for_credential's except branch.


Reply with @kilocode-bot fix it to have Kilo Code address this issue.



def _parse_iso_to_unix(ts_str: Optional[str]) -> Optional[float]:
"""Parse an ISO 8601 timestamp to a Unix timestamp."""
if not ts_str:
Expand Down Expand Up @@ -267,7 +300,9 @@ def set_usage_manager(self, usage_manager: "UsageManager") -> None:
self._usage_manager = usage_manager

def _resolve_api_base(self) -> str:
return os.getenv("UMANS_API_BASE", UMANS_API_BASE_DEFAULT).rstrip("/")
return _normalize_umans_api_base(
os.getenv("UMANS_API_BASE", UMANS_API_BASE_DEFAULT)
)

async def _fetch_usage_for_credential(
self, credential_path: str
Expand All @@ -282,9 +317,10 @@ async def _fetch_usage_for_credential(
UmansQuotaSnapshot with status "success" or "error".
"""
identifier = _get_credential_identifier(credential_path)
api_key = _resolve_umans_api_key(credential_path)
try:
headers = {
"Authorization": f"Bearer {credential_path}",
"Authorization": f"Bearer {api_key}",
"Accept": "application/json",
}
base = self._resolve_api_base()
Expand Down
57 changes: 55 additions & 2 deletions tests/test_umans_quota_tracker.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,11 @@
UmansQuotaTracker,
_detect_plan,
_get_credential_identifier,
_normalize_umans_api_base,
_parse_iso_to_unix,
_parse_usage_response,
_resolve_request_limit,
_resolve_umans_api_key,
_safe_int,
)

Expand Down Expand Up @@ -105,11 +107,26 @@ def test_get_credential_identifier_short_key_unmasked():
assert _get_credential_identifier("abcd") == "abcd"


def test_normalize_umans_api_base_strips_trailing_v1():
assert (
_normalize_umans_api_base("https://api.code.umans.ai/v1")
== "https://api.code.umans.ai"
)
assert _normalize_umans_api_base("https://api.code.umans.ai") == (
"https://api.code.umans.ai"
)


def test_resolve_umans_api_key_env_virtual_path():
with patch.dict(os.environ, {"UMANS_API_KEY_1": "secret-key"}, clear=False):
assert _resolve_umans_api_key("env://umans/1") == "secret-key"
assert _resolve_umans_api_key("sk-raw") == "sk-raw"


def test_parse_iso_to_unix_z():
ts = _parse_iso_to_unix("2026-06-22T05:41:43Z")
assert ts is not None
# 2026-06-22 05:41:43 UTC is after the current test run epoch
assert ts > time.time()
assert abs(ts - 1782106903.0) < 1.0


def test_detect_plan_code_pro_inferred():
Expand Down Expand Up @@ -268,6 +285,42 @@ async def _run():
asyncio.run(_run())


def test_fetch_usage_uses_normalized_base_url():
async def _run():
host = _TrackerHost()
captured = {}

async def fake_get(url, headers=None):
captured["url"] = url
captured["auth"] = headers.get("Authorization")
resp = MagicMock()
resp.raise_for_status = MagicMock()
resp.json = MagicMock(return_value=SAMPLE_CODE_PRO)
return resp

mock_client = MagicMock()
mock_client.get = AsyncMock(side_effect=fake_get)
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock(return_value=None)

with patch.dict(
os.environ,
{"UMANS_API_BASE": "https://api.code.umans.ai/v1"},
clear=False,
):
with patch(
"rotator_library.providers.utilities.umans_quota_tracker.httpx.AsyncClient",
return_value=mock_client,
):
snap = await host._fetch_usage_for_credential("sk-testkey")

assert captured["url"] == "https://api.code.umans.ai/v1/usage"
assert captured["auth"] == "Bearer sk-testkey"
assert snap.status == "success"

asyncio.run(_run())


def test_provider_get_model_quota_group():
provider = object.__new__(UmansProvider)
assert provider.get_model_quota_group("umans/kimi-k2.7") == "5h-requests"
Expand Down
Loading