PL-135277: log auth token identifier in proxy

src/skvaider/auth.py

@@ -7,7 +7,7 @@
 
 import svcs
 from argon2 import PasswordHasher
-from fastapi import Depends, HTTPException
+from fastapi import Depends, HTTPException, Request
 from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
 
 import aramaki
@@ -55,6 +55,7 @@ def add(self, key: str):
 
 
 async def verify_token(
+    request: Request,
     credentials: Annotated[HTTPAuthorizationCredentials, Depends(_bearer_auth)],
     services: svcs.fastapi.DepContainer,
 ) -> None:
@@ -67,6 +68,7 @@ async def verify_token(
         pass
     else:
         if credentials.credentials in admin_tokens.tokens:
+            request.state.token_id = "admin-token"
             return
 
     # XXX There's a lot of type issues going on here, because the mechanics of passing through
@@ -93,13 +95,15 @@ async def verify_token(
         try:
             assert isinstance(db_token["secret_hash"], str)
             hasher.verify(db_token["secret_hash"], client_token["secret"])
+            request.state.token_id = client_token["id"]
             cache.add(credentials.credentials)
         # We could specify explicit exceptions here but go the safe route and just catch all in case the lib addes one
         except Exception:
             raise HTTPException(401, detail="Bad authentication")
 
 
 async def verify_admin_token(
+    request: Request,
     credentials: Annotated[HTTPAuthorizationCredentials, Depends(_bearer_auth)],
     services: svcs.fastapi.DepContainer,
 ) -> None:
@@ -110,3 +114,4 @@ async def verify_admin_token(
         raise HTTPException(401, detail="No admin tokens configured")
     if credentials.credentials not in admin_tokens.tokens:
         raise HTTPException(401, detail="Bad authentication")
+    request.state.token_id = "admin-token"

src/skvaider/conftest.py

@@ -238,6 +238,7 @@ def _create(model: str = "test-model", stream: bool = False) -> MagicMock:
         req.state = MagicMock()
         req.state.model = model
         req.state.stream = stream
+        req.state.token_id = "test-token"
         req.headers.get.return_value = None
         return req
 

src/skvaider/logging.py

@@ -87,6 +87,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
         request.state.parallel_backend = 0
         request.state.backend_endpoint = ""
         request.state.response_headers = {}
+        request.state.token_id = "n/a"
 
         try:
             await self.app(
@@ -118,7 +119,7 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send):
                 p = request.state
                 # XXX when streaming: add time to time first token
                 self._logger.info(
-                    f'{anon_ip} {p.model} {backend} {p.request_id} {time_queue}/{time_server}/{process_time} {p.status_code} {p.tokens_prompt}/{p.tokens_completion}/{p.tokens_total} {stream}{debug_flag} {p.retries} {p.parallel_total}/{p.parallel_model}/{p.parallel_backend} "{request.method} {request.url.path}" '
+                    f'{anon_ip} {p.token_id} {p.model} {backend} {p.request_id} {time_queue}/{time_server}/{process_time} {p.status_code} {p.tokens_prompt}/{p.tokens_completion}/{p.tokens_total} {stream}{debug_flag} {p.retries} {p.parallel_total}/{p.parallel_model}/{p.parallel_backend} "{request.method} {request.url.path}" '
                 )
 
 

src/skvaider/tests/test_auth.py

@@ -1,10 +1,12 @@
 import base64
 import json
+from typing import Callable
 
 import fastapi.exceptions
 import pytest
 import svcs
 from argon2 import PasswordHasher
+from fastapi import Request
 from fastapi.security import HTTPAuthorizationCredentials
 
 from skvaider.auth import verify_token
@@ -13,16 +15,22 @@
 hasher = PasswordHasher()
 
 
-async def test_verify_token_incorrect_syntax(services: svcs.Container):
+async def test_verify_token_incorrect_syntax(
+    services: svcs.Container, mock_request_factory: Callable[..., Request]
+):
+    request = mock_request_factory()
     credentials = HTTPAuthorizationCredentials(
         scheme="Bearer", credentials="unknown"
     )
     with pytest.raises(fastapi.exceptions.HTTPException) as e:
-        await verify_token(credentials, services)
+        await verify_token(request, credentials, services)
     assert e.value.status_code == 401
 
 
-async def test_verify_token_unknown_user(services: svcs.Container):
+async def test_verify_token_unknown_user(
+    services: svcs.Container, mock_request_factory: Callable[..., Request]
+):
+    request = mock_request_factory()
     secret = "asdf"
     auth_token = base64.b64encode(
         json.dumps({"id": "user", "secret": secret}).encode("utf-8")
@@ -31,13 +39,16 @@ async def test_verify_token_unknown_user(services: svcs.Container):
         scheme="Bearer", credentials=auth_token
     )
     with pytest.raises(fastapi.exceptions.HTTPException) as e:
-        await verify_token(credentials, services)
+        await verify_token(request, credentials, services)
     assert e.value.status_code == 401
 
 
 async def test_verify_token_incorrect_password(
-    services: svcs.Container, token_db: DummyTokens
+    services: svcs.Container,
+    token_db: DummyTokens,
+    mock_request_factory: Callable[..., Request],
 ):
+    request = mock_request_factory()
     secret = "asdf"
     token_db.data["user"] = {"secret_hash": hasher.hash(secret)}
 
@@ -49,13 +60,16 @@ async def test_verify_token_incorrect_password(
     )
 
     with pytest.raises(fastapi.exceptions.HTTPException) as e:
-        await verify_token(credentials, services)
+        await verify_token(request, credentials, services)
     assert e.value.status_code == 401
 
 
 async def test_verify_token_correct_user_and_password(
-    services: svcs.Container, token_db: DummyTokens
+    services: svcs.Container,
+    token_db: DummyTokens,
+    mock_request_factory: Callable[..., Request],
 ):
+    request = mock_request_factory()
     secret = "asdf"
     token_db.data["user"] = {"secret_hash": hasher.hash(secret)}
 
@@ -65,4 +79,5 @@ async def test_verify_token_correct_user_and_password(
     credentials = HTTPAuthorizationCredentials(
         scheme="Bearer", credentials=auth_token
     )
-    await verify_token(credentials, services)
+    await verify_token(request, credentials, services)
+    assert request.state.token_id == "user"