The file /home/ubuntu/pipecat/src/bot/services/limit_checker.py has been updated successfully.

"""Usage limit checker — 3-layer credits: plan → topup → free."""

from loguru import logger

def can_generate(bot_session) -> str:
    """Check if user can generate content.

    Returns:
        "" — allowed
        "subscribe" — no plan, free exhausted, payments ON → show plans popup
        "topup" — has plan but exhausted, payments ON → show topup popup
        "cooldown" — limits exhausted, payments OFF → agent says wait
    """
    if bot_session.plan_remaining > 0:
        return ""
    if bot_session.topup_remaining > 0:
        return ""
    if bot_session.free_remaining > 0:
        return ""

    # Blocked — check if payments are enabled
    if not bot_session.payments_enabled:
        return "cooldown"

    if bot_session.subscription_id:
        return "topup"
    return "subscribe"

def use_credit(bot_session):
    """Decrement one credit after successful generation. Plan first, then topup, then free."""
    if bot_session.plan_remaining > 0:
        bot_session.plan_remaining -= 1
        bot_session.credits_used_plan += 1
        logger.debug(f"CREDIT_USED_PLAN: device={bot_session.device_id} plan_remaining={bot_session.plan_remaining}")
    elif bot_session.topup_remaining > 0:
        bot_session.topup_remaining -= 1
        bot_session.credits_used_topup += 1
        logger.debug(f"CREDIT_USED_TOPUP: device={bot_session.device_id} topup_remaining={bot_session.topup_remaining}")
    elif bot_session.free_remaining > 0:
        bot_session.free_remaining -= 1
        logger.debug(f"CREDIT_USED_FREE: device={bot_session.device_id} free_remaining={bot_session.free_remaining}")

async def send_limit_popup(bot_session, rtvi_processor, block_reason: str):
    """Send limit reached popup or cooldown message to the frontend.

    If block_reason is "cooldown", sends a cooldown message instead of payment popup.
    The agent will verbally tell the user to wait.
    """
    from src.service.usage_service import get_plans_for_device, get_topup_for_device, get_topups_for_content

    content_type = bot_session.content_type or "image_gen"
    device_id = bot_session.device_id

    if block_reason == "cooldown":
        cooldown_minutes = bot_session.cooldown_minutes
        logger.info(f"COOLDOWN_LIMIT: device={device_id} content={content_type} cooldown_minutes={cooldown_minutes}")
        await rtvi_processor.send_server_message({
            "type": "show_cooldown",
            "data": {
                "content_type": content_type,
                "cooldown_minutes": cooldown_minutes,
            },
        })
        return

    # Payment popup flow (existing behavior)
    from src.service.usage_service import _get_device_region, _get_payment_provider

    region = await _get_device_region(device_id)
    payment_provider = _get_payment_provider(region)

    try:
        if block_reason == "topup":
            plan_data = await get_topup_for_device(device_id, content_type, payment_provider)
            topups_list = await get_topups_for_content(content_type, payment_provider)
        else:
            plan_data = await get_plans_for_device(device_id, content_type, payment_provider)
            topups_list = []

        popup_type = "show_plans_popup" if block_reason == "subscribe" else "show_topup_popup"
        await rtvi_processor.send_server_message({
            "type": popup_type,
            "data": {
                "content_type": content_type,
                "reason": block_reason,
                "plan": plan_data,
                "topups": topups_list,
            },
        })
        logger.info(f"LIMIT_POPUP_SENT: device={device_id} reason={block_reason} content={content_type} provider={payment_provider} plan_id={plan_data.get('plan_id') or plan_data.get('topup_id', 'none')}")
    except Exception as e:
        logger.error(f"LIMIT_POPUP_FAILED: device={device_id} error={e}")
        popup_type = "show_plans_popup" if block_reason == "subscribe" else "show_topup_popup"
        await rtvi_processor.send_server_message({"type": popup_type, "content_type": content_type})

-        "subscribe" — no plan, free exhausted → show plans popup

-        "topup" — has plan but exhausted → show topup popup

+        "subscribe" — no plan, free exhausted, payments ON → show plans popup

+        "topup" — has plan but exhausted, payments ON → show topup popup

+        "cooldown" — limits exhausted, payments OFF → agent says wait

+    # Blocked — check if payments are enabled

+    if not bot_session.payments_enabled:

-    """Send limit reached popup with plan/topup data to the frontend.

+    """Send limit reached popup or cooldown message to the frontend.

-    Fetches plan data from DB so the pricing screen shows actual prices.

+    If block_reason is "cooldown", sends a cooldown message instead of payment popup.

+    The agent will verbally tell the user to wait.

-    # Determine payment provider from region

+        cooldown_minutes = bot_session.cooldown_minutes

+        logger.info(f"COOLDOWN_LIMIT: device={device_id} content={content_type} cooldown_minutes={cooldown_minutes}")

