Merge pull request #233 from MorpheusAIs/feature/usage-with-api-key

feat: allow get billing info using api key

Author: alex-sandrk

src/api/v1/billing/index.py

@@ -10,7 +10,7 @@
 
 from ....db.database import get_db_session
 from ....db.models import User, LedgerEntryType
-from ....dependencies import get_current_user, get_api_key_user
+from ....dependencies import get_current_user, get_user_jwt_or_api_key
 from ....services.billing_service import billing_service
 from ....crud import credits as credits_crud
 from ....schemas.billing import (
@@ -40,10 +40,12 @@
 async def get_balance(
     request: Request,
     db: AsyncSession = Depends(get_db_session),
-    current_user: User = Depends(get_current_user),
+    current_user: User = Depends(get_user_jwt_or_api_key),
 ):
     """
     Get current credit balance for the authenticated user.
+
+    Authenticate with either a Cognito JWT or an ``sk-…`` API key.
     
     Returns:
     - paid: Paid bucket balance (posted, holds, available)
@@ -132,11 +134,13 @@ async def list_transactions(
     from_date: Optional[datetime] = Query(default=None, alias="from"),
     to_date: Optional[datetime] = Query(default=None, alias="to"),
     db: AsyncSession = Depends(get_db_session),
-    current_user: User = Depends(get_current_user),
+    current_user: User = Depends(get_user_jwt_or_api_key),
 ):
     """
     Get paginated list of credit transactions (ledger entries).
-    
+
+    Authenticate with either a Cognito JWT or an ``sk-…`` API key.
+
     Parameters:
     - limit: Maximum number of items to return (1-∞)
     - offset: Number of items to skip
@@ -226,11 +230,13 @@ async def get_monthly_spending(
     year: int = Query(default=None, description="Year for spending data (defaults to current year)"),
     mode: SpendingModeEnum = Query(default=SpendingModeEnum.gross),
     db: AsyncSession = Depends(get_db_session),
-    current_user: User = Depends(get_current_user),
+    current_user: User = Depends(get_user_jwt_or_api_key),
 ):
     """
     Get monthly spending metrics for a year.
-    
+
+    Authenticate with either a Cognito JWT or an ``sk-…`` API key.
+
     Parameters:
     - year: Year to get spending for (defaults to current year)
     - mode: 
@@ -300,11 +306,13 @@ async def list_usage(
     to_date: Optional[datetime] = Query(default=None, alias="to"),
     model: Optional[str] = Query(default=None),
     db: AsyncSession = Depends(get_db_session),
-    current_user: User = Depends(get_current_user),
+    current_user: User = Depends(get_user_jwt_or_api_key),
 ):
     """
     Get paginated list of usage entries (posted usage charges only).
-    
+
+    Authenticate with either a Cognito JWT or an ``sk-…`` API key.
+
     Parameters:
     - limit: Maximum number of items to return (1-∞)
     - offset: Number of items to skip
@@ -374,11 +382,13 @@ async def list_usage_for_month(
     limit: int = Query(default=50, ge=1),
     offset: int = Query(default=0, ge=0),
     db: AsyncSession = Depends(get_db_session),
-    current_user: User = Depends(get_current_user),
+    current_user: User = Depends(get_user_jwt_or_api_key),
 ):
     """
     Get paginated list of usage entries for a specific month.
-    
+
+    Authenticate with either a Cognito JWT or an ``sk-…`` API key.
+
     Parameters:
     - year: Year
     - month: Month (1-12)

src/dependencies.py

@@ -553,11 +553,62 @@ async def get_current_api_key(
     return auth.api_key
 
 
+# ---------------------------------------------------------------------------
+# Union auth: accept either Cognito JWT or sk-… API key
+# ---------------------------------------------------------------------------
+
+async def get_user_jwt_or_api_key(
+    db: AsyncSession = Depends(get_db_session),
+    token: Optional[HTTPAuthorizationCredentials] = Depends(oauth2_scheme_optional),
+    api_key_str: Optional[str] = Security(api_key_header),
+) -> User:
+    """
+    Authenticate via either a Cognito JWT or an ``sk-…`` API key and return
+    the associated :class:`User`.
+
+    Both schemes are carried in the ``Authorization`` header; the scheme is
+    selected purely by the credential prefix:
+
+    * value starts with ``sk-`` → treated as API key, delegated to
+      :func:`get_api_key_auth`
+    * anything else → treated as a Cognito JWT, delegated to
+      :func:`get_current_user`
+
+    Intended for read-only endpoints (e.g. billing GETs) that should be
+    reachable from both the dashboard (JWT) and programmatic clients (key).
+    """
+    # Local testing bypass mirrors the underlying dependencies.
+    from src.core.local_testing import is_local_testing_mode, get_or_create_test_user
+    if is_local_testing_mode():
+        return await get_or_create_test_user(db)
+
+    # Both HTTPBearer and APIKeyHeader read the same Authorization header.
+    # Prefer the HTTPBearer-parsed credential (already stripped of "Bearer ");
+    # fall back to the raw header for clients that omit the scheme prefix.
+    raw = token.credentials if token else (api_key_str or "")
+    if raw.startswith("Bearer "):
+        raw = raw[7:]
+
+    if not raw:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Not authenticated",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    if raw.startswith("sk-"):
+        auth = await get_api_key_auth(api_key_str=raw)
+        return auth.user
+
+    return await get_current_user(db=db, token=token)
+
+
 # ---------------------------------------------------------------------------
 # Type aliases for commonly used dependency chains
 # ---------------------------------------------------------------------------
 CurrentUser = Annotated[User, Depends(get_current_user)]
 APIKeyUser = Annotated[User, Depends(get_api_key_user)]
 CurrentAPIKey = Annotated[APIKey, Depends(get_current_api_key)]
 APIKeyAuthentication = Annotated[APIKeyAuth, Depends(get_api_key_auth)]
+JwtOrApiKeyUser = Annotated[User, Depends(get_user_jwt_or_api_key)]
 DBSession = Annotated[AsyncSession, Depends(get_db_session)]