[codex] Add Hugging Face billing callback

api/api_key/schemas.py

@@ -61,8 +61,10 @@ def validate_object_type(_, __, type_):
         """
         Limit which types of objects we can manipulate with API keys.
         """
-        if type_ not in ("images", "chutes", "invocations"):
-            raise ValueError("Invalid object_type, must be one of images, chutes, invocations")
+        if type_ not in ("images", "chutes", "invocations", "billing"):
+            raise ValueError(
+                "Invalid object_type, must be one of images, chutes, invocations, billing"
+            )
         return type_
 
 

api/invocation/router.py

@@ -766,6 +766,7 @@ async def _stream_with_first_chunk():
                     request,
                     {
                         "X-Chutes-InvocationID": parent_invocation_id,
+                        "Inference-Id": parent_invocation_id,
                         "Cache-Control": "no-cache, no-transform",
                         "X-Accel-Buffering": "no",
                     },
@@ -815,6 +816,7 @@ async def _streamfile():
                         request,
                         {
                             "X-Chutes-InvocationID": parent_invocation_id,
+                            "Inference-Id": parent_invocation_id,
                             "Cache-Control": "no-cache, no-transform",
                             "X-Accel-Buffering": "no",
                         },
@@ -825,7 +827,11 @@ async def _streamfile():
                     content=result["text"],
                     media_type=result["content_type"],
                     headers=build_response_headers(
-                        request, {"X-Chutes-InvocationID": parent_invocation_id}
+                        request,
+                        {
+                            "X-Chutes-InvocationID": parent_invocation_id,
+                            "Inference-Id": parent_invocation_id,
+                        },
                     ),
                 )
             else:
@@ -837,6 +843,7 @@ async def _streamfile():
                         {
                             "Content-type": "application/json",
                             "X-Chutes-InvocationID": parent_invocation_id,
+                            "Inference-Id": parent_invocation_id,
                         },
                     ),
                 )

api/main.py

@@ -365,7 +365,11 @@ async def host_router_middleware(request: Request, call_next):
 
         if request.state.auth_method != "invoke":
             # Handle /users/me/* paths specially for OAuth scope checking
-            if request.url.path.startswith("/users/me"):
+            if request.url.path == "/partners/huggingface/billing":
+                request.state.auth_method = "read"
+                request.state.auth_object_type = "billing"
+                request.state.auth_object_id = "__self__"
+            elif request.url.path.startswith("/users/me"):
                 if "/balance" in request.url.path:
                     request.state.auth_object_type = "billing"
                 elif "/quota" in request.url.path:

api/payment/router.py

@@ -3,14 +3,15 @@
 """
 
 import orjson as json
+from decimal import Decimal, InvalidOperation, ROUND_HALF_UP
 from typing import Optional
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from fastapi import APIRouter, status, HTTPException, Depends, Request
 from sqlalchemy.ext.asyncio import AsyncSession
 from api.gpu import SUPPORTED_GPUS, COMPUTE_UNIT_PRICE_BASIS, COMPUTE_MIN
 from api.config import settings
 from api.fmv.fetcher import get_fetcher
-from api.database import get_db_session
+from api.database import get_db_session, get_inv_session
 from api.user.util import refund_deposit
 from api.user.schemas import User
 from api.user.service import get_current_user
@@ -24,6 +25,34 @@ class ReturnDepositArgs(BaseModel):
     address: str
 
 
+class HuggingFaceBillingArgs(BaseModel):
+    requestIds: list[str] = Field(default_factory=list, max_length=10000)
+
+
+class HuggingFaceBillingItem(BaseModel):
+    requestId: str
+    costNanoUsd: int
+
+
+class HuggingFaceBillingResponse(BaseModel):
+    requests: list[HuggingFaceBillingItem]
+
+
+NANO_USD = Decimal("1000000000")
+
+
+def usd_to_nano_usd(value) -> int:
+    """
+    Convert a USD amount to HF's integer nano-USD billing unit.
+    """
+    try:
+        usd_amount = Decimal(str(value if value is not None else 0))
+    except (InvalidOperation, ValueError):
+        return 0
+    nano_usd = (usd_amount * NANO_USD).quantize(Decimal("1"), rounding=ROUND_HALF_UP)
+    return max(int(nano_usd), 0)
+
+
 @router.get("/daily_revenue_summary")
 async def get_daily_revenue_summary(
     days: Optional[int] = 90, db: AsyncSession = Depends(get_db_session)
@@ -145,6 +174,65 @@ async def get_pricing():
     }
 
 
+@router.post(
+    "/partners/huggingface/billing",
+    response_model=HuggingFaceBillingResponse,
+)
+async def huggingface_billing(
+    args: HuggingFaceBillingArgs,
+    current_user: User = Depends(get_current_user()),
+):
+    """
+    Hugging Face billing callback.
+
+    HF sends request IDs copied from the Inference-Id response header. Those IDs
+    are Chutes parent_invocation_id values, and the persisted metrics["b"] value
+    is the actual USD amount charged to the authenticated Chutes account.
+    Missing/unknown request IDs intentionally return 0.
+    """
+    request_ids = args.requestIds or []
+    unique_ids = list(dict.fromkeys(request_ids))
+    if not unique_ids:
+        return HuggingFaceBillingResponse(requests=[])
+
+    async with get_inv_session() as session:
+        result = await session.execute(
+            text(
+                """
+                SELECT
+                    parent_invocation_id,
+                    SUM(
+                        CASE
+                            WHEN jsonb_typeof(metrics->'b') = 'number'
+                            THEN GREATEST((metrics->>'b')::numeric, 0)
+                            ELSE 0
+                        END
+                    ) AS usd_amount
+                FROM invocations
+                WHERE user_id = :user_id
+                  AND parent_invocation_id = ANY(:request_ids)
+                GROUP BY parent_invocation_id
+                """
+            ),
+            {
+                "user_id": current_user.user_id,
+                "request_ids": unique_ids,
+            },
+        )
+        rows = result.fetchall()
+
+    costs_by_id = {row.parent_invocation_id: usd_to_nano_usd(row.usd_amount) for row in rows}
+    return HuggingFaceBillingResponse(
+        requests=[
+            HuggingFaceBillingItem(
+                requestId=request_id,
+                costNanoUsd=costs_by_id.get(request_id, 0),
+            )
+            for request_id in request_ids
+        ]
+    )
+
+
 @router.post("/return_developer_deposit")
 async def return_developer_deposit(
     args: ReturnDepositArgs,

tests/unit/test_huggingface_billing.py

@@ -0,0 +1,25 @@
+from api.api_key.schemas import APIKeyScope, Action
+from api.payment.router import usd_to_nano_usd
+
+
+def test_usd_to_nano_usd_is_deterministic():
+    assert usd_to_nano_usd("0") == 0
+    assert usd_to_nano_usd("0.000218875") == 218875
+    assert usd_to_nano_usd("1.234567891") == 1234567891
+    assert usd_to_nano_usd("0.0000000005") == 1
+    assert usd_to_nano_usd("-1") == 0
+    assert usd_to_nano_usd("not-a-number") == 0
+
+
+def test_api_keys_can_be_scoped_to_billing_read():
+    scope = APIKeyScope(object_type="billing", action=Action.READ)
+    assert scope.object_type == "billing"
+
+
+def test_unknown_api_key_scope_is_still_rejected():
+    try:
+        APIKeyScope(object_type="huggingface", action=Action.READ)
+    except ValueError as exc:
+        assert "billing" in str(exc)
+    else:
+        raise AssertionError("unknown API key scope should be rejected")