+        await rtvi_processor.send_server_message({

+                "content_type": content_type,

+                "cooldown_minutes": cooldown_minutes,

+    # Payment popup flow (existing behavior)

-        # Fallback: send without data (old behavior)

"""Usage limit checker — 3-layer credits: plan → topup → free."""

from loguru import logger

def can_generate(bot_session) -> str:
    """Check if user can generate content.

    Returns:
        "" — allowed
        "subscribe" — no plan, free exhausted → show plans popup
        "topup" — has plan but exhausted → show topup popup
    """
    if bot_session.plan_remaining > 0:
        return ""
    if bot_session.topup_remaining > 0:
        return ""
    if bot_session.free_remaining > 0:
        return ""

    # Blocked
    if bot_session.subscription_id:
        return "topup"
    return "subscribe"

def use_credit(bot_session):
    """Decrement one credit after successful generation. Plan first, then topup, then free."""
    if bot_session.plan_remaining > 0:
        bot_session.plan_remaining -= 1
        bot_session.credits_used_plan += 1
        logger.debug(f"CREDIT_USED_PLAN: device={bot_session.device_id} plan_remaining={bot_session.plan_remaining}")
    elif bot_session.topup_remaining > 0:
        bot_session.topup_remaining -= 1
        bot_session.credits_used_topup += 1
        logger.debug(f"CREDIT_USED_TOPUP: device={bot_session.device_id} topup_remaining={bot_session.topup_remaining}")
    elif bot_session.free_remaining > 0:
        bot_session.free_remaining -= 1
        logger.debug(f"CREDIT_USED_FREE: device={bot_session.device_id} free_remaining={bot_session.free_remaining}")

async def send_limit_popup(bot_session, rtvi_processor, block_reason: str):
    """Send limit reached popup with plan/topup data to the frontend.

    Fetches plan data from DB so the pricing screen shows actual prices.
    """
    from src.service.usage_service import get_plans_for_device, get_topup_for_device, get_topups_for_content

    content_type = bot_session.content_type or "image_gen"
    device_id = bot_session.device_id

    # Determine payment provider from region
    from src.service.usage_service import _get_device_region, _get_payment_provider
    region = await _get_device_region(device_id)
    payment_provider = _get_payment_provider(region)

    try:
        if block_reason == "topup":
            plan_data = await get_topup_for_device(device_id, content_type, payment_provider)
            topups_list = await get_topups_for_content(content_type, payment_provider)
        else:
            plan_data = await get_plans_for_device(device_id, content_type, payment_provider)
            topups_list = []

        popup_type = "show_plans_popup" if block_reason == "subscribe" else "show_topup_popup"
        await rtvi_processor.send_server_message({
            "type": popup_type,
            "data": {
                "content_type": content_type,
                "reason": block_reason,
                "plan": plan_data,
                "topups": topups_list,
            },
        })
        logger.info(f"LIMIT_POPUP_SENT: device={device_id} reason={block_reason} content={content_type} provider={payment_provider} plan_id={plan_data.get('plan_id') or plan_data.get('topup_id', 'none')}")
    except Exception as e:
        logger.error(f"LIMIT_POPUP_FAILED: device={device_id} error={e}")
        # Fallback: send without data (old behavior)
        popup_type = "show_plans_popup" if block_reason == "subscribe" else "show_topup_popup"
        await rtvi_processor.send_server_message({"type": popup_type, "content_type": content_type})


Now add `payments_enabled` and `cooldown_minutes` to BotSessionState and load from DB at startup:


123	class BotSessionState:
124	    def __init__(self, session_id: str, device_id: str, bot_start_time: float):
125	        self.session_id = session_id
126	        self.device_id = device_id
127	        self.chunk_counter = 0
128	        self.user_stopped_speaking_time: Optional[float] = None
129	        self.current_user_id: Optional[str] = None
130	        self.greeting_sent = False
131	        self.user_spoke_first = False
132	        self.audio_recorder = None
133	        self.user_joined = False
134	        self.bot_joined = False
135	        self.bot_idle = False
136	        self.no_user_timeout_task = None
137	        self.participants: dict = {}
138	        self.session_start_time = bot_start_time
139	        self.bot_start_time = bot_start_time
140	        self.app_review_eligible = False
141	        self.vad_analyzer = None
142	        self.audio_buffer = None
143	        self._last_upload_time: float = 0
144	        self.uploaded_image: Optional[dict] = None
145	        self.last_image: Optional[dict] = None
146	        self.last_generated_image: Optional[dict] = None
147	        self.image_batch: list = []  # NEW: Batch of images uploaded together
148	        self.max_batch_size: int = 2  # NEW: Max images per batch
149	        self.session_created_at_start = False
150	        self.video_enabled = False  # NEW: Track if video generation is enabled for this session
151	        self.is_processing = False  # NEW: In-memory lock for video/image processing (prevents concurrent requests)
152	        self.generated_videos: List[Dict] = []  # Store generated videos to save at session end
153	        self.generated_ppts: List[Dict] = []    # Store generated PPTs to save at session end
154	        self.generated_music: List[Dict] = []   # Store generated music to save at session end
155
156	        # Pipeline service references for explicit cleanup (prevents memory leaks)
157	        self.stt_service = None       # STT service (Deepgram/OpenAI WebSocket)
158	        self.llm_service = None       # LLM service (Gemini/OpenAI HTTP client)
159	        self.transport = None         # SmallWebRTC transport
160	        self.pipeline = None          # Pipeline object
161	        self.context_aggregator = None  # LLM context (grows with conversation)
162	        self.transport_type: str = "smallwebrtc"  # "smallwebrtc" or "daily"
163
164	        # Background task tracking for proper cancellation
165	        self._background_tasks: List[asyncio.Task] = []
166
167	        # Persona
168	        self.persona_id: Optional[str] = None
169	        self.persona_name: Optional[str] = None
170	        self.content_type: Optional[str] = None  # image_gen, video_gen, ad_video, ppt_gen, music_gen
171
172	        # App key

class BotSessionState:
    def __init__(self, session_id: str, device_id: str, bot_start_time: float):
        self.session_id = session_id
        self.device_id = device_id
        self.chunk_counter = 0
        self.user_stopped_speaking_time: Optional[float] = None
        self.current_user_id: Optional[str] = None
        self.greeting_sent = False
        self.user_spoke_first = False
        self.audio_recorder = None
        self.user_joined = False
        self.bot_joined = False
        self.bot_idle = False
        self.no_user_timeout_task = None
        self.participants: dict = {}
        self.session_start_time = bot_start_time
        self.bot_start_time = bot_start_time
        self.app_review_eligible = False
        self.vad_analyzer = None
        self.audio_buffer = None
        self._last_upload_time: float = 0
        self.uploaded_image: Optional[dict] = None
        self.last_image: Optional[dict] = None
        self.last_generated_image: Optional[dict] = None
        self.image_batch: list = []  # NEW: Batch of images uploaded together
        self.max_batch_size: int = 2  # NEW: Max images per batch
        self.session_created_at_start = False
        self.video_enabled = False  # NEW: Track if video generation is enabled for this session
        self.is_processing = False  # NEW: In-memory lock for video/image processing (prevents concurrent requests)
        self.generated_videos: List[Dict] = []  # Store generated videos to save at session end
        self.generated_ppts: List[Dict] = []    # Store generated PPTs to save at session end
        self.generated_music: List[Dict] = []   # Store generated music to save at session end

        # Pipeline service references for explicit cleanup (prevents memory leaks)
        self.stt_service = None       # STT service (Deepgram/OpenAI WebSocket)
        self.llm_service = None       # LLM service (Gemini/OpenAI HTTP client)
        self.transport = None         # SmallWebRTC transport
        self.pipeline = None          # Pipeline object
        self.context_aggregator = None  # LLM context (grows with conversation)
        self.transport_type: str = "smallwebrtc"  # "smallwebrtc" or "daily"

        # Background task tracking for proper cancellation
        self._background_tasks: List[asyncio.Task] = []

        # Persona
        self.persona_id: Optional[str] = None
        self.persona_name: Optional[str] = None
        self.content_type: Optional[str] = None  # image_gen, video_gen, ad_video, ppt_gen, music_gen

        # App key


172	        # App key
173	        self.app_key: Optional[str] = None
174	        self.language: Optional[str] = None
175	        self.source: Optional[str] = None
176	        self.app_version: Optional[int] = None  # v5+ sends this to differentiate from prod
177
178	        # Per-content credits (3-layer: plan → topup → free)
179	        self.free_remaining: int = 0
180	        self.plan_remaining: int = 0
181	        self.topup_remaining: int = 0
182	        self.subscription_id: Optional[str] = None
183	        self.plan_content_type: Optional[str] = None  # tracks if credits from specific or 'all' plan
184	        self.credits_used_plan: int = 0   # track for session-end persist
185	        self.credits_used_topup: int = 0  # track for session-end persist
186	        self.session_saved: bool = False  # flag to prevent double-saving
187
188	    def get_session_info(self) -> str:
189	        """Get formatted session info for logging."""
190	        info = f"session_id={self.session_id} device_id={self.device_id}"
191	        if self.persona_id:
192	            info += f" persona={self.persona_id}"
193	            if self.persona_name:
194	                info += f" persona_name={self.persona_name}"
195	        return info
196
197	    def add_participant(self, participant_id: str, name: str, participant_type: str, joined_at: float):
198	        self.participants[participant_id] = {
199	            "name": name,
200	            "type": participant_type,
201	            "joined_at": joined_at

        # App key
        self.app_key: Optional[str] = None
        self.language: Optional[str] = None
        self.source: Optional[str] = None
        self.app_version: Optional[int] = None  # v5+ sends this to differentiate from prod

        # Per-content credits (3-layer: plan → topup → free)
        self.free_remaining: int = 0
        self.plan_remaining: int = 0
        self.topup_remaining: int = 0
        self.subscription_id: Optional[str] = None
        self.plan_content_type: Optional[str] = None  # tracks if credits from specific or 'all' plan
        self.credits_used_plan: int = 0   # track for session-end persist
        self.credits_used_topup: int = 0  # track for session-end persist
        self.session_saved: bool = False  # flag to prevent double-saving

    def get_session_info(self) -> str:
        """Get formatted session info for logging."""
        info = f"session_id={self.session_id} device_id={self.device_id}"
        if self.persona_id:
            info += f" persona={self.persona_id}"
            if self.persona_name:
                info += f" persona_name={self.persona_name}"
        return info

    def add_participant(self, participant_id: str, name: str, participant_type: str, joined_at: float):
        self.participants[participant_id] = {
            "name": name,
            "type": participant_type,
            "joined_at": joined_at


Edit

        self.credits_used_plan: int = 0   # track for session-end persist
        self.credits_used_topup: int = 0  # track for session-end persist
        self.session_saved: bool = False  # flag to prevent double-saving

        self.credits_used_plan: int = 0   # track for session-end persist
        self.credits_used_topup: int = 0  # track for session-end persist
        self.session_saved: bool = False  # flag to prevent double-saving

        # Payments config (loaded from DB at session start)
        self.payments_enabled: bool = True  # True = show payment popup, False = cooldown message
        self.cooldown_minutes: int = 30     # default cooldown when payments off


The file /home/ubuntu/pipecat/src/bot/core/config.py has been updated successfully.

        self.credits_used_plan: int = 0   # track for session-end persist
        self.credits_used_topup: int = 0  # track for session-end persist
        self.session_saved: bool = False  # flag to prevent double-saving

        # Payments config (loaded from DB at session start)
        self.payments_enabled: bool = True  # True = show payment popup, False = cooldown message
        self.cooldown_minutes: int = 30     # default cooldown when payments off

import asyncio
import threading
import time
from datetime import datetime
from typing import Optional, List, Dict
from loguru import logger

from src.config import constants

__all__ = [
    'load_system_instruction',
    'cleanup_session',
    'BotSessionState',
    'ChunkState',
    'TranscriptManager',
]

class ChunkState:
    def __init__(self):
        self.counter: int = 0
        self.chunks: List[str] = []
        self.session_id: Optional[str] = None

    def reset(self, session_id: str):
        self.session_id = session_id
        self.counter = 0
        self.chunks = []

    def clear(self):
        """Clear all chunk data for memory cleanup."""
        self.counter = 0
        self.chunks.clear()
        self.session_id = None

    def add_chunk(self, chunk_key: str) -> int:
        chunk_num = self.counter
        self.counter += 1
        self.chunks.append(chunk_key)
        return chunk_num

    def get_chunk_key(self, session_id: str, chunk_num: int) -> str:
        return f"maya-c/{session_id}/chunks/chunk_{chunk_num:03d}.wav"

    @property
    def total_chunks(self) -> int:
        return self.counter

    @property
    def chunk_list(self) -> List[str]:
        return self.chunks.copy()

class TranscriptManager:
    def __init__(self):
        self._messages: List[Dict] = []

    def clear(self):
        """Clear all messages for memory cleanup."""
        self._messages.clear()

    def add_message(self, role: str, content: str, msg_type: str = None, **kwargs):
        entry = {
            "timestamp": datetime.now().isoformat(),
            "role": role,
            "content": content
        }
        if msg_type:
            entry["type"] = msg_type
        entry.update(kwargs)
        self._messages.append(entry)

    def add_user_message(self, content: str):
        self.add_message("user", content)

    def add_assistant_message(self, content: str):
        self.add_message("assistant", content)

    def add_system_message(self, content: str):
        self.add_message("system", content)

    def add_generation(self, generation_type: str, content: Dict):
        self._messages.append({
            "timestamp": datetime.now().isoformat(),
            "role": "system",
            "type": "generation",
            "generation_type": generation_type,
            "content": content
        })

    def add_image_upload(self, mime_type: str, size_bytes: int, s3_url: str = None):
        self._messages.append({
            "timestamp": datetime.now().isoformat(),
            "role": "user",
            "type": "image_upload",
            "content": {"mime_type": mime_type, "size_bytes": size_bytes, "s3_url": s3_url}
        })

    @property
    def messages(self) -> List[Dict]:
        return self._messages

    @property
    def count(self) -> int:
        return len(self._messages)

    def get_user_messages(self) -> List[Dict]:
        return [m for m in self._messages if m.get("role") == "user"]

    def get_assistant_messages(self) -> List[Dict]:
        return [m for m in self._messages if m.get("role") == "assistant"]

    def get_generations(self) -> List[Dict]:
        return [m for m in self._messages if m.get("type") == "generation"]

    def get_image_generations_count(self) -> int:
        return len([m for m in self._messages if m.get("generation_type") == "image"])

    def get_image_edits_count(self) -> int:
        return len([m for m in self._messages if m.get("generation_type") == "image_edit"])

class BotSessionState:
    def __init__(self, session_id: str, device_id: str, bot_start_time: float):
        self.session_id = session_id
        self.device_id = device_id
        self.chunk_counter = 0
        self.user_stopped_speaking_time: Optional[float] = None
        self.current_user_id: Optional[str] = None
        self.greeting_sent = False
        self.user_spoke_first = False
        self.audio_recorder = None
        self.user_joined = False
        self.bot_joined = False
        self.bot_idle = False
        self.no_user_timeout_task = None
        self.participants: dict = {}
        self.session_start_time = bot_start_time
        self.bot_start_time = bot_start_time
        self.app_review_eligible = False
        self.vad_analyzer = None
        self.audio_buffer = None
        self._last_upload_time: float = 0
        self.uploaded_image: Optional[dict] = None
        self.last_image: Optional[dict] = None
        self.last_generated_image: Optional[dict] = None
        self.image_batch: list = []  # NEW: Batch of images uploaded together
        self.max_batch_size: int = 2  # NEW: Max images per batch
        self.session_created_at_start = False
        self.video_enabled = False  # NEW: Track if video generation is enabled for this session
        self.is_processing = False  # NEW: In-memory lock for video/image processing (prevents concurrent requests)
        self.generated_videos: List[Dict] = []  # Store generated videos to save at session end
        self.generated_ppts: List[Dict] = []    # Store generated PPTs to save at session end
        self.generated_music: List[Dict] = []   # Store generated music to save at session end

        # Pipeline service references for explicit cleanup (prevents memory leaks)
        self.stt_service = None       # STT service (Deepgram/OpenAI WebSocket)
        self.llm_service = None       # LLM service (Gemini/OpenAI HTTP client)
        self.transport = None         # SmallWebRTC transport
        self.pipeline = None          # Pipeline object
        self.context_aggregator = None  # LLM context (grows with conversation)
        self.transport_type: str = "smallwebrtc"  # "smallwebrtc" or "daily"

        # Background task tracking for proper cancellation
        self._background_tasks: List[asyncio.Task] = []

        # Persona
        self.persona_id: Optional[str] = None
        self.persona_name: Optional[str] = None
        self.content_type: Optional[str] = None  # image_gen, video_gen, ad_video, ppt_gen, music_gen

        # App key
        self.app_key: Optional[str] = None
        self.language: Optional[str] = None
        self.source: Optional[str] = None
        self.app_version: Optional[int] = None  # v5+ sends this to differentiate from prod

        # Per-content credits (3-layer: plan → topup → free)
        self.free_remaining: int = 0
        self.plan_remaining: int = 0
        self.topup_remaining: int = 0
        self.subscription_id: Optional[str] = None
        self.plan_content_type: Optional[str] = None  # tracks if credits from specific or 'all' plan
        self.credits_used_plan: int = 0   # track for session-end persist
        self.credits_used_topup: int = 0  # track for session-end persist
        self.session_saved: bool = False  # flag to prevent double-saving

    def get_session_info(self) -> str:
        """Get formatted session info for logging."""
        info = f"session_id={self.session_id} device_id={self.device_id}"
        if self.persona_id:
            info += f" persona={self.persona_id}"
            if self.persona_name:
                info += f" persona_name={self.persona_name}"
        return info

    def add_participant(self, participant_id: str, name: str, participant_type: str, joined_at: float):
        self.participants[participant_id] = {
            "name": name,
            "type": participant_type,
            "joined_at": joined_at
        }

    def get_participant(self, participant_id: str) -> dict:
        return self.participants.get(participant_id, {})

    def can_upload_image(self, cooldown: float = 2.0) -> bool:
        now = time.time()
        if now - self._last_upload_time < cooldown:
            return False
        self._last_upload_time = now
        return True

    def set_uploaded_image(self, image_info: dict):
        """Set single uploaded image (backward compatible)."""
        # Clear stale generated/edited images when user uploads fresh
        self.last_generated_image = None
        self.image_batch = []
        # Set new upload
        self.uploaded_image = image_info
        self.last_image = image_info

    def set_image_batch(self, images: list):
        """Store a batch of images uploaded together."""
        # Clear stale generated/edited images when user uploads fresh batch
        self.last_generated_image = None
        self.uploaded_image = None
        # Set new batch
        self.image_batch = images[:self.max_batch_size]  # Limit to max
        if self.image_batch:
            self.last_image = self.image_batch[0]
            self.uploaded_image = self.image_batch[0]  # Backward compatibility

    def get_image_batch(self) -> list:
        """Get current image batch."""
        return self.image_batch

    def get_batch_size(self) -> int:
        """Get number of images in current batch."""
        return len(self.image_batch)

    def clear_image_batch(self):
        """Clear image batch."""
        self.image_batch = []
        self.uploaded_image = None

    def set_generated_image(self, image_info: dict):
        """Store generated image for iteration."""
        self.last_generated_image = image_info
        self.last_image = image_info

    def add_generated_video(self, video_type: str, metadata: Dict, prompt: str = None):
        """Add a generated video to the list for saving at session end."""
        self.generated_videos.append({
            "video_type": video_type,
            "metadata": metadata,
            "prompt": prompt
        })

    def get_generated_videos(self) -> List[Dict]:
        """Get all generated videos for this session."""
        return self.generated_videos

    def add_generated_ppt(self, metadata: Dict, topic: str = None):
        """Add a generated PPT to the list for saving at session end."""
        self.generated_ppts.append({
            "metadata": metadata,
            "topic": topic
        })

    def get_generated_ppts(self) -> List[Dict]:
        """Get all generated PPTs for this session."""
        return self.generated_ppts

    def add_generated_music(self, metadata: Dict, prompt: str = None):
        """Add a generated music track to the list for saving at session end."""
        self.generated_music.append({
            "metadata": metadata,
            "prompt": prompt
        })

    def get_generated_music(self) -> List[Dict]:
        """Get all generated music for this session."""
        return self.generated_music

    def set_edited_image(self, image_info: dict):
        """Store edited image for iteration and clear uploads."""
        self.last_generated_image = image_info
        self.last_image = image_info
        # Clear uploads after edit (ready for new uploads)
        self.uploaded_image = None
        self.image_batch = []

    def clear_uploaded_image(self):
        """Clear single uploaded image (backward compatible)."""
        self.uploaded_image = None

    def clear_all_images(self):
        """Clear ALL image data for memory cleanup at session end."""
        self.uploaded_image = None
        self.last_image = None
        self.last_generated_image = None
        self.image_batch = []

    def get_image_to_edit(self) -> Optional[dict]:
        """Get the image to edit (supports iteration)."""
        if self.last_image:
            return self.last_image
        return self.uploaded_image or self.last_generated_image

    def get_uploaded_image(self) -> Optional[dict]:
        """Get the currently uploaded image (for i2v video generation).
        Also returns last_generated_image if no uploaded image is available,
        since users may want to create video from a generated image too.
        """
        return self.uploaded_image or self.last_generated_image

    def track_background_task(self, task: asyncio.Task):
        """Track a fire-and-forget background task for cleanup cancellation."""
        self._background_tasks.append(task)
        # Auto-remove when done to avoid accumulation
        def _on_task_done(t):
            try:
                self._background_tasks.remove(t)
            except ValueError:
                pass  # Already removed by cancel_background_tasks
        task.add_done_callback(_on_task_done)

    async def cancel_background_tasks(self):
        """Cancel all tracked background tasks. Call during session cleanup."""
        if not self._background_tasks:
            return

        tasks_to_cancel = list(self._background_tasks)
        for task in tasks_to_cancel:
            if not task.done():
                task.cancel()

        # Wait briefly for cancellation to propagate
        if tasks_to_cancel:
            await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
            logger.debug(f"CLEANUP: cancelled {len(tasks_to_cancel)} background tasks session_id={self.session_id}")

        self._background_tasks.clear()

    def cancel_timeout_task(self):
        if self.no_user_timeout_task and not self.no_user_timeout_task.done():
            self.no_user_timeout_task.cancel()

    def can_generate_images(self) -> bool:
        """All characters can generate images."""
        return True

    def can_edit_images(self) -> bool:
        """All characters can edit images."""
        return True

async def load_system_instruction(
    language: Optional[str] = None,
    persona_prompt: Optional[str] = None,
    persona_id: Optional[str] = None,
) -> str:
    """Assemble the system prompt: global + language_instruction + language_intro + personality.

    Fetches global prompt from DB (character_prompt), language instruction and
    intro from configurations_v2, and personality from persona's system_prompt.
    """
    import os
    from src.repository.config_v2_repository import config_v2_repo
    from src.repository.config_repository import config_repo

    lang_code = language or "hi"

    # 1. Global prompt (from DB character_prompt, fallback to file)
    global_prompt = await config_repo.get_active_prompt("character_prompt")
    if not global_prompt:
        try:
            char_path = constants.CHARACTER_PROMPT_PATH
            if os.path.exists(char_path):
                with open(char_path, "r", encoding="utf-8") as f:
                    global_prompt = f.read().strip()
        except Exception as e:
            logger.error(f"CHARACTER_PROMPT_FILE_FALLBACK_FAILED: {e}")

    if not global_prompt:
        global_prompt = constants.DEFAULT_SYSTEM_PROMPT

    # 2. Language instruction (from configurations_v2)
    language_instruction = await config_v2_repo.get_language_prompt(lang_code)

    # 3. Language intro (from configurations_v2, per persona per language)
    language_intro = ""
    if persona_id:
        language_intro = await config_v2_repo.get_language_intro(lang_code, persona_id)

    # 4. Personality (persona's system_prompt)
    personality = persona_prompt or ""

    # Assemble: global + language + intro + personality
    parts = [global_prompt]
    if language_instruction:
        parts.append(f"LANGUAGE:\n{language_instruction}")
    if language_intro:
        parts.append(f"INTRO (say this when user joins):\n{language_intro}")
    if personality:
        parts.append(f"WHAT YOU CAN HELP WITH:\n{personality}")

    full_prompt = "\n\n".join(parts)
    logger.info(f"PROMPT_LOADED: persona={persona_id} lang={lang_code} global={len(global_prompt)} lang_inst={len(language_instruction)} intro={len(language_intro)} personality={len(personality)} total={len(full_prompt)}")
    return full_prompt

async def cleanup_session(session_state, bot_session: BotSessionState, session_id: str,
                         transcript_manager: "TranscriptManager" = None,
                         chunk_state: "ChunkState" = None):
    """Comprehensive session cleanup to free memory."""
    threads_at_start = threading.active_count()
    logger.info(f"CLEANUP_START: session_id={session_id} threads={threads_at_start}")

    # 1. Cancel all tracked background tasks first (prevents them from holding references)
    try:
        await bot_session.cancel_background_tasks()
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: cancel_background_tasks failed: {e}")

    # 2. Stop queue processors
    try:
        await session_state.stop_queue_processor()
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: stop_queue_processor failed: {e}")

    # 3. Close TTS service (HTTP session + connector)
    if session_state.tts_service:
        try:
            await session_state.tts_service.close()
        except Exception as e:
            logger.error(f"CLEANUP_ERROR: tts_service.close failed: {e}")
        finally:
            session_state.tts_service = None

    # 4. Close STT service (WebSocket connections to Deepgram/OpenAI)
    if bot_session.stt_service:
        try:
            if hasattr(bot_session.stt_service, 'close'):
                await bot_session.stt_service.close()
            elif hasattr(bot_session.stt_service, 'cleanup'):
                await bot_session.stt_service.cleanup()
            logger.debug(f"CLEANUP: STT service closed session_id={session_id}")
        except Exception as e:
            logger.error(f"CLEANUP_ERROR: stt_service close failed: {e}")
        finally:
            bot_session.stt_service = None

    # 5. Close LLM service (HTTP client connections to Gemini/OpenAI)
    if bot_session.llm_service:
        try:
            if hasattr(bot_session.llm_service, 'close'):
                await bot_session.llm_service.close()
            elif hasattr(bot_session.llm_service, 'cleanup'):
                await bot_session.llm_service.cleanup()
            logger.debug(f"CLEANUP: LLM service closed session_id={session_id}")
        except Exception as e:
            logger.error(f"CLEANUP_ERROR: llm_service close failed: {e}")
        finally:
            bot_session.llm_service = None

    # 6. Clear LLM context aggregator (conversation history - can be several MB)
    try:
        if bot_session.context_aggregator:
            bot_session.context_aggregator = None
            logger.debug(f"CLEANUP: context aggregator released session_id={session_id}")
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: context_aggregator clear failed: {e}")

    # 7. Cleanup pipeline first (this releases transport and all native threads)
    # IMPORTANT: pipeline.cleanup() must be called BEFORE setting pipeline=None
    # This ensures native threads (WebRTC, aiortc, etc.) are properly released
    # Each session creates ~7-8 native threads that MUST be cleaned up to prevent thread leaks
    pipeline_cleanup_success = False
    try:
        if bot_session.pipeline:
            # Pipeline.cleanup() internally calls transport.cleanup() and releases all native resources
            if hasattr(bot_session.pipeline, 'cleanup'):
                await bot_session.pipeline.cleanup()
                pipeline_cleanup_success = True
                logger.debug(f"CLEANUP: pipeline.cleanup() called session_id={session_id}")
            bot_session.pipeline = None
            logger.debug(f"CLEANUP: pipeline released session_id={session_id}")
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: pipeline cleanup failed: {e}", exc_info=True)
        # Still try to release reference even if cleanup failed
        try:
            bot_session.pipeline = None
        except Exception:
            pass

    # 8. Cleanup transport explicitly (in case pipeline.cleanup() didn't handle it or failed)
    # Transport native threads must be released via transport.cleanup().
    # Without this, threads accumulate (~147 sessions worth = 1130 threads = ~7GB memory)
    try:
        if bot_session.transport:
            # Only call transport.cleanup() if pipeline.cleanup() didn't succeed
            # (pipeline.cleanup() typically calls transport.cleanup() internally)
            if not pipeline_cleanup_success:
                if hasattr(bot_session.transport, 'cleanup'):
                    await bot_session.transport.cleanup()
                    logger.debug(f"CLEANUP: transport.cleanup() called (pipeline cleanup failed) session_id={session_id}")
                elif hasattr(bot_session.transport, 'close'):
                    await bot_session.transport.close()
                    logger.debug(f"CLEANUP: transport.close() called (pipeline cleanup failed) session_id={session_id}")
            bot_session.transport = None
            logger.debug(f"CLEANUP: transport released session_id={session_id}")
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: transport cleanup failed: {e}", exc_info=True)
        # Still try to release reference even if cleanup failed
        try:
            bot_session.transport = None
        except Exception:
            pass

    # 9. Destroy VAD analyzer (ONNX model - frees C++ allocations)
    if bot_session.vad_analyzer:
        try:
            del bot_session.vad_analyzer
            logger.info(f"VAD_DESTROYED: session_id={session_id} (memory freed)")
        except Exception as e:
            logger.error(f"CLEANUP_ERROR: VAD destruction failed: {e}")
        finally:
            bot_session.vad_analyzer = None

    # 10. Clear all image base64 data (2-10MB per session)
    try:
        bot_session.clear_all_images()
        logger.debug(f"CLEANUP: images cleared session_id={session_id}")
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: clear_all_images failed: {e}")

    # 11. Clear audio buffer
    try:
        if bot_session.audio_buffer:
            bot_session.audio_buffer = None
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: audio_buffer clear failed: {e}")

    # 12. Clear participants dict
    try:
        bot_session.participants.clear()
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: participants clear failed: {e}")

    # 13. Fallback save if participant_left didn't fire
    if not bot_session.session_saved and (bot_session.generated_videos or bot_session.generated_ppts or bot_session.generated_music or bot_session.credits_used_plan > 0 or bot_session.credits_used_topup > 0):
        try:
            from src.repository.session_repository import session_repo
            videos_to_save = bot_session.get_generated_videos()
            ppts_to_save = bot_session.get_generated_ppts()
            music_to_save = bot_session.get_generated_music()
            await session_repo.save_session_complete(
                session_id=session_id,
                device_id=bot_session.device_id,
                session_data={
                    "character_id": bot_session.persona_id,
                    "image_generations_count": 0,
                    "image_edits_count": 0,
                    "duration_seconds": int(time.time() - bot_session.bot_start_time),
                    "started_at": int(bot_session.bot_start_time),
                    "ended_at": int(time.time()),
                    "status": "completed",
                    "metadata": {"source": bot_session.source or "app", "save_type": "cleanup_fallback"}
                },
                videos=videos_to_save,
                ppts=ppts_to_save,
                music=music_to_save
            )
            bot_session.session_saved = True
            logger.info(f"CLEANUP_FALLBACK_SAVED: session_id={session_id} persona={bot_session.persona_id} videos={len(videos_to_save)} ppts={len(ppts_to_save)} music={len(music_to_save)}")

            # Persist credits too
            if bot_session.credits_used_plan > 0 and bot_session.subscription_id:
                from src.repository.subscription_repository import subscription_repo
                await subscription_repo.add_usage(bot_session.subscription_id, bot_session.credits_used_plan)
            if bot_session.credits_used_topup > 0 and bot_session.content_type:
                from src.repository.topup_repository import topup_repo
                await topup_repo.deduct_credits(bot_session.device_id, bot_session.content_type, bot_session.credits_used_topup)
        except Exception as e:
            logger.error(f"CLEANUP_FALLBACK_SAVE_FAILED: session_id={session_id} error={e}")

    # Clear generated media lists
    try:
        bot_session.generated_videos.clear()
        bot_session.generated_ppts.clear()
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: generated media clear failed: {e}")

    # 14. Clear transcript messages (only after saved to S3)
    try:
        if transcript_manager:
            transcript_manager.clear()
            logger.debug(f"CLEANUP: transcript cleared session_id={session_id}")
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: transcript clear failed: {e}")

    # 15. Clear chunk state (only after manifest saved)
    try:
        if chunk_state:
            chunk_state.clear()
            logger.debug(f"CLEANUP: chunks cleared session_id={session_id}")
    except Exception as e:
        logger.error(f"CLEANUP_ERROR: chunk_state clear failed: {e}")

    threads_at_end = threading.active_count()
    logger.info(f"CLEANUP_END: session_id={session_id} threads_at_start={threads_at_start} threads_at_end={threads_at_end} delta={threads_at_end - threads_at_start}")

+        # Payments config (loaded from DB at session start)

+        self.payments_enabled: bool = True  # True = show payment popup, False = cooldown message

+        self.cooldown_minutes: int = 30     # default cooldown when payments off


Now load the payments config at session start in `main.py`:


logger.info(f"USAGE_STATE: session_id={session_id} device_id={device_id} persona={persona_id} content={bot_session.content_type} free={bot_session.free_remaining}/{free_count} plan={bot_session.plan_remaining} topup={bot_session.topup_remaining} transport={transport_type}")

        # Load payments config
        from src.repository.config_v2_repository import config_v2_repo
        payments_config = await config_v2_repo.get_payments_config()
        bot_session.payments_enabled = payments_config.get("payments_enabled", True)
        cooldowns = payments_config.get("cooldowns", {})
        bot_session.cooldown_minutes = cooldowns.get(bot_session.content_type, 30)

        logger.info(f"USAGE_STATE: session_id={session_id} device_id={device_id} persona={persona_id} content={bot_session.content_type} free={bot_session.free_remaining}/{free_count} plan={bot_session.plan_remaining} topup={bot_session.topup_remaining} payments={bot_session.payments_enabled} cooldown={bot_session.cooldown_minutes}m transport={transport_type}")


The file /home/ubuntu/pipecat/src/bot/core/main.py has been updated successfully.

        # Load payments config
        from src.repository.config_v2_repository import config_v2_repo
        payments_config = await config_v2_repo.get_payments_config()
        bot_session.payments_enabled = payments_config.get("payments_enabled", True)
        cooldowns = payments_config.get("cooldowns", {})
        bot_session.cooldown_minutes = cooldowns.get(bot_session.content_type, 30)

        logger.info(f"USAGE_STATE: session_id={session_id} device_id={device_id} persona={persona_id} content={bot_session.content_type} free={bot_session.free_remaining}/{free_count} plan={bot_session.plan_remaining} topup={bot_session.topup_remaining} payments={bot_session.payments_enabled} cooldown={bot_session.cooldown_minutes}m transport={transport_type}")

import asyncio
import threading
import time
from dotenv import load_dotenv
from loguru import logger

from pipecat.pipeline.runner import PipelineRunner

from src.config import constants
from src.bot.core.session_state import SessionState

from src.bot.core.config import load_system_instruction, cleanup_session, BotSessionState, ChunkState, TranscriptManager
from src.bot.core.pipeline import create_vad_analyzer, create_transport, create_daily_transport, create_stt_service, create_llm_service, create_tts_service, create_context_and_tools, create_processors, create_llm_text_processor, create_pipeline, create_task, create_gemini_live_service, create_gemini_live_context, create_gemini_live_pipeline
from src.bot.core.session_handlers import create_audio_data_handler, create_transcript_handler, create_client_message_handler, create_cleanup_if_no_user, create_idle_timeout_handler, create_app_review_trigger, create_first_client_connected_handler, create_client_disconnected_handler, create_client_ready_handler, create_daily_participant_joined_handler, create_daily_first_participant_joined_handler, create_daily_participant_left_handler
from src.bot.services.image_service import create_cf_image_handler
from src.bot.services.video_service import create_video_handler, create_video_from_image_handler
from src.bot.services.ppt_service import create_ppt_handler
from src.bot.services.music_service import create_music_handler

load_dotenv()

async def run_bot_webrtc_global_agent(webrtc_connection, device_id: str = None, session_id: str = None, version: str = None, character_id: str = None, app_key: str = None, language: str = None, source: str = None, persona_id: str = None, transport_type: str = "smallwebrtc", room_url: str = None, token: str = None, app_version: int = None):
    bot_start_time = time.time()
    device_id = device_id or "common"

    # Accept character_id as persona_id for backward compat
    if not persona_id and character_id:
        persona_id = character_id

    session_state = SessionState()
    await session_state.start_queue_processor()

    # Initialize bot session
    bot_session = BotSessionState(session_id, device_id, bot_start_time)
    bot_session.persona_id = persona_id
    bot_session.app_key = app_key
    bot_session.language = language
    bot_session.source = source
    bot_session.app_version = app_version

    # Load persona info
    persona_prompt = None
    persona_voice = None
    persona_tools = []
    if persona_id:
        from src.repository.persona_repository import persona_repo
        persona = await persona_repo.get_persona_by_id(persona_id)
        if persona:
            bot_session.persona_name = persona.get("name")
            bot_session.content_type = persona.get("content_type")
            persona_prompt = persona.get("system_prompt")
            persona_voice = persona.get("tts_voice")
            persona_tools = persona.get("tools", [])
            logger.info(f"SESSION_START: persona={persona_id} name={persona.get('name')} content_type={bot_session.content_type} tools={persona_tools} transport={transport_type}")
        else:
            logger.warning(f"PERSONA_NOT_FOUND: persona={persona_id}")
    else:
        logger.info(f"SESSION_START: no_persona_id device_id={device_id} transport={transport_type}")

    # Fetch device info early — use DB language instead of frontend-sent value
    if device_id != "common":
        from src.repository.device_repository import device_repo
        device_info = await device_repo.get_device(device_id)
        if device_info and device_info.get("language"):
            language = device_info["language"]
            bot_session.language = language
            logger.info(f"LANGUAGE_FROM_DB: device={device_id} language={language}")
    else:
        device_info = None

    # Load system instruction
    system_instruction = await load_system_instruction(
        language=language,
        persona_prompt=persona_prompt,
        persona_id=persona_id,
    )

    # Load usage state (persona-level free limits, content-level plan/topup)
    if persona_id and bot_session.content_type and device_id != "common":
        from src.service.usage_service import get_usage_state
        from src.repository.usage_repository import usage_repo
        from src.repository.plan_repository import plan_repo
        import hashlib

        device_region = (device_info.get("region") or "india") if device_info else "india"
        payment_provider = "google_play" if device_region == "global" else "razorpay"

        usage = await get_usage_state(device_id, bot_session.content_type, payment_provider)
        bot_session.plan_remaining = usage["plan_remaining"]
        bot_session.topup_remaining = usage["topup_remaining"]
        bot_session.subscription_id = usage["subscription_id"]
        bot_session.plan_content_type = usage.get("plan_content_type")

        if device_region == "india":
            # India: always use persona-level free limits (no change from old behavior)
            free_count = persona.get("free_count", 0) if persona else 0
            free_window = persona.get("free_window_hours", 24) if persona else 24
        else:
            # Global: use region_free_limits with A/B percentage
            # Group A (hash < ab_percentage): fully free, no limits
            # Group B (hash >= ab_percentage): gets free_count free, then paywall
            region_limits = await plan_repo.get_region_free_limits(device_region)
            region_limit = region_limits.get(bot_session.content_type)
            if region_limit:
                ab_pct = region_limit.get("ab_percentage", 100)
                device_h = int(hashlib.sha256(device_id.encode()).hexdigest(), 16) % 100
                if device_h < ab_pct:
                    # Group A: fully free — set unlimited
                    free_count = 999999
                    free_window = 24
                else:
                    # Group B: limited free
                    free_count = region_limit["free_count"]
                    free_window = region_limit["free_window_hours"]
            else:
                free_count = persona.get("free_count", 0) if persona else 0
                free_window = persona.get("free_window_hours", 24) if persona else 24

        if free_count > 0:
            if free_count == 999999:
                bot_session.free_remaining = 999999  # unlimited
            else:
                used_free = await usage_repo.count_persona_generations(device_id, persona_id, bot_session.content_type, free_window)
                bot_session.free_remaining = max(0, free_count - used_free)
        else:
            bot_session.free_remaining = 0

        logger.info(f"USAGE_STATE: session_id={session_id} device_id={device_id} persona={persona_id} content={bot_session.content_type} free={bot_session.free_remaining}/{free_count} plan={bot_session.plan_remaining} topup={bot_session.topup_remaining} transport={transport_type}")

    # Determine which tools to enable
    video_enabled = "generate_video" in persona_tools or "generate_video_from_image" in persona_tools

    ppt_enabled = "generate_ppt" in persona_tools
    music_enabled = "generate_music" in persona_tools
    bot_session.video_enabled = video_enabled

    use_gemini_live = language not in (None, "hi")
    bot_session.use_gemini_live = use_gemini_live
    logger.info(f"PIPELINE_MODE: language={language} use_gemini_live={use_gemini_live}")

    vad_analyzer = create_vad_analyzer(session_id)
    bot_session.vad_analyzer = vad_analyzer

    # Create services and store references IMMEDIATELY on bot_session after each creation.
    # This ensures cleanup_session can release them even if a later create_* call throws.
    if transport_type == "daily":
        transport = create_daily_transport(room_url, token, vad_analyzer)
        logger.info(f"TRANSPORT_CREATED: type=daily session_id={session_id}")
    else:
        transport = create_transport(webrtc_connection, vad_analyzer)
        logger.info(f"TRANSPORT_CREATED: type=smallwebrtc session_id={session_id}")
    bot_session.transport = transport
    bot_session.transport_type = transport_type

    rtvi, transcript, audio_buffer = create_processors()
    transcript_manager = TranscriptManager()
    chunk_state = ChunkState()

    if use_gemini_live:
        # Gemini Live: single service handles STT + LLM + TTS
        llm = create_gemini_live_service(system_instruction, include_video=video_enabled, include_ppt=ppt_enabled, include_music=music_enabled, language=language)
        bot_session.llm_service = llm

        context_aggregator = create_gemini_live_context()
        bot_session.context_aggregator = context_aggregator

        logger.info(f"GEMINI_LIVE_SESSION: session_id={session_id} language={language} persona={persona_id} transport={transport_type}")

        pipeline = create_gemini_live_pipeline(transport, llm, context_aggregator, rtvi, audio_buffer, transcript)
    else:
        # Traditional pipeline: separate STT → LLM → TTS (mobile app only, web always uses Gemini Live)
        stt = create_stt_service()
        bot_session.stt_service = stt

        llm = create_llm_service(system_instruction)
        bot_session.llm_service = llm

        tts = create_tts_service(speaker=persona_voice)
        session_state.tts_service = tts

        context_aggregator = create_context_and_tools(
            llm,
            system_instruction,
            include_video=video_enabled,
            include_ppt=ppt_enabled,
            include_music=music_enabled,
        )
        bot_session.context_aggregator = context_aggregator

        llm_text_processor = create_llm_text_processor()
        pipeline = create_pipeline(transport, stt, context_aggregator, rtvi, llm, tts, audio_buffer, transcript, llm_text_processor)

    bot_session.pipeline = pipeline
    task = create_task(pipeline, rtvi, bot_session)

    audio_buffer.event_handler("on_audio_data")(create_audio_data_handler(bot_session, chunk_state))

    trigger_app_review_if_eligible = create_app_review_trigger(bot_session, rtvi)

    # Register tools based on persona config or legacy defaults
    # timeout_secs=120: image/video generation can take 20-60s; default 10s causes premature timeout
    # cancel_on_interruption=False: user may speak while waiting; don't cancel generation
    _tool_handlers = {
        "cf_image": lambda: create_cf_image_handler(
            session_state, rtvi, bot_session, transcript_manager, task, llm, trigger_app_review_if_eligible
        ),
        "generate_video": lambda: create_video_handler(
            session_state, rtvi, bot_session, transcript_manager, task, llm
        ),
        "generate_video_from_image": lambda: create_video_from_image_handler(
            session_state, rtvi, bot_session, transcript_manager, task, llm
        ),
        "generate_ppt": lambda: create_ppt_handler(
            session_state, rtvi, bot_session, transcript_manager, task, llm
        ),
        "generate_music": lambda: create_music_handler(
            session_state, rtvi, bot_session, transcript_manager, task, llm
        ),
    }

    # Register tools from persona config (or default cf_image if no persona)
    tools_to_register = persona_tools if persona_tools else ["cf_image"]
    if video_enabled and not persona_tools:
        tools_to_register.extend(["generate_video", "generate_video_from_image"])

    for tool_name in tools_to_register:
        if tool_name in _tool_handlers:
            try:
                llm.register_function(tool_name, _tool_handlers[tool_name](), timeout_secs=120, cancel_on_interruption=False)
                logger.info(f"FUNCTION_REGISTERED: {tool_name}")
            except Exception as e:
                logger.error(f"Failed to register {tool_name} handler: {e}", exc_info=True)

    transcript.event_handler("on_transcript_update")(create_transcript_handler(bot_session, transcript_manager, rtvi))
    rtvi.event_handler("on_client_ready")(create_client_ready_handler())
    rtvi.event_handler("on_client_message")(create_client_message_handler(bot_session, rtvi, task, transcript_manager))

    cleanup_if_no_user = create_cleanup_if_no_user(bot_session, session_state, task)

    if transport_type == "daily":
        # Daily events: on_participant_joined, on_first_participant_joined, on_participant_left
        transport.event_handler("on_participant_joined")(create_daily_participant_joined_handler(bot_session, task))
        transport.event_handler("on_first_participant_joined")(create_daily_first_participant_joined_handler(bot_session, transport, rtvi, task, audio_buffer, chunk_state))
        transport.event_handler("on_participant_left")(create_daily_participant_left_handler(bot_session, session_state, transcript_manager, chunk_state, task))
    else:
        # SmallWebRTC events: on_client_connected, on_client_disconnected
        # NOTE: pipecat only supports one handler per event — the merged handler
        # includes both client_connected and first_client_connected logic.
        transport.event_handler("on_client_connected")(create_first_client_connected_handler(bot_session, rtvi, task, audio_buffer, chunk_state))
        transport.event_handler("on_client_disconnected")(create_client_disconnected_handler(bot_session, session_state, transcript_manager, chunk_state, task))
    task.event_handler("on_idle_timeout")(create_idle_timeout_handler(bot_session))

    runner = PipelineRunner(handle_sigint=False)

    bot_session.no_user_timeout_task = asyncio.create_task(cleanup_if_no_user())
    logger.info(f"NO_USER_TIMEOUT_STARTED: session_id={session_id} transport={transport_type} timeout={constants.NO_USER_TIMEOUT}s")

    try:
        await runner.run(task)
    except Exception as e:
        logger.error(f"PIPELINE_ERROR: session_id={session_id} transport={transport_type} error={e}", exc_info=True)
        await task.cancel()
    finally:
        threads_before = threading.active_count()
        logger.info(f"CLEANUP_TRIGGER: session_id={session_id} transport={transport_type} reason=runner_finally threads={threads_before}")
        logger.info(f"RUNNER_CLEANUP_START: session_id={session_id} transport={transport_type} threads={threads_before}")

        bot_session.cancel_timeout_task()
        if bot_session.no_user_timeout_task:
            try:
                await bot_session.no_user_timeout_task
            except asyncio.CancelledError:
                pass

        # Timeout cleanup to prevent stuck native threads from blocking forever.
        # If pipeline.cleanup() or transport.cleanup() hangs on C++ thread joins,
        # we force-release references after 30s so the task always completes.
        try:
            await asyncio.wait_for(
                cleanup_session(session_state, bot_session, session_id, transcript_manager, chunk_state),
                timeout=30.0
            )
            threads_after = threading.active_count()
            logger.info(f"CLEANUP_THREADS: session_id={session_id} transport={transport_type} threads_before={threads_before} threads_after={threads_after} delta={threads_after - threads_before}")
        except asyncio.TimeoutError:
            threads_after = threading.active_count()
            logger.error(
                f"CLEANUP_TIMEOUT: session_id={session_id} transport={transport_type} cleanup hung for 30s, forcing release threads_before={threads_before} threads_after={threads_after}"
            )
            # Force-null everything to drop references even if native cleanup hung
            bot_session.pipeline = None
            bot_session.transport = None
            bot_session.vad_analyzer = None
            bot_session.stt_service = None
            bot_session.llm_service = None
            bot_session.audio_buffer = None
            bot_session.context_aggregator = None

        logger.info(f"RUNNER_CLEANUP_COMPLETE: session_id={session_id} transport={transport_type} threads={threading.active_count()}")

-        logger.info(f"USAGE_STATE: session_id={session_id} device_id={device_id} persona={persona_id} content={bot_session.content_type} free={bot_session.free_remaining}/{free_count} plan={bot_session.plan_remaining} topup={bot_session.topup_remaining} transport={transport_type}")

+        from src.repository.config_v2_repository import config_v2_repo

+        payments_config = await config_v2_repo.get_payments_config()

+        bot_session.payments_enabled = payments_config.get("payments_enabled", True)

+        cooldowns = payments_config.get("cooldowns", {})

+        bot_session.cooldown_minutes = cooldowns.get(bot_session.content_type, 30)

+        logger.info(f"USAGE_STATE: session_id={session_id} device_id={device_id} persona={persona_id} content={bot_session.content_type} free={bot_session.free_remaining}/{free_count} plan={bot_session.plan_remaining} topup={bot_session.topup_remaining} payments={bot_session.payments_enabled} cooldown={bot_session.cooldown_minutes}m transport={transport_type}")


Now I need to make the agent say the cooldown message. Let me check how the image service handles the block:


160	                "status": "busy",
161	                "assistant_instruction": "Image generation/edit is already in progress. Do NOT call cf_image again. Wait silently. Reply in the user's language.",
162	            }))
163	            return
164
165	        # Check generation limits
166	        if not is_web:
167	            from src.bot.services.limit_checker import can_generate, send_limit_popup
168	            block_reason = can_generate(bot_session)
169	            if block_reason:
170	                await send_limit_popup(bot_session, rtvi_processor, block_reason)
171	                await params.result_callback(json.dumps({
172	                    "status": "limit_reached",
173	                    "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call cf_image again.",
174	                }))
175	                return
176
177	        if operation == "generate":
178	            return await _handle_generation(
179	                params, prompt, style, request_id, session_info,
180	                state, rtvi_processor, bot_session, transcript_manager,
181	                cf_client, replicate_client, task, trigger_review_callback,
182	                gemini_client=gemini_client,
183	                openai_client=openai_client,
184	            )

                "status": "busy",
                "assistant_instruction": "Image generation/edit is already in progress. Do NOT call cf_image again. Wait silently. Reply in the user's language.",
            }))
            return

        # Check generation limits
        if not is_web:
            from src.bot.services.limit_checker import can_generate, send_limit_popup
            block_reason = can_generate(bot_session)
            if block_reason:
                await send_limit_popup(bot_session, rtvi_processor, block_reason)
                await params.result_callback(json.dumps({
                    "status": "limit_reached",
                    "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call cf_image again.",
                }))
                return

        if operation == "generate":
            return await _handle_generation(
                params, prompt, style, request_id, session_info,
                state, rtvi_processor, bot_session, transcript_manager,
                cf_client, replicate_client, task, trigger_review_callback,
                gemini_client=gemini_client,
                openai_client=openai_client,
            )


