hyf-backend/th_agenter/core/context.py

"""
HTTP请求上下文管理，如：获取当前登录用户信息及Token信息
"""

from contextvars import ContextVar
from typing import Optional
import threading
from ..models.user import User
from loguru import logger

# Context variable to store current user
current_user_context: ContextVar[Optional[dict]] = ContextVar('current_user', default=None)

# Thread-local storage as backup
_thread_local = threading.local()


class UserContext:
    """User context manager for accessing current user globally."""
    
    @staticmethod
    def set_current_user(user: User, canLog: bool = False) -> None:
        """Set current user in context."""
        if canLog:
            logger.info(f"[UserContext] - 设置用户上下文 {user.username} (ID: {user.id})")
        
        # Store user information as a dictionary instead of the SQLAlchemy model
        user_dict = {
            'id': user.id,
            'username': user.username,
            'email': user.email,
            'full_name': user.full_name,
            'is_active': user.is_active
        }
        
        # Set in ContextVar
        current_user_context.set(user_dict)
        
        # Also set in thread-local as backup
        _thread_local.current_user = user_dict
        
        # Verify it was set
        verify_user = current_user_context.get()
        if canLog:
            logger.info(f"[UserContext] - 验证 - ContextVar 用户: {verify_user.get('username') if verify_user else None}")
    
    @staticmethod
    def set_current_user_with_token(user: User, canLog: bool = False):
        """Set current user in context and return token for cleanup."""
        if canLog:
            logger.info(f"[UserContext] - 设置用户上下文 {user.username} (ID: {user.id})")  
        
        # Store user information as a dictionary instead of the SQLAlchemy model
        user_dict = {
            'id': user.id,
            'username': user.username,
            'email': user.email,
            'full_name': user.full_name,
            'is_active': user.is_active
        }
        
        # Set in ContextVar and get token
        token = current_user_context.set(user_dict)
        
        # Also set in thread-local as backup
        _thread_local.current_user = user_dict
        
        # Verify it was set
        verify_user = current_user_context.get()
        if canLog:
            logger.info(f"[UserContext] - 验证 - ContextVar 用户: {verify_user.get('username') if verify_user else None}")
        
        return token
    
    @staticmethod
    def reset_current_user_token(token):
        """Reset current user context using token."""
        logger.info("[UserContext] - Resetting user context using token")
        
        # Reset ContextVar using token
        current_user_context.reset(token)
        
        # Clear thread-local as well
        if hasattr(_thread_local, 'current_user'):
            delattr(_thread_local, 'current_user')
    
    @staticmethod
    def get_current_user() -> Optional[dict]:
        """Get current user from context."""        
        # Try ContextVar first
        user = current_user_context.get()        
        if user:
            # logger.info(f"[UserContext] - 取得当前用户为 ContextVar 用户: {user.get('username') if user else None}")
            return user
        
        # Fallback to thread-local
        user = getattr(_thread_local, 'current_user', None)        
        if user:
            # logger.info(f"[UserContext] - 取得当前用户为线程本地用户: {user.get('username') if user else None}")
            return user
        
        logger.error("[UserContext] - 上下文未找到当前用户 (neither ContextVar nor thread-local)")
        return None
    
    @staticmethod
    def get_current_user_id() -> Optional[int]:
        """Get current user ID from context."""
        try:
            user = UserContext.get_current_user()
            return user.get('id') if user else None
        except Exception as e:
            logger.error(f"[UserContext] - Error getting current user ID: {e}")
            return None
    
    @staticmethod
    def clear_current_user(canLog: bool = False) -> None:
        """Clear current user from context."""
        if canLog:
            logger.info("[UserContext] - 清除当前用户上下文")
        
        current_user_context.set(None)
        if hasattr(_thread_local, 'current_user'):
            delattr(_thread_local, 'current_user')
    
    @staticmethod
    def require_current_user() -> dict:
        """Get current user from context, raise exception if not found."""
        # Use the same logic as get_current_user to check both ContextVar and thread-local
        user = UserContext.get_current_user()
        if user is None:
            from fastapi import HTTPException, status
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="No authenticated user in context"
            )
        return user
    
    @staticmethod
    def require_current_user_id() -> int:
        """Get current user ID from context, raise exception if not found."""
        user = UserContext.require_current_user()
        return user.get('id')