PABPAT/TCI_Shield

1"""2Session Manager for TCI Shield FastAPI backend.3 4Manages in-memory session store for active conversations.5Each session holds a Strands Agent instance and conversation state.6Sessions expire after 30 minutes of inactivity.7Expired sessions are saved to DynamoDB before deletion.8"""9 10import uuid11import logging12from datetime import datetime, timedelta13 14import boto315from apscheduler.schedulers.asyncio import AsyncIOScheduler16from strands import Agent17from strands.models.bedrock import BedrockModel18 19from backend.core.enums import SessionStatus20 21logger = logging.getLogger(__name__)22 23# ============================================================24# SECTION 1 -- CONFIGURATION25# ============================================================26 27SESSION_TTL_MINUTES = 3028CLEANUP_INTERVAL_MINUTES = 529ABANDONED_TABLE = "tci_abandoned_sessions"30 31# ============================================================32# SECTION 2 -- IN-MEMORY SESSION STORE33# ============================================================34 35# Structure:36# {37#     "session_id": {38#         "agent":        Agent instance,39#         "tci_session":  dict (underwriting session state),40#         "created_at":   datetime,41#         "last_active":  datetime,42#         "status":       SessionStatus43#     }44# }45active_sessions: dict = {}46 47# ============================================================48# SECTION 3 -- SCHEDULER49# ============================================================50 51scheduler = AsyncIOScheduler()52 53# ============================================================54# SECTION 4 -- SESSION OPERATIONS55# ============================================================56 57def create_session(agent: Agent, tci_session: dict) -> str:58    """59    Creates a new session and stores it in active_sessions.60 61    Args:62        agent       : Strands Agent instance for this conversation63        tci_session : underwriting session state from tci_agent.py64 65    Returns:66        session_id : unique UUID string for this session67    """68    session_id = str(uuid.uuid4())69    now = datetime.now()70 71    active_sessions[session_id] = {72        "agent":       agent,73        "tci_session": tci_session,74        "created_at":  now,75        "last_active": now,76        "status":      SessionStatus.active,77    }78 79    logger.info(f"Session created: {session_id}")80    return session_id81 82 83def get_session(session_id: str) -> dict | None:84    """85    Retrieves a session by ID.86 87    Returns:88        session dict if found and active, None if not found or expired89    """90    session = active_sessions.get(session_id)91    if not session:92        logger.warning(f"Session not found: {session_id}")93        return None94 95    # Check if session has expired96    elapsed = datetime.now() - session["last_active"]97    if elapsed > timedelta(minutes=SESSION_TTL_MINUTES):98        logger.warning(f"Session expired: {session_id}")99        return None100 101    return session102 103 104def update_last_active(session_id: str):105    """106    Updates the last_active timestamp for a session.107    Called on every message received.108    """109    if session_id in active_sessions:110        active_sessions[session_id]["last_active"] = datetime.now()111        logger.debug(f"Session last_active updated: {session_id}")112 113 114def delete_session(session_id: str):115    """116    Deletes a session from active_sessions.117    Should be called after saving to DynamoDB.118    """119    if session_id in active_sessions:120        del active_sessions[session_id]121        logger.info(f"Session deleted from memory: {session_id}")122 123 124def get_session_count() -> int:125    """Returns number of active sessions."""126    return len(active_sessions)127 128# ============================================================129# SECTION 5 -- SESSION PERSISTENCE130# ============================================================131 132def _save_abandoned_session(session_id: str, session: dict):133    """134    Saves an abandoned session to DynamoDB tci_abandoned_sessions table.135    Called before deleting an expired session.136 137    Args:138        session_id : the session UUID139        session    : the full session dict140    """141    try:142        dynamodb = boto3.resource("dynamodb", region_name="us-east-1")143 144        # Create table if it doesn't exist145        try:146            client = boto3.client("dynamodb", region_name="us-east-1")147            client.describe_table(TableName=ABANDONED_TABLE)148        except client.exceptions.ResourceNotFoundException:149            dynamodb.create_table(150                TableName=ABANDONED_TABLE,151                KeySchema=[{"AttributeName": "session_id", "KeyType": "HASH"}],152                AttributeDefinitions=[{"AttributeName": "session_id", "AttributeType": "S"}],153                BillingMode="PAY_PER_REQUEST"154            )155            logger.info(f"Created table: {ABANDONED_TABLE}")156 157        table = dynamodb.Table(ABANDONED_TABLE)158 159        # Extract serialisable session data (Agent instance can't be serialised)160        tci_session = session.get("tci_session", {})161 162        # Remove non-serialisable objects163        safe_session = {164            k: v for k, v in tci_session.items()165            if isinstance(v, (str, int, float, bool, list, dict, type(None)))166        }167 168        table.put_item(Item={169            "session_id":    session_id,170            "status":        SessionStatus.expired,171            "created_at":    session["created_at"].isoformat(),172            "last_active":   session["last_active"].isoformat(),173            "abandoned_at":  datetime.now().isoformat(),174            "tci_session":   safe_session,175        })176 177        logger.info(f"Abandoned session saved to DynamoDB: {session_id}")178 179    except Exception as e:180        logger.error(f"Failed to save abandoned session {session_id}: {e}")181 182 183# ============================================================184# SECTION 6 -- CLEANUP JOB185# ============================================================186 187async def cleanup_expired_sessions():188    """189    Scheduled job that runs every CLEANUP_INTERVAL_MINUTES.190    Finds sessions inactive for more than SESSION_TTL_MINUTES.191    Saves them to DynamoDB then deletes from memory.192    """193    now = datetime.now()194    expired = []195 196    for session_id, session in active_sessions.items():197        elapsed = now - session["last_active"]198        if elapsed > timedelta(minutes=SESSION_TTL_MINUTES):199            expired.append(session_id)200 201    if expired:202        logger.info(f"Cleanup job: found {len(expired)} expired sessions")203 204    for session_id in expired:205        session = active_sessions.get(session_id)206        if session:207            _save_abandoned_session(session_id, session)208            delete_session(session_id)209            logger.info(f"Cleaned up expired session: {session_id}")210 211 212# ============================================================213# SECTION 7 -- SCHEDULER LIFECYCLE214# ============================================================215 216def start_scheduler():217    """218    Starts the APScheduler cleanup job.219    Called from main.py on application startup.220    """221    scheduler.add_job(222        cleanup_expired_sessions,223        trigger="interval",224        minutes=CLEANUP_INTERVAL_MINUTES,225        id="cleanup_expired_sessions",226        replace_existing=True,227    )228    scheduler.start()229    logger.info(f"Session cleanup scheduler started — runs every {CLEANUP_INTERVAL_MINUTES} minutes")230 231 232def stop_scheduler():233    """234    Stops the APScheduler.235    Called from main.py on application shutdown.236    """237    if scheduler.running:238        scheduler.shutdown()239        logger.info("Session cleanup scheduler stopped")