The `assistant_instruction` is what the agent speaks. I need to change it for cooldown to include the wait time. Let me update the image service, video service, music service, and ppt service:


            block_reason = can_generate(bot_session)
            if block_reason:
                await send_limit_popup(bot_session, rtvi_processor, block_reason)
                await params.result_callback(json.dumps({
                    "status": "limit_reached",
                    "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call cf_image again.",
                }))

            block_reason = can_generate(bot_session)
            if block_reason:
                await send_limit_popup(bot_session, rtvi_processor, block_reason)
                if block_reason == "cooldown":
                    cooldown = bot_session.cooldown_minutes
                    instruction = f"User ka limit khatam ho gaya hai. Unko politely batao ki {cooldown} minutes baad try karo. Do NOT call cf_image again."
                else:
                    instruction = "User has reached their generation limit. Tell them politely. Do NOT call cf_image again."
                await params.result_callback(json.dumps({
                    "status": "limit_reached",
                    "assistant_instruction": instruction,
                }))


The file /home/ubuntu/pipecat/src/bot/services/image_service.py has been updated successfully.

            block_reason = can_generate(bot_session)
            if block_reason:
                await send_limit_popup(bot_session, rtvi_processor, block_reason)
                if block_reason == "cooldown":
                    cooldown = bot_session.cooldown_minutes
                    instruction = f"User ka limit khatam ho gaya hai. Unko politely batao ki {cooldown} minutes baad try karo. Do NOT call cf_image again."
                else:
                    instruction = "User has reached their generation limit. Tell them politely. Do NOT call cf_image again."
                await params.result_callback(json.dumps({
                    "status": "limit_reached",
                    "assistant_instruction": instruction,
                }))

