From a01dd2106c6abe22be50196e21cbb11e853359be Mon Sep 17 00:00:00 2001 From: Priyanshu-u07 Date: Sun, 25 Oct 2026 19:48:48 +0530 Subject: [PATCH 1/2] Add usage history per API key and model The usage endpoint returns only today's totals per key, summed across models. Add /config/quota/usage/history: daily requests and tokens per API key and model over a date range, for the caller's org, with revoked keys included. One join through api_keys replaces a query per key. Signed-off-by: Priyanshu-u07 --- src/api_gateway/management/configuration.py | 69 +++++++++++- src/api_gateway/schemas/config.py | 15 ++- src/tests/api_gateway/test_usage_history.py | 112 ++++++++++++++++++++ 3 files changed, 194 insertions(+), 2 deletions(-) create mode 100644 src/tests/api_gateway/test_usage_history.py diff --git a/src/api_gateway/management/configuration.py b/src/api_gateway/management/configuration.py index 0e093298..3370056d 100644 --- a/src/api_gateway/management/configuration.py +++ b/src/api_gateway/management/configuration.py @@ -1,8 +1,9 @@ from fastapi import APIRouter, Depends, HTTPException, Request from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy import func from sqlalchemy.future import select from typing import List, Optional -from datetime import datetime, timezone +from datetime import date, datetime, timedelta, timezone import asyncio import logging @@ -24,6 +25,7 @@ def utcnow_naive(): ConfigUpdateRequest, ConfigResponse, UsageStatsResponse, + UsageHistoryRow, ) from api_gateway.management.dependencies import ( get_current_user_context, @@ -426,6 +428,71 @@ async def get_usage_stats(request: Request, db: AsyncSession = Depends(get_db)): return stats +# Bounds one query; a spend page asks for a month at a time. +MAX_HISTORY_DAYS = 366 + + +@router.get("/config/quota/usage/history", response_model=List[UsageHistoryRow]) +async def get_usage_history( + request: Request, + start: Optional[date] = None, + end: Optional[date] = None, + db: AsyncSession = Depends(get_db), +): + """Daily usage per API key and model in the caller's org, revoked keys included. + + Days are the dates usage was recorded under. Writes use the server's local + date (PolicyEngine.persist_usage_db), so the default range does too. + """ + user_ctx = get_current_user_context(request) + authz_service.require_permission(user_ctx, PermissionEnum.ORG_VIEW) + if not user_ctx.org_id: + return [] + + end = end or date.today() + start = start or end - timedelta(days=29) + if start > end: + raise HTTPException(status_code=400, detail="start must not be after end") + if (end - start).days >= MAX_HISTORY_DAYS: + raise HTTPException( + status_code=400, detail=f"range is limited to {MAX_HISTORY_DAYS} days" + ) + + # usage_stats has no org column; API-key usage is stored as "apikey:". + result = await db.execute( + select( + DBUsage.date, + DBApiKey.id, + DBApiKey.name, + DBApiKey.prefix, + DBApiKey.is_active, + DBUsage.model, + DBUsage.request_count, + DBUsage.prompt_tokens, + DBUsage.completion_tokens, + DBUsage.total_tokens, + ) + .join(DBApiKey, DBUsage.user_id == func.concat("apikey:", DBApiKey.id)) + .where(DBApiKey.org_id == user_ctx.org_id, DBUsage.date.between(start, end)) + .order_by(DBUsage.date, DBApiKey.name, DBUsage.model) + ) + return [ + UsageHistoryRow( + day=row.date, + key_id=row.id, + key_name=row.name, + key_prefix=row.prefix, + key_active=row.is_active is not False, + model=row.model, + requests=row.request_count or 0, + prompt_tokens=row.prompt_tokens or 0, + completion_tokens=row.completion_tokens or 0, + total_tokens=row.total_tokens or 0, + ) + for row in result + ] + + @router.get("/config/{policy_type}", response_model=ConfigResponse) async def get_config( policy_type: str, diff --git a/src/api_gateway/schemas/config.py b/src/api_gateway/schemas/config.py index 311f78e1..f1bef0e6 100644 --- a/src/api_gateway/schemas/config.py +++ b/src/api_gateway/schemas/config.py @@ -1,6 +1,6 @@ from pydantic import BaseModel, ConfigDict from typing import Optional, Dict, Any -from datetime import datetime +from datetime import date, datetime class ConfigUpdateRequest(BaseModel): policy_type: str @@ -18,3 +18,16 @@ class UsageStatsResponse(BaseModel): key_prefix: str requests: int tokens: int + +class UsageHistoryRow(BaseModel): + """One API key's usage of one model on one day.""" + day: date + key_id: str + key_name: str + key_prefix: str + key_active: bool + model: str + requests: int + prompt_tokens: int + completion_tokens: int + total_tokens: int diff --git a/src/tests/api_gateway/test_usage_history.py b/src/tests/api_gateway/test_usage_history.py new file mode 100644 index 00000000..b699db14 --- /dev/null +++ b/src/tests/api_gateway/test_usage_history.py @@ -0,0 +1,112 @@ +"""Usage history: daily usage per API key and model, scoped to the caller's org. + +The checks on the query itself need Postgres and skip without TEST_DATABASE_URL. +""" +import os +import uuid +from datetime import date, timedelta +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from api_gateway.management import configuration +from api_gateway.schemas.auth import PermissionEnum + +pytestmark = pytest.mark.asyncio + + +async def _call(org_id="org-a", **params): + """Returns (require_permission mock, result or the HTTPException raised).""" + ctx = SimpleNamespace(org_id=org_id) + require = MagicMock() + with patch.object(configuration, "get_current_user_context", return_value=ctx), \ + patch.object(configuration.authz_service, "require_permission", require): + try: + return require, await configuration.get_usage_history(MagicMock(), db=AsyncMock(), **params) + except HTTPException as e: + return require, e + + +async def test_it_requires_org_view(): + require, _ = await _call(start=date(2026, 10, 1), end=date(2026, 10, 2)) + assert require.call_args.args[1] == PermissionEnum.ORG_VIEW + + +async def test_a_user_without_an_org_gets_nothing(): + _, result = await _call(org_id=None) + assert result == [] + + +async def test_start_after_end_is_rejected(): + _, result = await _call(start=date(2026, 10, 2), end=date(2026, 10, 1)) + assert isinstance(result, HTTPException) and result.status_code == 400 + + +async def test_a_range_longer_than_the_limit_is_rejected(): + end = date(2026, 10, 25) + _, result = await _call(start=end - timedelta(days=configuration.MAX_HISTORY_DAYS), end=end) + assert isinstance(result, HTTPException) and result.status_code == 400 + + +_DSN = os.getenv("TEST_DATABASE_URL") + + +@pytest.mark.skipif(not _DSN, reason="needs TEST_DATABASE_URL") +async def test_history_is_per_key_and_model_in_the_callers_org_only(): + from sqlalchemy import text + from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine + + engine = create_async_engine(_DSN.replace("postgresql://", "postgresql+asyncpg://")) + org, other = f"org-{uuid.uuid4()}", f"org-{uuid.uuid4()}" + live, revoked, foreign = (str(uuid.uuid4()) for _ in range(3)) + day = date(2026, 10, 20) + + rows = [ + (f"apikey:{live}", day, "qwen", 3, 30, 10), + (f"apikey:{live}", day, "llama", 1, 5, 5), + (f"apikey:{revoked}", day, "qwen", 2, 20, 2), + (f"apikey:{foreign}", day, "qwen", 9, 99, 9), # another org + (f"apikey:{live}", day - timedelta(days=40), "qwen", 7, 7, 7), # outside the range + ("some-user-id", day, "qwen", 4, 4, 4), # sandbox usage, not an API key + ] + try: + async with engine.begin() as conn: + for key_id, owner, active in ((live, org, True), (revoked, org, False), (foreign, other, True)): + await conn.execute( + text("INSERT INTO api_keys (id, name, key_hash, prefix, org_id, is_active) " + "VALUES (:id, :name, :hash, 'sk-', :org, :active)"), + {"id": key_id, "name": f"key-{key_id[:4]}", "hash": key_id, "org": owner, "active": active}, + ) + for user_id, d, model, n, p, c in rows: + await conn.execute( + text("INSERT INTO usage_stats (id, user_id, date, model, request_count, " + "prompt_tokens, completion_tokens, total_tokens) " + "VALUES (:id, :u, :d, :m, :n, :p, :c, :t)"), + {"id": str(uuid.uuid4()), "u": user_id, "d": d, "m": model, + "n": n, "p": p, "c": c, "t": p + c}, + ) + + ctx = SimpleNamespace(org_id=org) + async with AsyncSession(engine) as db: + with patch.object(configuration, "get_current_user_context", return_value=ctx), \ + patch.object(configuration.authz_service, "require_permission"): + result = await configuration.get_usage_history( + MagicMock(), start=day - timedelta(days=7), end=day, db=db, + ) + + got = {(r.key_id, r.model): (r.requests, r.total_tokens, r.key_active) for r in result} + assert got == { + (live, "qwen"): (3, 40, True), + (live, "llama"): (1, 10, True), + (revoked, "qwen"): (2, 22, False), + } + assert all(r.day == day for r in result) + finally: + async with engine.begin() as conn: + await conn.execute(text("DELETE FROM usage_stats WHERE user_id = ANY(:u)"), + {"u": [r[0] for r in rows]}) + await conn.execute(text("DELETE FROM api_keys WHERE id = ANY(:k)"), + {"k": [live, revoked, foreign]}) + await engine.dispose() From e978358cf69a19f4dbb109a3c95d863b81ace4f4 Mon Sep 17 00:00:00 2001 From: Priyanshu-u07 Date: Sun, 25 Oct 2026 21:05:41 +0530 Subject: [PATCH 2/2] Move the usage history query into usage_history_rows So the history can be read for an org in-process, not only through the endpoint. The endpoint keeps its checks and calls it; behaviour is unchanged. Signed-off-by: Priyanshu-u07 --- src/api_gateway/management/configuration.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/api_gateway/management/configuration.py b/src/api_gateway/management/configuration.py index 3370056d..1e3fadbf 100644 --- a/src/api_gateway/management/configuration.py +++ b/src/api_gateway/management/configuration.py @@ -458,6 +458,13 @@ async def get_usage_history( status_code=400, detail=f"range is limited to {MAX_HISTORY_DAYS} days" ) + return await usage_history_rows(db, user_ctx.org_id, start, end) + + +async def usage_history_rows( + db: AsyncSession, org_id: str, start: date, end: date +) -> List[UsageHistoryRow]: + """Daily usage per API key and model for one org between two days, inclusive.""" # usage_stats has no org column; API-key usage is stored as "apikey:". result = await db.execute( select( @@ -473,7 +480,7 @@ async def get_usage_history( DBUsage.total_tokens, ) .join(DBApiKey, DBUsage.user_id == func.concat("apikey:", DBApiKey.id)) - .where(DBApiKey.org_id == user_ctx.org_id, DBUsage.date.between(start, end)) + .where(DBApiKey.org_id == org_id, DBUsage.date.between(start, end)) .order_by(DBUsage.date, DBApiKey.name, DBUsage.model) ) return [