New enterprise feature

services/demo_session/app/api/__init__.py

@@ -3,5 +3,6 @@
 from .demo_sessions import router as demo_sessions_router
 from .demo_accounts import router as demo_accounts_router
 from .demo_operations import router as demo_operations_router
+from .internal import router as internal_router
 
-__all__ = ["demo_sessions_router", "demo_accounts_router", "demo_operations_router"]
+__all__ = ["demo_sessions_router", "demo_accounts_router", "demo_operations_router", "internal_router"]

services/demo_session/app/api/demo_sessions.py

@@ -5,6 +5,7 @@ Demo Sessions API - Atomic CRUD operations on DemoSession model
 from fastapi import APIRouter, Depends, HTTPException, Path, Query, Request
 from typing import Optional
 from uuid import UUID
+from datetime import datetime, timezone
 import structlog
 import jwt
 
@@ -54,6 +55,41 @@ async def _background_cloning_task(session_id: str, session_obj_id: UUID, base_t
                 error=str(e),
                 exc_info=True
             )
+            # Attempt to update session status to failed if possible
+            try:
+                from app.core.database import db_manager
+                from app.models import DemoSession
+                from sqlalchemy import select, update
+
+                # Try to update the session directly in DB to mark it as failed
+                async with db_manager.session_factory() as update_db:
+                    from app.models import DemoSessionStatus
+                    update_result = await update_db.execute(
+                        update(DemoSession)
+                        .where(DemoSession.id == session_obj_id)
+                        .values(status=DemoSessionStatus.FAILED, cloning_completed_at=datetime.now(timezone.utc))
+                    )
+                    await update_db.commit()
+            except Exception as update_error:
+                logger.error(
+                    "Failed to update session status to FAILED after background task error",
+                    session_id=session_id,
+                    error=str(update_error)
+                )
+
+
+def _handle_task_result(task, session_id: str):
+    """Handle the result of the background cloning task"""
+    try:
+        # This will raise the exception if the task failed
+        task.result()
+    except Exception as e:
+        logger.error(
+            "Background cloning task failed with exception",
+            session_id=session_id,
+            error=str(e),
+            exc_info=True
+        )
 
 
 @router.post(
@@ -77,6 +113,7 @@ async def create_demo_session(
         session_manager = DemoSessionManager(db, redis)
         session = await session_manager.create_session(
             demo_account_type=request.demo_account_type,
+            subscription_tier=request.subscription_tier,
             user_id=request.user_id,
             ip_address=ip_address,
             user_agent=user_agent
@@ -92,10 +129,14 @@ async def create_demo_session(
         base_tenant_id = demo_config.get("base_tenant_id", str(session.base_demo_tenant_id))
 
         # Start cloning in background task with session ID (not session object)
-        asyncio.create_task(
+        # Store task reference in case we need to track it
+        task = asyncio.create_task(
             _background_cloning_task(session.session_id, session.id, base_tenant_id)
         )
 
+        # Add error handling for the task to prevent silent failures
+        task.add_done_callback(lambda t: _handle_task_result(t, session.session_id))
+
         # Generate session token
         session_token = jwt.encode(
             {
@@ -104,8 +145,8 @@ async def create_demo_session(
                 "demo_account_type": request.demo_account_type,
                 "exp": session.expires_at.timestamp()
             },
-            "demo-secret-key",
-            algorithm="HS256"
+            settings.JWT_SECRET_KEY,
+            algorithm=settings.JWT_ALGORITHM
         )
 
         return {

services/demo_session/app/api/internal.py

@@ -0,0 +1,82 @@
+"""
+Internal API for Demo Session Service
+Handles internal service-to-service operations
+"""
+
+from fastapi import APIRouter, Depends, HTTPException, Header
+from sqlalchemy.ext.asyncio import AsyncSession
+import structlog
+
+from app.core import get_db, settings
+from app.core.redis_wrapper import get_redis, DemoRedisWrapper
+from app.services.data_cloner import DemoDataCloner
+
+logger = structlog.get_logger()
+router = APIRouter()
+
+
+async def verify_internal_api_key(x_internal_api_key: str = Header(None)):
+    """Verify internal API key for service-to-service communication"""
+    required_key = settings.INTERNAL_API_KEY
+    if x_internal_api_key != required_key:
+        logger.warning("Unauthorized internal API access attempted")
+        raise HTTPException(status_code=403, detail="Invalid internal API key")
+    return True
+
+
+@router.post("/internal/demo/cleanup")
+async def cleanup_demo_session_internal(
+    cleanup_request: dict,
+    db: AsyncSession = Depends(get_db),
+    redis: DemoRedisWrapper = Depends(get_redis),
+    _: bool = Depends(verify_internal_api_key)
+):
+    """
+    Internal endpoint to cleanup demo session data for a specific tenant
+    Used by rollback mechanisms
+    """
+    try:
+        tenant_id = cleanup_request.get('tenant_id')
+        session_id = cleanup_request.get('session_id')
+
+        if not all([tenant_id, session_id]):
+            raise HTTPException(
+                status_code=400, 
+                detail="Missing required parameters: tenant_id, session_id"
+            )
+
+        logger.info(
+            "Internal cleanup requested", 
+            tenant_id=tenant_id,
+            session_id=session_id
+        )
+
+        data_cloner = DemoDataCloner(db, redis)
+        
+        # Delete session data for this tenant
+        await data_cloner.delete_session_data(
+            str(tenant_id),
+            session_id
+        )
+
+        logger.info(
+            "Internal cleanup completed",
+            tenant_id=tenant_id,
+            session_id=session_id
+        )
+
+        return {
+            "status": "completed",
+            "tenant_id": tenant_id,
+            "session_id": session_id
+        }
+
+    except Exception as e:
+        logger.error(
+            "Internal cleanup failed", 
+            error=str(e),
+            tenant_id=cleanup_request.get('tenant_id'),
+            session_id=cleanup_request.get('session_id'),
+            exc_info=True
+        )
+        raise HTTPException(status_code=500, detail=f"Failed to cleanup demo session: {str(e)}")