"""Image generation and editing service using CF Flux + Replicate fallback."""

import asyncio
import base64
import time
import json
import uuid
from datetime import datetime
from loguru import logger

from pipecat.processors.frameworks.rtvi import RTVIProcessor
from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.services.llm_service import FunctionCallParams

from src.bot.core.session_state import SessionState
from src.core.storage import upload_to_s3
from src.config import constants
from src.repository.config_repository import config_repo

from src.service.cf_flux_client import CFFluxClient
from src.service.replicate_flux_client import ReplicateFluxClient
from src.service.gemini_image_client import GeminiImageClient
from src.service.openai_image_client import OpenAIImageClient
from src.service.firered_client import FireRedClient
from src.bot.core.rtvi_events import RTVIEventSender
from src.bot.core.helpers import handle_generation_error
from src.service.prompt_writer import write_image_prompt
from src.service.image_describer import describe_image_safe

from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from src.bot.core.config import BotSessionState, TranscriptManager

cf_image_function = FunctionSchema(
    name="cf_image",
    description=(
        "Generate or edit images. This is the ONLY image function. Use for ALL image tasks: "
        "1) Image generation - create new images from text descriptions (scenes, objects, landscapes, people, gods, deities, religious images, text banners, quotes, etc.) "
        "2) Image edits - modify uploaded images (background changes, filters, style changes, add celebrities, add text, object removal/addition, color adjustments, merging images) "
        "CRITICAL: Do NOT call this function again if generation/edit is already in progress."
    ),
    properties={
        "operation": {
            "type": "string",
            "enum": ["generate", "edit"],
            "description": "Operation type: 'generate' for text-to-image, 'edit' for image modification"
        },
        "prompt": {
            "type": "string",
            "description": "For 'generate': detailed image description. For 'edit': edit instruction (e.g., 'change background to beach', 'make it brighter', 'add sunglasses')"
        },
        "style": {
            "type": "string",
            "description": "Image style (e.g., 'photorealistic', 'artistic', 'cartoon', 'abstract')",
            "default": "photorealistic"
        }
    },
    required=["operation", "prompt"]
)

