fix(gateway): /usage billing route follows the most recent model too

`_persisted_billing_route` (idle `/usage` account-limits lookup) was the last
reader of the lifetime-dominant route, so it queried the retired provider's
account after a switch. Point it at `get_recent_session_model_route` and
delete the dominant query, which no longer has a caller.
This commit is contained in:
teknium1
2026-09-13 12:47:06 -07:00
committed by Teknium
parent 1387c64cac
commit f26335752b
3 changed files with 10 additions and 28 deletions
+4 -4
View File
@@ -612,14 +612,14 @@ class GatewayStatusCommandsMixin:
return t("gateway.usage.no_data")
async def _persisted_billing_route(self, source):
"""``(provider, base_url)`` from the SessionDB row / dominant route when no agent is resident."""
"""``(provider, base_url)`` from the SessionDB row / most recent route when no agent is resident."""
async def _rows():
entry = await self.async_session_store.get_or_create_session(source)
persisted = await self._session_db.get_session(entry.session_id) or {}
route = await self._session_db.get_dominant_session_model_route(entry.session_id)
route = await self._session_db.get_recent_session_model_route(entry.session_id)
return persisted, route if isinstance(route, dict) else {}
persisted, dominant = await _quiet(_rows, ({}, {}))
row = dominant if dominant.get("billing_provider") else persisted
persisted, recent = await _quiet(_rows, ({}, {}))
row = recent if recent.get("billing_provider") else persisted
return row.get("billing_provider"), row.get("billing_base_url")
async def _handle_insights_command(self, event: MessageEvent) -> str:
+4 -22
View File
@@ -744,29 +744,11 @@ class SessionSessionsMixin:
)
return self._session_row_dict(row) if row else None
def get_dominant_session_model_route(self, session_id: str) -> Optional[Dict[str, Any]]:
"""Main-loop model route that served most API calls (``session_model_usage`` keeps the coherent
per-call tuple; ``sessions`` mixes route changes)."""
self.flush_token_counts()
row = self._read_one(
"""SELECT model, billing_provider, billing_base_url, billing_mode,
api_call_count
FROM session_model_usage
WHERE session_id = ?
AND task = ''
AND model <> 'unknown'
AND billing_provider <> ''
ORDER BY api_call_count DESC,
(input_tokens + output_tokens + cache_read_tokens +
cache_write_tokens + reasoning_tokens) DESC,
last_seen DESC
LIMIT 1""",
(session_id,),
)
return dict(row) if row else None
def get_recent_session_model_route(self, session_id: str) -> Optional[Dict[str, Any]]:
"""Most recently used main-loop model route as one coherent per-call tuple."""
"""Most recently used main-loop model route as one coherent per-call tuple
(``session_model_usage`` keeps model+provider together; ``sessions`` mixes route changes).
Recency, not lifetime call count: on a long session a route retired weeks ago can hold the
highest ``api_call_count`` forever, and /status and /usage would keep calling it current."""
self.flush_token_counts()
row = self._read_one(
"""SELECT model, billing_provider, billing_base_url, billing_mode,
+2 -2
View File
@@ -163,14 +163,14 @@ class TestUsageAccountSection:
assert "📈 **Account limits**" in result
@pytest.mark.asyncio
async def test_usage_command_prefers_dominant_persisted_route(self, monkeypatch):
async def test_usage_command_prefers_recent_persisted_route(self, monkeypatch):
runner = _make_runner(SK)
runner._session_db = AsyncSessionDB(MagicMock())
runner._session_db._db.get_session.return_value = {
"billing_provider": "nous",
"billing_base_url": "https://inference-api.nousresearch.com/v1/",
}
runner._session_db._db.get_dominant_session_model_route.return_value = {
runner._session_db._db.get_recent_session_model_route.return_value = {
"model": "z-ai/glm-5.2",
"billing_provider": "nvidia",
"billing_base_url": "https://integrate.api.nvidia.com/v1/",