tatlock/src/core/context.py

"""
Request context using ContextVar for async-safe user/conversation tracking.

ContextVar provides task-local storage that automatically propagates through
async calls, eliminating the need to thread user identity through every function.

Usage:
    # At request entry (router):
    token = current_user.set(request.user or get_default_user())
    try:
        await service.process(request)
    finally:
        current_user.reset(token)

    # Anywhere in the codebase:
    from src.core.context import get_user
    user = get_user()  # Returns current request's user
"""

from contextvars import ContextVar


def get_default_user() -> str:
    """
    Get default user from config (environment-aware).

    - development/testing: llm_tester (isolated test scope)
    - production: jpmschweitzer (real user)
    """
    # Import here to avoid circular dependency
    from src.core.config import config

    return config.effective_default_user


# Request-scoped context variables (async-safe, isolated per request)
# Note: ContextVar default is evaluated at definition, so we use a sentinel
# and resolve the real default in get_user()
_USER_NOT_SET = "__user_not_set__"
current_user: ContextVar[str] = ContextVar("current_user", default=_USER_NOT_SET)
current_conversation: ContextVar[str | None] = ContextVar("current_conversation", default=None)


def apply_tenant_guard(user: str) -> str:
    """
    Enforce tenant isolation at request-context resolution.

    In non-production environments the production tenant must never be
    the effective user - a request that explicitly asks for it is forced
    to the reserved test tenant instead (with a loud log line).

    Comparison happens on the *sanitized* form of the user: every local
    namespace (Qdrant collection, Redis key) is derived through
    sanitize_user_id(), so any raw variant that collides with the
    production tenant after sanitization ("JPMSchweitzer",
    "jpmschweitzer.", " jpmschweitzer", ...) would otherwise resolve to
    the production namespaces. Those variants are forced too.
    """
    # Import here to avoid circular dependency
    from src.core.config import PRODUCTION_TENANT, TEST_TENANT, Environment, config
    from src.core.multi_tenancy import sanitize_user_id

    if config.ENVIRONMENT != Environment.PRODUCTION and sanitize_user_id(user) == sanitize_user_id(
        PRODUCTION_TENANT
    ):
        from src.core.logging_config import get_logger

        get_logger(__name__).warning(
            "tenant_guard_forced",
            environment=config.ENVIRONMENT.value,
            requested_tenant=user,
            forced_tenant=TEST_TENANT,
        )
        return TEST_TENANT
    return user


def get_user() -> str:
    """
    Get current user from request context.

    Returns:
        User identifier for the current request.
        Falls back to environment-aware default if not set.
        In non-production environments the production tenant is never
        returned - the tenant guard forces the reserved test tenant.

    Example:
        user = get_user()  # "llm_tester" (dev) or "jpmschweitzer" (prod)
    """
    user = current_user.get()
    if user == _USER_NOT_SET:
        return get_default_user()
    return apply_tenant_guard(user)


def get_conversation_id() -> str | None:
    """
    Get current conversation ID from request context.

    Returns:
        Conversation ID if set, None otherwise.

    Example:
        conv_id = get_conversation_id()  # "conv_abc123" or None
    """
    return current_conversation.get()


class RequestContext:
    """
    Context manager for setting request-scoped context.

    Provides a cleaner alternative to manual token management.

    Usage:
        async with RequestContext(user="alice", conversation_id="conv_123"):
            # All code here sees user="alice"
            result = await some_service.process()
    """

    def __init__(
        self,
        user: str | None = None,
        conversation_id: str | None = None,
    ):
        """
        Initialize request context.

        Args:
            user: User identifier (defaults to environment-aware user if None)
            conversation_id: Conversation ID (optional)
        """
        self.user = user or get_default_user()
        self.conversation_id = conversation_id
        self._user_token = None
        self._conv_token = None

    async def __aenter__(self) -> "RequestContext":
        """Set context variables on entry."""
        self._user_token = current_user.set(self.user)
        self._conv_token = current_conversation.set(self.conversation_id)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Reset context variables on exit."""
        if self._user_token is not None:
            current_user.reset(self._user_token)
        if self._conv_token is not None:
            current_conversation.reset(self._conv_token)

    def __enter__(self) -> "RequestContext":
        """Sync context manager entry (for non-async code)."""
        self._user_token = current_user.set(self.user)
        self._conv_token = current_conversation.set(self.conversation_id)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Sync context manager exit."""
        if self._user_token is not None:
            current_user.reset(self._user_token)
        if self._conv_token is not None:
            current_conversation.reset(self._conv_token)