# ── Shared helpers ──────────────────────────────────────────────

async def _generate_with_fallback(cf_client, replicate_client, operation_id, cf_func, replicate_func):
    """Try CF client, fall back to Replicate on failure. Returns CFFluxResult."""
    try:
        return await cf_func()
    except Exception as cf_err:
        logger.warning(f"CF_{operation_id}_FAILED, retrying with Replicate: {cf_err}")
        try:
            result = await replicate_func()
            logger.info(f"REPLICATE_FALLBACK_SUCCESS: {operation_id}")
            return result
        except Exception as rep_err:
            logger.error(f"REPLICATE_FALLBACK_FAILED: {operation_id} {rep_err}")
            if not replicate_client.is_available():
                raise cf_err from rep_err
            raise rep_err from cf_err

def _get_provider(model: str) -> str:
    """Get provider name from model string."""
    if not model:
        return "cf_flux"
    if "replicate" in model:
        return "replicate"
    if "firered" in model:
        return "firered"
    return "cf_flux"

def _build_s3_key(bot_session, sub_folder: str, file_id: str, extension: str = ".png") -> str:
    """Build S3 key: device_id/session_id/sub_folder/file_id.ext"""
    device_id = bot_session.device_id or 'common'
    session_id = bot_session.session_id or 'unknown'
    return f"{device_id}/{session_id}/{sub_folder}/{file_id}{extension}"

