auth.py

"""Session identity and OAuth wiring for per-user document isolation.

Owner identity resolution follows Design 4's principle that the authorization
boundary must exist *before* retrieval, never in prompts:

- Google OAuth (when OAUTH_GOOGLE_* env vars are configured) maps a verified
  email to the stable owner id ``google:<email>``.
- Without OAuth, sessions degrade to ``anon:<chainlit-session-id>`` so local
  development stays frictionless while uploads remain isolated per browser
  session.
- ``AUTH_REQUIRED=true`` hard-refuses unauthenticated sessions instead of
  degrading to anonymous owners.

Chainlit is imported lazily so this module stays cheap to unit-test and can
never trigger Chainlit's import-time OAuth provider validation outside app.py.
"""

import os
from typing import Any, Dict, Optional

from observability import get_logger

logger = get_logger("auth")


def google_oauth_configured() -> bool:
    """True when both Google OAuth env vars are present."""
    return bool(os.getenv("OAUTH_GOOGLE_CLIENT_ID")) and bool(
        os.getenv("OAUTH_GOOGLE_CLIENT_SECRET")
    )


def auth_required() -> bool:
    """When true, sessions without an authenticated identity are refused."""
    return os.getenv("AUTH_REQUIRED", "false").strip().lower() in ("1", "true", "yes", "on")


def build_oauth_user(provider_id: str, raw_user_data: Dict[str, Any]):
    """Map provider profile data to a Chainlit user with a stable identifier.

    Returns ``None`` to reject the login when the provider is unsupported or
    no verified email is available - isolation depends on stable ids.
    """
    if provider_id != "google":
        logger.warning(f"Unsupported OAuth provider requested: {provider_id}")
        return None
    from chainlit import User

    email = (raw_user_data.get("email") or "").strip().lower()
    if not email:
        logger.warning("Google OAuth payload missing email; rejecting login")
        return None
    return User(
        identifier=f"google:{email}",
        display_name=raw_user_data.get("name") or email,
        metadata={"provider": "google", "email": email},
    )


def register_oauth() -> bool:
    """Register Chainlit's oauth_callback when Google OAuth is configured.

    The ``@cl.oauth_callback`` decorator raises at decoration time when no
    provider env vars exist, so registration is conditional. Returns True
    when the callback was registered.
    """
    if not google_oauth_configured():
        logger.info(
            "OAuth not configured (missing OAUTH_GOOGLE_*); "
            "sessions will be isolated anonymously"
        )
        return False

    import chainlit as cl

    @cl.oauth_callback
    async def _oauth_callback(
        provider_id: str,
        token: str,
        raw_user_data: Dict[str, Any],
        default_user,
        id_token: Optional[str] = None,
    ):
        user = build_oauth_user(provider_id, raw_user_data)
        if user is not None:
            # Registry write is best-effort; login never depends on storage.
            try:
                from store import get_store

                get_store().upsert_user(
                    user.identifier,
                    email=user.metadata.get("email"),
                    name=user.display_name,
                    provider=provider_id,
                )
            except Exception as e:
                logger.warning(f"Could not persist user record: {e}")
        return user

    logger.info("Google OAuth login enabled")
    return True


def resolve_owner_id(user: Any = None, session_id: Optional[str] = None) -> Optional[str]:
    """Resolve the requester's owner id for retrieval isolation.

    Precedence: authenticated user identifier, then anonymous session id.
    Returns None only when authentication is required but absent (the caller
    must refuse the request) or when there is nothing to key on.
    """
    identifier = getattr(user, "identifier", None)
    if identifier:
        return str(identifier)

    if auth_required():
        return None

    if session_id:
        return f"anon:{session_id}"
    return None