feat(auth): capture per-credential identity for usage attribution (#13606)

backend/onyx/auth/users.py

@@ -116,7 +116,7 @@ from onyx.configs.constants import (
     MilestoneRecordType,
     OnyxRedisLocks,
 )
-from onyx.db.api_key import fetch_user_for_api_key
+from onyx.db.api_key import fetch_api_key_auth_result
 from onyx.db.auth import (
     SQLAlchemyUserAdminDB,
     get_access_token_db,
@@ -132,7 +132,7 @@ from onyx.db.engine.sql_engine import (
     get_session_with_current_tenant,
     get_session_with_tenant,
 )
-from onyx.db.enums import AccountType, Permission
+from onyx.db.enums import AccountType, PatType, Permission
 from onyx.db.models import AccessToken, OAuthAccount, Persona, User
 from onyx.db.pat import resolve_pat
 from onyx.db.users import (
@@ -168,9 +168,12 @@ from shared_configs.configs import (
 )
 from shared_configs.contextvars import (
     CURRENT_TENANT_ID_CONTEXTVAR,
+    CURRENT_USAGE_CREDENTIAL_CONTEXTVAR,
     CURRENT_USER_ID_CONTEXTVAR,
+    UsageCredentialIdentity,
     get_current_tenant_id,
 )
+from shared_configs.enums import UsageCredentialType
 
 logger = setup_logger()
 
@@ -1943,6 +1946,10 @@ async def _check_for_saml_and_jwt(
                 user = await _get_or_create_user_from_jwt(
                     payload, request, async_db_session
                 )
+                if user is not None:
+                    request.state.usage_credential = UsageCredentialIdentity(
+                        UsageCredentialType.JWT
+                    )
 
     return user
 
@@ -2012,6 +2019,11 @@ async def _resolve_optional_user(
     user: User | None,
     user_manager: BaseUserManager[User, uuid.UUID],
 ) -> User | None:
+    if user is not None:
+        request.state.usage_credential = UsageCredentialIdentity(
+            UsageCredentialType.SESSION
+        )
+
     if user := await _check_for_saml_and_jwt(request, user, async_db_session):
         # If user is already set, _check_for_saml_and_jwt returns the same user object
         await _maybe_refresh_oauth_tokens(user, async_db_session, user_manager)
@@ -2025,8 +2037,26 @@ async def _resolve_optional_user(
                 # Expose the token's scopes so require_permission can cap the
                 # request to them.
                 request.state.token_scopes = pat.scopes
+                request.state.usage_credential = UsageCredentialIdentity(
+                    (
+                        UsageCredentialType.CRAFT_PAT
+                        if pat.pat_type == PatType.CRAFT
+                        else UsageCredentialType.PAT
+                    ),
+                    str(pat.pat_id),
+                    pat.pat_name,
+                    pat.pat_display,
+                )
         elif hashed_api_key := get_hashed_api_key_from_request(request):
-            user = await fetch_user_for_api_key(hashed_api_key, async_db_session)
+            api_key = await fetch_api_key_auth_result(hashed_api_key, async_db_session)
+            if api_key is not None:
+                user = api_key.user
+                request.state.usage_credential = UsageCredentialIdentity(
+                    UsageCredentialType.API_KEY,
+                    str(api_key.api_key_id),
+                    api_key.api_key_name,
+                    api_key.api_key_display,
+                )
     except ValueError:
         logger.warning("Issue with validating authentication token")
         return None
@@ -2060,9 +2090,13 @@ async def optional_user(
         user_manager,
     )
     token = CURRENT_USER_ID_CONTEXTVAR.set(str(user.id) if user is not None else None)
+    credential_token = CURRENT_USAGE_CREDENTIAL_CONTEXTVAR.set(
+        getattr(request.state, "usage_credential", None)
+    )
     try:
         yield user
     finally:
+        CURRENT_USAGE_CREDENTIAL_CONTEXTVAR.reset(credential_token)
         CURRENT_USER_ID_CONTEXTVAR.reset(token)
 
 

backend/onyx/db/api_key.py

@@ -1,9 +1,10 @@
 import uuid
+from typing import NamedTuple
 
 from fastapi_users.password import PasswordHelper
 from sqlalchemy import delete, select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import Session, joinedload
+from sqlalchemy.orm import Session, contains_eager, joinedload
 
 from onyx.auth.api_key import (
     ApiKeyDescriptor,
@@ -28,6 +29,13 @@ from shared_configs.contextvars import get_current_tenant_id
 logger = setup_logger()
 
 
+class ApiKeyAuthResult(NamedTuple):
+    user: User
+    api_key_id: int
+    api_key_name: str | None
+    api_key_display: str
+
+
 def get_api_key_email_pattern() -> str:
     return DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN
 
@@ -66,6 +74,33 @@ async def fetch_user_for_api_key(
     )
 
 
+async def fetch_api_key_auth_result(
+    hashed_api_key: str, async_db_session: AsyncSession
+) -> ApiKeyAuthResult | None:
+    row = (
+        (
+            await async_db_session.execute(
+                select(ApiKey)
+                .join(ApiKey.user)
+                # a lazy row.user under an AsyncSession raises MissingGreenlet
+                .options(contains_eager(ApiKey.user))
+                .where(ApiKey.hashed_api_key == hashed_api_key)
+            )
+        )
+        .scalars()
+        .unique()
+        .one_or_none()
+    )
+    if row is None:
+        return None
+    return ApiKeyAuthResult(
+        user=row.user,
+        api_key_id=row.id,
+        api_key_name=row.name,
+        api_key_display=row.api_key_display,
+    )
+
+
 def get_api_key_fake_email(
     name: str,
     unique_id: str,

backend/onyx/db/pat.py

@@ -30,6 +30,10 @@ class PatAuthResult(NamedTuple):
 
     user: User
     scopes: list[Permission] | None
+    pat_id: int
+    pat_name: str
+    pat_display: str
+    pat_type: PatType
 
 
 async def resolve_pat(
@@ -70,7 +74,14 @@ async def resolve_pat(
     _schedule_pat_last_used_update(hashed_token, now)
     # None (no scopes) = unrestricted; a stored list is parsed to Permissions.
     scopes = parse_permission_values(pat.scopes) if pat.scopes is not None else None
-    return PatAuthResult(user=pat.user, scopes=scopes)
+    return PatAuthResult(
+        user=pat.user,
+        scopes=scopes,
+        pat_id=pat.id,
+        pat_name=pat.name,
+        pat_display=pat.token_display,
+        pat_type=pat.pat_type,
+    )
 
 
 def _schedule_pat_last_used_update(hashed_token: str, now: datetime) -> None:

backend/shared_configs/contextvars.py

@@ -1,6 +1,8 @@
 import contextvars
+from typing import NamedTuple
 
 from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA
+from shared_configs.enums import UsageCredentialType
 
 # Context variable for the current tenant id
 CURRENT_TENANT_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = (
@@ -35,6 +37,18 @@ CURRENT_USER_ID_CONTEXTVAR: contextvars.ContextVar[str | None] = contextvars.Con
 )
 
 
+class UsageCredentialIdentity(NamedTuple):
+    credential_type: UsageCredentialType
+    credential_id: str | None = None
+    credential_name: str | None = None
+    credential_display: str | None = None
+
+
+CURRENT_USAGE_CREDENTIAL_CONTEXTVAR: contextvars.ContextVar[
+    UsageCredentialIdentity | None
+] = contextvars.ContextVar("current_usage_credential", default=None)
+
+
 def get_current_tenant_id() -> str:
     tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
     if tenant_id is None:
@@ -55,3 +69,7 @@ def get_current_tenant_id() -> str:
 def get_current_user_id() -> str | None:
     """Requesting user's id, or None outside a per-request context."""
     return CURRENT_USER_ID_CONTEXTVAR.get()
+
+
+def get_current_usage_credential() -> UsageCredentialIdentity | None:
+    return CURRENT_USAGE_CREDENTIAL_CONTEXTVAR.get()

backend/shared_configs/enums.py

@@ -40,3 +40,14 @@ class WebContentProviderType(str, Enum):
 class TracingProviderType(str, Enum):
     BRAINTRUST = "braintrust"
     LANGFUSE = "langfuse"
+
+
+class UsageCredentialType(str, Enum):
+    """CRAFT_PAT is separate from PAT because Craft sandbox tokens are
+    agent-driven yet carry the owning human's user_id."""
+
+    SESSION = "session"
+    JWT = "jwt"
+    PAT = "pat"
+    CRAFT_PAT = "craft_pat"
+    API_KEY = "api_key"

backend/tests/external_dependency_unit/db/test_api_key_auth_result.py

@@ -0,0 +1,39 @@
+import pytest
+from sqlalchemy.orm import Session
+
+from onyx.auth.api_key import hash_api_key
+from onyx.auth.schemas import UserRole
+from onyx.db.api_key import fetch_api_key_auth_result, insert_api_key, remove_api_key
+from onyx.db.engine.async_sql_engine import get_async_session_context_manager
+from onyx.server.api_key.models import APIKeyArgs
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("tenant_context")
+async def test_fetch_api_key_auth_result_loads_user(db_session: Session) -> None:
+    descriptor = insert_api_key(
+        db_session, APIKeyArgs(name="auth-result", role=UserRole.BASIC), user_id=None
+    )
+    assert descriptor.api_key is not None
+
+    try:
+        async with get_async_session_context_manager() as async_session:
+            result = await fetch_api_key_auth_result(
+                hash_api_key(descriptor.api_key), async_session
+            )
+
+        assert result is not None
+        # Accessed outside the async session: only possible if eagerly loaded.
+        assert result.user.id == descriptor.user_id
+        assert result.user.role == UserRole.BASIC
+        assert result.api_key_id == descriptor.api_key_id
+        assert result.api_key_name == "auth-result"
+        assert result.api_key_display == descriptor.api_key_display
+
+        async with get_async_session_context_manager() as async_session:
+            assert (
+                await fetch_api_key_auth_result("not-a-real-hash", async_session)
+                is None
+            )
+    finally:
+        remove_api_key(db_session, descriptor.api_key_id)

backend/tests/unit/onyx/auth/test_usage_credential_identity.py

@@ -0,0 +1,115 @@
+from typing import Any, cast
+
+import pytest
+from fastapi import Request
+
+from onyx.auth import users
+from onyx.db.enums import PatType
+from onyx.db.models import User
+from onyx.db.pat import PatAuthResult
+from shared_configs.contextvars import UsageCredentialIdentity
+from shared_configs.enums import UsageCredentialType
+
+
+def _bare_request(headers: list[tuple[bytes, bytes]] | None = None) -> Request:
+    return Request(
+        {"type": "http", "method": "GET", "path": "/", "headers": headers or []}
+    )
+
+
+async def _no_oauth_refresh(*_: Any, **__: Any) -> None:
+    return None
+
+
+async def _resolve(request: Request, user: User | None) -> User | None:
+    return await users._resolve_optional_user(
+        request,
+        cast(Any, None),
+        user,
+        cast(Any, None),
+    )
+
+
+@pytest.fixture(autouse=True)
+def _patch_auth_helpers(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(users, "_maybe_refresh_oauth_tokens", _no_oauth_refresh)
+    monkeypatch.setattr(users, "JWT_PUBLIC_KEY_URL", None)
+
+
+@pytest.mark.asyncio
+async def test_cookie_session_user_is_session_credential() -> None:
+    request = _bare_request()
+    user = cast(User, object())
+
+    assert await _resolve(request, user) is user
+    assert request.state.usage_credential == UsageCredentialIdentity(
+        UsageCredentialType.SESSION
+    )
+
+
+@pytest.mark.asyncio
+async def test_unauthenticated_request_has_no_credential() -> None:
+    request = _bare_request()
+
+    assert await _resolve(request, None) is None
+    assert getattr(request.state, "usage_credential", None) is None
+
+
+@pytest.mark.parametrize(
+    "pat_type, expected_credential_type",
+    [
+        (PatType.USER, UsageCredentialType.PAT),
+        (PatType.CRAFT, UsageCredentialType.CRAFT_PAT),
+    ],
+)
+@pytest.mark.asyncio
+async def test_pat_type_selects_credential_type(
+    monkeypatch: pytest.MonkeyPatch,
+    pat_type: PatType,
+    expected_credential_type: UsageCredentialType,
+) -> None:
+    pat_user = cast(User, object())
+    pat = PatAuthResult(
+        user=pat_user,
+        scopes=None,
+        pat_id=7,
+        pat_name="my-token",
+        pat_display="onyx_pat_...abcd",
+        pat_type=pat_type,
+    )
+
+    async def fake_resolve_pat(*_: Any, **__: Any) -> PatAuthResult:
+        return pat
+
+    monkeypatch.setattr(users, "get_hashed_pat_from_request", lambda _: "hashed")
+    monkeypatch.setattr(users, "resolve_pat", fake_resolve_pat)
+    request = _bare_request()
+
+    assert await _resolve(request, None) is pat_user
+    assert request.state.usage_credential == UsageCredentialIdentity(
+        expected_credential_type, "7", "my-token", "onyx_pat_...abcd"
+    )
+
+
+@pytest.mark.asyncio
+async def test_external_jwt_user_is_jwt_credential(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Drives the real _check_for_saml_and_jwt — stubbing it out tests nothing."""
+    jwt_user = cast(User, object())
+
+    async def fake_verify(_: str) -> dict[str, str]:
+        return {"sub": "someone"}
+
+    async def fake_get_or_create(*_: Any, **__: Any) -> User:
+        return jwt_user
+
+    monkeypatch.setattr(users, "JWT_PUBLIC_KEY_URL", "https://idp.example/jwks")
+    monkeypatch.setattr(users, "verify_jwt_token", fake_verify)
+    monkeypatch.setattr(users, "_get_or_create_user_from_jwt", fake_get_or_create)
+    request = _bare_request([(b"authorization", b"Bearer some-token")])
+
+    assert await _resolve(request, None) is jwt_user
+    assert request.state.usage_credential == UsageCredentialIdentity(
+        UsageCredentialType.JWT
+    )