def _update_session_image(bot_session, result, s3_url, image_id, description, is_edit=False):
    """Update session state with generated/edited image."""
    image_info = {
        "data": base64.b64encode(result.image_bytes).decode(),
        "mime_type": result.mime_type,
        "s3_url": s3_url,
        "image_id": image_id,
        "description": f"Edited: {description}" if is_edit else description,
        "generated_at": datetime.now().isoformat(),
        "source": "cf_edit" if is_edit else "cf_generate"
    }
    if is_edit:
        bot_session.set_edited_image(image_info)
    else:
        bot_session.set_generated_image(image_info)

async def _trigger_review(callback):
    """Safely trigger app review callback."""
    if callback:
        try:
            await callback()
        except Exception as e:
            logger.debug(f"App review trigger failed: {e}")

# ── Handler factory ─────────────────────────────────────────────

def create_cf_image_handler(
    state: SessionState,
    rtvi_processor: RTVIProcessor,
    bot_session: "BotSessionState",
    transcript_manager: "TranscriptManager",
    task=None,
    llm_service=None,
    trigger_review_callback=None
):
    """Create CF Flux image generation/editing handler."""
    is_web = (bot_session.source == "web")
    cf_client = CFFluxClient()
    replicate_client = ReplicateFluxClient()
    firered_client = FireRedClient()
    gemini_client = None
    openai_client = None

    async def handle_cf_image(params: FunctionCallParams):
        request_id = str(uuid.uuid4())[:8]
        session_info = bot_session.get_session_info()

        args = params.arguments
        operation = args.get("operation", "generate")
        prompt = args.get("prompt", "")
        style = args.get("style", "photorealistic")

        # Block duplicate calls while generation/edit is in progress
        if bot_session.is_processing:
            logger.warning(f"CF_IMAGE_DUPLICATE_BLOCKED: {session_info} request_id={request_id} operation={operation}")
            await params.result_callback(json.dumps({
                "status": "busy",
                "assistant_instruction": "Image generation/edit is already in progress. Do NOT call cf_image again. Wait silently. Reply in the user's language.",
            }))
            return

        # Check generation limits
        if not is_web:
            from src.bot.services.limit_checker import can_generate, send_limit_popup
            block_reason = can_generate(bot_session)
            if block_reason:
                await send_limit_popup(bot_session, rtvi_processor, block_reason)
                await params.result_callback(json.dumps({
                    "status": "limit_reached",
                    "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call cf_image again.",
                }))
                return

        if operation == "generate":
            return await _handle_generation(
                params, prompt, style, request_id, session_info,
                state, rtvi_processor, bot_session, transcript_manager,
                cf_client, replicate_client, task, trigger_review_callback,
                gemini_client=gemini_client,
                openai_client=openai_client,
            )
        else:
            return await _handle_edit(
                params, prompt, style, request_id, session_info,
                state, rtvi_processor, bot_session, transcript_manager,
                cf_client, replicate_client, task, llm_service, trigger_review_callback,
                gemini_client=gemini_client,
                openai_client=openai_client,
                firered_client=firered_client,
            )

    return handle_cf_image

