diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index 8725837871..2cc14005e6 100644 --- a/backend/onyx/auth/users.py +++ b/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) diff --git a/backend/onyx/db/api_key.py b/backend/onyx/db/api_key.py index e696cb3fa3..0b7d453c79 100644 --- a/backend/onyx/db/api_key.py +++ b/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, diff --git a/backend/onyx/db/pat.py b/backend/onyx/db/pat.py index 0efef6ec00..7430750a1c 100644 --- a/backend/onyx/db/pat.py +++ b/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: diff --git a/backend/shared_configs/contextvars.py b/backend/shared_configs/contextvars.py index 67fbd2de72..8b900d7160 100644 --- a/backend/shared_configs/contextvars.py +++ b/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() diff --git a/backend/shared_configs/enums.py b/backend/shared_configs/enums.py index 9a61d92176..a1204f1267 100644 --- a/backend/shared_configs/enums.py +++ b/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" diff --git a/backend/tests/external_dependency_unit/db/test_api_key_auth_result.py b/backend/tests/external_dependency_unit/db/test_api_key_auth_result.py new file mode 100644 index 0000000000..3245729e9c --- /dev/null +++ b/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) diff --git a/backend/tests/unit/onyx/auth/test_usage_credential_identity.py b/backend/tests/unit/onyx/auth/test_usage_credential_identity.py new file mode 100644 index 0000000000..e1f5529506 --- /dev/null +++ b/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 + )