# ── Generation ──────────────────────────────────────────────────

async def _handle_generation(
    params, prompt, style, request_id, session_info,
    state, rtvi_processor, bot_session, transcript_manager,
    cf_client, replicate_client, task, trigger_review_callback,
    gemini_client=None,
    openai_client=None,
):
    """Handle image generation."""
    image_id = f"image_{int(time.time())}"
    rtvi_events = RTVIEventSender(rtvi_processor, session_info)

    async def generate_image_async():
        start_time = time.time()
        logger.debug(f"CF_IMAGE_GENERATION_STARTED: image_id={image_id} {session_info} request_id={request_id}")

        try:
            await rtvi_events.send_generation_started(
                "image", image_id, constants.IMAGE_LOADER_TIME_SECONDS, description=prompt
            )

            if gemini_client:
                result = await gemini_client.generate_image(prompt=prompt)
            elif openai_client:
                result = await openai_client.generate_image(prompt=prompt)
            else:
                result = await _generate_with_fallback(
                    cf_client, replicate_client, image_id,
                    cf_func=lambda: cf_client.generate_image(prompt=prompt, width=1024, height=1024, num_steps=25, guidance=3.5),
                    replicate_func=lambda: replicate_client.generate_image(prompt=prompt)
                )

            s3_key = _build_s3_key(bot_session, "images", image_id)
            s3_url = await upload_to_s3(result.image_bytes, s3_key, result.mime_type)

            _update_session_image(bot_session, result, s3_url, image_id, prompt)

            provider = _get_provider(result.model)
            total_latency = time.time() - start_time
            logger.info(
                f"CF_IMAGE_GENERATED: image_id={image_id} {session_info} request_id={request_id} "
                f"provider={result.model} total_latency={total_latency:.3f}s api_latency={result.api_latency:.3f}s s3_url={s3_url}"
            )

            await rtvi_events.send_generation_complete(
                "image", image_id, s3_url, provider,
                model=result.model, description=prompt, style=style
            )
            transcript_manager.add_generation("image", {
                "image_id": image_id, "description": prompt, "s3_url": s3_url,
                "provider": provider, "model": result.model, "style": style,
                "latency_seconds": round(total_latency, 3),
                "api_latency_seconds": round(result.api_latency, 3)
            })
            await rtvi_events.send_generation_ended("image", image_id)
            await _trigger_review(trigger_review_callback)

            # Decrement credit
            from src.bot.services.limit_checker import use_credit
            use_credit(bot_session)

            # Notify LLM that image is ready (traditional pipeline runs in background)
            if task and not getattr(bot_session, 'use_gemini_live', False):
                from pipecat.frames.frames import LLMMessagesAppendFrame
                await task.queue_frames([
                    LLMMessagesAppendFrame(
                        messages=[{"role": "user", "content": "[System: Image generation is complete. The image is now on screen and ready for editing. If user asks to edit, use cf_image with operation=edit. Do NOT ask them to upload an image.]"}],
                        run_llm=False
                    )
                ])

            return {"status": "completed", "image_url": s3_url, "image_id": image_id, "provider": provider, "model": result.model,
                    "assistant_instruction": "Image is ready! Tell the user their image is generated and they can see it on screen. Be enthusiastic! Reply in the user's language."}

        except Exception as e:
            return await handle_generation_error(
                e, image_id, request_id, session_info, rtvi_events, task,
                generation_type="image", description=prompt,
                language=bot_session.language,
            )

    # Gemini Live: await result directly so it waits before speaking.
    # Traditional pipeline: queue in background (HTTP LLM blocks on function calls anyway).
    use_sync = getattr(bot_session, 'use_gemini_live', False)

    if use_sync:
        # Gemini Live: run synchronously, return actual result
        bot_session.is_processing = True
        try:
            result_data = await generate_image_async()
            await params.result_callback(json.dumps(
                result_data if result_data else {"status": "completed", "image_id": image_id}
            ))
        finally:
            bot_session.is_processing = False
    else:
        # Traditional pipeline: queue in background
        await state.add_to_queue("image", generate_image_async, {"image_id": image_id, "description": prompt, "style": style})
        await params.result_callback(json.dumps({
            "status": "in_progress",
            "assistant_instruction": "Image generation has STARTED for this request. Say a short waiting message in the user's language. CRITICAL: Do NOT call cf_image again for this same request. Reply in the user's language.",
            "image_id": image_id, "description": prompt
        }))

# ── Editing ─────────────────────────────────────────────────────

async def _handle_edit(
    params, prompt, style, request_id, session_info,
    state, rtvi_processor, bot_session, transcript_manager,
    cf_client, replicate_client, task, llm_service, trigger_review_callback,
    gemini_client=None,
    openai_client=None,
    firered_client=None,
):
    """Handle image editing."""
    edit_id = f"edit_{int(time.time())}"
    rtvi_events = RTVIEventSender(rtvi_processor, session_info)

    logger.info(f"CF_IMAGE_EDIT_FUNCTION_CALLED: {session_info} edit_id={edit_id} request_id={request_id}")

    # Validate image exists
    image_to_edit = bot_session.get_image_to_edit()
    if not image_to_edit:
        logger.warning(f"CF_IMAGE_EDIT_NO_IMAGE: {session_info} edit_id={edit_id}")
        from pipecat.frames.frames import TTSSpeakFrame
        no_image_msg = constants.ERROR_NO_IMAGE_FOR_EDIT if bot_session.language in (None, "hi") else constants.ERROR_NO_IMAGE_FOR_EDIT_EN
        await llm_service.push_frame(TTSSpeakFrame(no_image_msg))
        await params.result_callback(json.dumps({
            "status": "error", "message": "No image available. User needs to generate or upload an image first. Reply in the user's language.", "edit_id": edit_id
        }))
        return

    # Prepare image description
    use_sync = getattr(bot_session, 'use_gemini_live', False)
    image_source = image_to_edit.get("source", "unknown")

    if use_sync:
        # Gemini Live: skip describe entirely (saves 8-10s latency)
        image_description = None
        batch_size = 0
        logger.info(f"CF_IMAGE_EDIT_SKIP_DESCRIBE: edit_id={edit_id} (gemini_live)")
    else:
        try:
            image_batch = bot_session.get_image_batch()
            batch_size = bot_session.get_batch_size()
            has_batch = batch_size >= 2
            image_description = await _get_edit_description(image_to_edit, image_batch, has_batch, image_source, edit_id)
        except Exception as e:
            logger.error(f"CF_IMAGE_EDIT_DESCRIPTION_FAILED: edit_id={edit_id} error={e}")
            image_description = None
            batch_size = 0

    logger.info(f"CF_IMAGE_EDIT_STARTING: {session_info} edit_id={edit_id} image_source={image_source} batch_size={batch_size} has_description={bool(image_description)}")

    # Send RTVI started FIRST - user sees loader immediately
    await rtvi_events.send_generation_started("image_edit", edit_id, constants.IMAGE_LOADER_TIME_SECONDS, edit_instruction=prompt)

    # Gemini Live: skip prompt writer (adds latency, direct request is fine)
    if use_sync:
        enhanced_prompt = prompt
        logger.info(f"CF_PROMPT_DIRECT: edit_id={edit_id} len={len(prompt)} (gemini_live, skipped prompt_writer)")
    else:
        try:
            enhanced_prompt = await write_image_prompt(prompt, image_description)
            logger.info(f"CF_PROMPT_ENHANCED: edit_id={edit_id} original_len={len(prompt)} enhanced_len={len(enhanced_prompt)}")
        except Exception as e:
            logger.error(f"CF_PROMPT_ENHANCE_FAILED: edit_id={edit_id} error={e}")
            enhanced_prompt = prompt

    async def edit_image_async():
        start_time = time.time()
        try:
            image_data_b64 = image_to_edit.get("data")
            input_mime_type = image_to_edit.get("mime_type", "image/jpeg")

            # Prepare batch images
            current_batch = bot_session.get_image_batch()
            second_image_b64 = None
            is_batch_edit = len(current_batch) >= 2
            original_image_urls = []

            if is_batch_edit:
                image_data_b64 = current_batch[0].get("data")
                second_image_b64 = current_batch[1].get("data")
                original_image_urls = [img.get("s3_url") for img in current_batch if img.get("s3_url")]
                logger.info(f"CF_IMAGE_EDIT_BATCH: edit_id={edit_id} batch_size={len(current_batch)}")

            if gemini_client:
                result = await gemini_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt, reference_image_b64=second_image_b64)
            elif openai_client:
                result = await openai_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt)
            elif firered_client and not is_batch_edit and await firered_client.should_use(bot_session.device_id or "common"):
                # DB-configured % of devices: FireRed for single-image edits, CF as fallback
                try:
                    logger.info(f"FIRERED_EDIT_ATTEMPT: edit_id={edit_id} device={bot_session.device_id}")
                    result = await firered_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt)
                except Exception as fr_err:
                    logger.warning(f"FIRERED_EDIT_FAILED, falling back to CF: edit_id={edit_id} error={fr_err}")
                    result = await _generate_with_fallback(
                        cf_client, replicate_client, edit_id,
                        cf_func=lambda: cf_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt, reference_image_b64=second_image_b64, width=1024, height=1024, num_steps=25),
                        replicate_func=lambda: replicate_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt, reference_image_b64=second_image_b64)
                    )
            else:
                result = await _generate_with_fallback(
                    cf_client, replicate_client, edit_id,
                    cf_func=lambda: cf_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt, reference_image_b64=second_image_b64, width=1024, height=1024, num_steps=25),
                    replicate_func=lambda: replicate_client.edit_image(image_base64=image_data_b64, prompt=enhanced_prompt, reference_image_b64=second_image_b64)
                )

            s3_key = _build_s3_key(bot_session, "edits", edit_id)
            s3_url = await upload_to_s3(result.image_bytes, s3_key, result.mime_type)

            # Upload original image for single edits
            original_s3_url = None
            if not is_batch_edit:
                original_s3_url = await _upload_original_image(image_data_b64, input_mime_type, bot_session, edit_id, session_info)

            _update_session_image(bot_session, result, s3_url, edit_id, prompt, is_edit=True)

            provider = _get_provider(result.model)
            total_latency = time.time() - start_time
            logger.info(
                f"CF_IMAGE_EDITED: edit_id={edit_id} {session_info} request_id={request_id} "
                f"provider={provider} total_latency={total_latency:.3f}s api_latency={result.api_latency:.3f}s s3_url={s3_url}"
            )

            original_url_value = original_image_urls if is_batch_edit and original_image_urls else original_s3_url

            await rtvi_events.send_generation_complete(
                "image_edit", edit_id, s3_url, provider,
                edit_instruction=prompt, original_image_url=original_url_value
            )

            transcript_entry = {
                "edit_id": edit_id, "edit_instruction": prompt, "edited_s3_url": s3_url,
                "original_s3_url": original_url_value, "provider": provider,
                "latency_seconds": round(total_latency, 3), "api_latency_seconds": round(result.api_latency, 3),
            }
            if config_repo.is_image_describer_enabled() and image_description:
                transcript_entry["image_description"] = image_description
            if config_repo.is_prompt_writer_enabled() and enhanced_prompt != prompt:
                transcript_entry["gpt_prompt"] = enhanced_prompt
            transcript_manager.add_generation("image_edit", transcript_entry)

            await rtvi_events.send_generation_ended("image_edit", edit_id)
            await _trigger_review(trigger_review_callback)

            # Increment in-memory counter (edits count toward image limit)
            from src.bot.services.limit_checker import use_credit
            use_credit(bot_session)

            # Notify LLM that edit is ready (traditional pipeline runs in background)
            if task and not getattr(bot_session, 'use_gemini_live', False):
                from pipecat.frames.frames import LLMMessagesAppendFrame
                await task.queue_frames([
                    LLMMessagesAppendFrame(
                        messages=[{"role": "user", "content": "[System: Image edit is complete. The edited image is now on screen and ready for further editing. If user asks to edit again, use cf_image with operation=edit. Do NOT ask them to upload an image.]"}],
                        run_llm=False
                    )
                ])

            return {"status": "completed", "edit_id": edit_id, "edited_image_url": s3_url, "original_image_url": original_url_value, "provider": provider,
                    "assistant_instruction": "Edit is done! Tell the user their edited image is ready and they can see it on screen. Ask what other changes they want. Reply in the user's language."}

        except Exception as e:
            return await handle_generation_error(
                e, edit_id, request_id, session_info, rtvi_events, task,
                generation_type="edit",
                language=bot_session.language,
            )

    if use_sync:
        logger.info(f"CF_IMAGE_EDIT_SYNC: {session_info} edit_id={edit_id}")
        bot_session.is_processing = True
        try:
            result_data = await edit_image_async()
            await params.result_callback(json.dumps(
                result_data if result_data else {"status": "completed", "edit_id": edit_id}
            ))
        finally:
            bot_session.is_processing = False
    else:
        logger.info(f"CF_IMAGE_EDIT_QUEUEING: {session_info} edit_id={edit_id}")
        await state.add_to_queue("image", edit_image_async, {"edit_id": edit_id, "edit_instruction": prompt})
        await params.result_callback(json.dumps({
            "status": "in_progress",
            "assistant_instruction": "Image edit has STARTED for this request. Say a short waiting message in the user's language. CRITICAL: Do NOT call cf_image again for this same request. Reply in the user's language.",
            "edit_id": edit_id, "edit_instruction": prompt
        }))

# ── Edit helpers ────────────────────────────────────────────────

async def _get_edit_description(image_to_edit, image_batch, has_batch, image_source, edit_id):
    """Get image description for editing, calling Gemini if needed."""
    image_description = image_to_edit.get("description")
    needs_fresh = (
        not image_description
        or image_description.startswith("Edited:")
        or image_source in ["edit", "cf_edit", "generate", "cf_generate"]
    )

    if needs_fresh and not has_batch:
        image_data_b64 = image_to_edit.get("data")
        if image_data_b64:
            logger.info(f"CF_GETTING_FRESH_DESCRIPTION: edit_id={edit_id} source={image_source}")
            image_description = await describe_image_safe(image_data_b64, image_to_edit.get("mime_type", "image/jpeg"))
    elif has_batch and image_batch:
        descriptions = []
        for img in image_batch:
            desc = img.get("description")
            if desc and not desc.startswith("Edited:"):
                descriptions.append(desc)
            elif img.get("data"):
                fresh_desc = await describe_image_safe(img.get("data"), img.get("mime_type", "image/jpeg"))
                if fresh_desc:
                    descriptions.append(fresh_desc)
        image_description = " | ".join(descriptions) if descriptions else None

    return image_description

async def _upload_original_image(image_data_b64, mime_type, bot_session, edit_id, session_info):
    """Upload original image to S3 before editing."""
    try:
        extension = ".jpg" if "jpeg" in mime_type else ".png"
        s3_key = _build_s3_key(bot_session, "edits", f"{edit_id}_original", extension)
        return await upload_to_s3(base64.b64decode(image_data_b64), s3_key, mime_type)
    except Exception as e:
        logger.error(f"CF_IMAGE_EDIT_ORIGINAL_UPLOAD_FAILED: {session_info} edit_id={edit_id} error={e}")
        return None

+                if block_reason == "cooldown":

+                    cooldown = bot_session.cooldown_minutes

+                    instruction = f"User ka limit khatam ho gaya hai. Unko politely batao ki {cooldown} minutes baad try karo. Do NOT call cf_image again."

+                    instruction = "User has reached their generation limit. Tell them politely. Do NOT call cf_image again."

-                    "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call cf_image again.",

+                    "assistant_instruction": instruction,


Now video service:


205	        logger.info(f"{label}_REQUEST: {session_info} request_id={request_id} prompt={prompt[:50]}...")
206
207	        # Check limits
208	        from src.bot.services.limit_checker import can_generate as _can_gen, send_limit_popup
209	        block_reason = _can_gen(bot_session)
210	        if block_reason:
211	            await send_limit_popup(bot_session, rtvi_processor, block_reason)
212	            await params.result_callback(json.dumps({
213	                "status": "limit_reached",
214	                "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call generate_video again.",
215	            }))
216	            return
217
218	        hi = _is_hindi(bot_session.language)
219
220	        if not prompt:
221	            no_prompt_msg = "Kaisa video chahiye aapko? Please batao kya banau." if hi else "What kind of video would you like? Please describe what to create."
222	            await _reject(params, llm_service, session_info, label, "need_more_info", "no_prompt",
223	                no_prompt_msg,
224	                "User was asked for video description via voice. Wait for their response.")
225	            return
226
227	        if bot_session.is_processing:
228	            busy_msg = "Abhi ek video ya image already process ho rahi hai. Please thoda wait karo." if hi else "A video or image is already being processed. Please wait a moment."
229	            await _reject(params, llm_service, session_info, label, "busy", "already_processing",

        logger.info(f"{label}_REQUEST: {session_info} request_id={request_id} prompt={prompt[:50]}...")

        # Check limits
        from src.bot.services.limit_checker import can_generate as _can_gen, send_limit_popup
        block_reason = _can_gen(bot_session)
        if block_reason:
            await send_limit_popup(bot_session, rtvi_processor, block_reason)
            await params.result_callback(json.dumps({
                "status": "limit_reached",
                "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call generate_video again.",
            }))
            return

        hi = _is_hindi(bot_session.language)

        if not prompt:
            no_prompt_msg = "Kaisa video chahiye aapko? Please batao kya banau." if hi else "What kind of video would you like? Please describe what to create."
            await _reject(params, llm_service, session_info, label, "need_more_info", "no_prompt",
                no_prompt_msg,
                "User was asked for video description via voice. Wait for their response.")
            return

        if bot_session.is_processing:
            busy_msg = "Abhi ek video ya image already process ho rahi hai. Please thoda wait karo." if hi else "A video or image is already being processed. Please wait a moment."
            await _reject(params, llm_service, session_info, label, "busy", "already_processing",


        block_reason = _can_gen(bot_session)
        if block_reason:
            await send_limit_popup(bot_session, rtvi_processor, block_reason)
            await params.result_callback(json.dumps({
                "status": "limit_reached",
                "assistant_instruction": "User has reached their generation limit. Tell them politely. Do NOT call generate_video again.",
            }))
            return

        hi = _is_hindi(bot_session.language)

        if not prompt:
            no_prompt_msg = "Kaisa video chahiye aapko? Please batao kya banau." if hi else "What kind of video would you like? Please describe what to create."

        block_reason = _can_gen(bot_session)
        if block_reason:
            await send_limit_popup(bot_session, rtvi_processor, block_reason)
            if block_reason == "cooldown":
                cooldown = bot_session.cooldown_minutes
                instruction = f"User ka limit khatam ho gaya hai. Unko politely batao ki {cooldown} minutes baad try karo. Do NOT call generate_video again."
            else:
                instruction = "User has reached their generation limit. Tell them politely. Do NOT call generate_video again."
            await params.result_callback(json.dumps({
                "status": "limit_reached",
                "assistant_instruction": instruction,
            }))
            return

        hi = _is_hindi(bot_session.language)

        if not prompt:
            no_prompt_msg = "Kaisa video chahiye aapko? Please batao kya banau." if hi else "What kind of video would you like? Please describe what to create."