pusher.py

import struct
import asyncio
import contextvars
import json
import time
from collections import deque
from datetime import datetime, timezone
from typing import Dict, Optional, Set

from fastapi import APIRouter
from fastapi.websockets import WebSocketDisconnect, WebSocket
from starlette.websockets import WebSocketState

import database.conversations as conversations_db
from database import users as users_db
from database.redis_db import get_cached_user_geolocation
from models.conversation import Conversation
from models.conversation_enums import ConversationStatus
from utils.conversations.factory import deserialize_conversation
from models.geolocation import Geolocation
from utils.apps import is_audio_bytes_app_enabled
from utils.app_integrations import (
    trigger_realtime_integrations,
    trigger_realtime_audio_bytes,
    trigger_external_integrations,
)
from utils.conversations.location import async_get_google_maps_location
from utils.byok import set_byok_keys
from utils.conversations.process_conversation import process_conversation
from utils.executors import storage_executor
from utils.webhooks import (
    send_audio_bytes_developer_webhook,
    realtime_transcript_webhook,
    get_audio_bytes_webhook_seconds,
)
from utils.other.storage import upload_audio_chunk, upload_audio_chunks_batch
from utils.metrics import PUSHER_ACTIVE_WS_CONNECTIONS
from utils.speaker_identification import extract_speaker_samples
import logging

logger = logging.getLogger(__name__)

router = APIRouter()

# Constants for speaker sample extraction
SPEAKER_SAMPLE_PROCESS_INTERVAL = 15.0
SPEAKER_SAMPLE_MIN_AGE = 120.0

# Constants for private cloud sync
PRIVATE_CLOUD_SYNC_PROCESS_INTERVAL = 1.0
PRIVATE_CLOUD_CHUNK_DURATION = 60.0
PRIVATE_CLOUD_BATCH_MAX_AGE = 60.0  # seconds — flush batch if oldest chunk exceeds this age
PRIVATE_CLOUD_SYNC_MAX_RETRIES = 3

# Queue size limits
PRIVATE_CLOUD_QUEUE_MAX_SIZE = 20  # ~18MB/connection max (30 conns × 18MB = 540MB) — prevents OOM with headroom
SPEAKER_SAMPLE_QUEUE_WARN_SIZE = 100

# Constants for transcript queue batching
TRANSCRIPT_QUEUE_FLUSH_INTERVAL = 1.0  # seconds
TRANSCRIPT_QUEUE_WARN_SIZE = 50

# Constants for audio bytes queue
AUDIO_BYTES_QUEUE_WARN_SIZE = 20


async def _process_conversation_task(
    uid: str,
    conversation_id: str,
    language: str,
    websocket: WebSocket,
    byok_keys: Optional[Dict[str, str]] = None,
):
    """Process a conversation and send result back to _listen via websocket.

    `byok_keys` is forwarded from the listen service. When present, LLM and
    STT calls made inside process_conversation route through the user's own
    provider keys instead of Omi's env keys.
    """
    if byok_keys:
        set_byok_keys(byok_keys)
    try:
        conversation_data = conversations_db.get_conversation(uid, conversation_id)
        if not conversation_data:
            # Send error response
            response = {"conversation_id": conversation_id, "error": "conversation_not_found"}
            data = bytearray()
            data.extend(struct.pack("I", 201))
            data.extend(bytes(json.dumps(response), "utf-8"))
            await websocket.send_bytes(data)
            return

        conversation = deserialize_conversation(conversation_data)

        if conversation.status != ConversationStatus.processing:
            conversations_db.update_conversation_status(uid, conversation.id, ConversationStatus.processing)
            conversation.status = ConversationStatus.processing

        try:
            # Geolocation
            geolocation = get_cached_user_geolocation(uid)
            if geolocation:
                geolocation = Geolocation(**geolocation)
                conversation.geolocation = await async_get_google_maps_location(
                    geolocation.latitude, geolocation.longitude
                )

            # Run in default executor (not critical_executor) because process_conversation
            # is a coordinator that submits child tasks to critical_executor — nesting both
            # in the same pool causes deadlock under concurrent load.
            # Copy the current context (which holds BYOK keys) into the worker thread so
            # LLM client proxies see the right key when resolving per-request.
            loop = asyncio.get_running_loop()
            ctx = contextvars.copy_context()
            conversation = await loop.run_in_executor(
                None, lambda: ctx.run(process_conversation, uid, language, conversation)
            )
            messages = await trigger_external_integrations(uid, conversation)
        except Exception as e:
            logger.error(f"Error processing conversation: {e} {uid} {conversation_id}")
            conversations_db.set_conversation_as_discarded(uid, conversation.id)
            conversation.discarded = True
            messages = []

        # Send success response back (minimal - transcribe will fetch from DB)
        response = {"conversation_id": conversation_id, "success": True}
        data = bytearray()
        data.extend(struct.pack("I", 201))
        data.extend(bytes(json.dumps(response), "utf-8"))
        await websocket.send_bytes(data)

    except Exception as e:
        logger.error(f"Error in _process_conversation_task: {e} {uid} {conversation_id}")
        response = {"conversation_id": conversation_id, "error": str(e)}
        data = bytearray()
        data.extend(struct.pack("I", 201))
        data.extend(bytes(json.dumps(response), "utf-8"))
        try:
            await websocket.send_bytes(data)
        except Exception:
            pass


async def _websocket_util_trigger(
    websocket: WebSocket,
    uid: str,
    sample_rate: int = 8000,
):
    logger.info(f'_websocket_util_trigger {uid}')

    try:
        await websocket.accept()
    except RuntimeError as e:
        logger.error(e)
        await websocket.close(code=1011, reason="Dirty state")
        return

    websocket_active = True
    websocket_close_code = 1000

    # audio bytes
    audio_bytes_webhook_delay_seconds = get_audio_bytes_webhook_seconds(uid)
    audio_bytes_trigger_delay_seconds = 4
    has_audio_apps_enabled = is_audio_bytes_app_enabled(uid)
    private_cloud_sync_enabled = users_db.get_user_private_cloud_sync_enabled(uid)
    cached_protection_level = users_db.get_data_protection_level(uid) if private_cloud_sync_enabled else None

    # Track background tasks to cancel on cleanup (prevents memory leaks from fire-and-forget tasks)
    bg_tasks: Set[asyncio.Task] = set()

    def spawn(coro) -> asyncio.Task:
        """Create a tracked background task that will be cancelled on cleanup."""
        task = asyncio.create_task(coro)
        bg_tasks.add(task)

        def on_done(t):
            bg_tasks.discard(t)
            if t.cancelled():
                return
            exc = t.exception()
            if exc:
                logger.error(f"Unhandled exception in background task: {exc} {uid}")

        task.add_done_callback(on_done)
        return task

    # Bounded queues — prevent unbounded memory growth during backpressure
    speaker_sample_queue: deque = deque(maxlen=SPEAKER_SAMPLE_QUEUE_WARN_SIZE)
    transcript_queue: deque = deque(maxlen=TRANSCRIPT_QUEUE_WARN_SIZE)
    audio_bytes_queue: deque = deque(maxlen=AUDIO_BYTES_QUEUE_WARN_SIZE)

    # private_cloud_queue caps at PRIVATE_CLOUD_QUEUE_MAX_SIZE to prevent OOM kills.
    # An OOM kill loses ALL queued data for ALL users on the pod — dropping the oldest
    # chunk for one user is strictly better than killing the pod.
    private_cloud_queue: deque = deque(maxlen=PRIVATE_CLOUD_QUEUE_MAX_SIZE)
    audio_bytes_event = asyncio.Event()  # Signals when items are added for instant wake

    async def process_private_cloud_queue():
        """Background task that batches private cloud sync uploads by conversation_id.

        Chunks are accumulated per conversation and flushed when:
        - The batch reaches 60s of audio data, or
        - The oldest chunk in the batch exceeds PRIVATE_CLOUD_BATCH_MAX_AGE, or
        - The websocket disconnects (shutdown flush).
        """
        nonlocal websocket_active

        # Pending batches keyed by conversation_id
        pending: Dict[str, dict] = {}

        def _add_to_batch(chunk_info: dict):
            conv_id = chunk_info['conversation_id']
            if conv_id not in pending:
                pending[conv_id] = {
                    'data': bytearray(),
                    'conversation_id': conv_id,
                    'timestamp': chunk_info['timestamp'],  # oldest chunk timestamp
                    'queued_at': time.monotonic(),
                    'retries': 0,
                }
            batch = pending[conv_id]
            batch['data'].extend(chunk_info['data'])

        async def _flush_batch(conv_id: str):
            """Upload a batched chunk and update audio files."""
            batch = pending.pop(conv_id, None)
            if not batch or len(batch['data']) == 0:
                return
            chunk_data = bytes(batch['data'])
            del batch['data']  # free bytearray immediately — chunk_data holds the bytes copy
            timestamp = batch['timestamp']
            retries = batch.get('retries', 0)
            try:
                chunks_to_upload = [{'data': chunk_data, 'timestamp': timestamp}]
                loop = asyncio.get_running_loop()
                await loop.run_in_executor(
                    storage_executor, upload_audio_chunks_batch, chunks_to_upload, uid, conv_id, cached_protection_level
                )
                del chunks_to_upload
                try:
                    audio_files = await loop.run_in_executor(
                        storage_executor, conversations_db.create_audio_files_from_chunks, uid, conv_id
                    )
                    if audio_files:
                        await loop.run_in_executor(
                            storage_executor,
                            conversations_db.update_conversation,
                            uid,
                            conv_id,
                            {'audio_files': [af.dict() for af in audio_files]},
                        )
                except Exception as e:
                    logger.error(f"Error updating audio files: {e} {uid} {conv_id}")
            except Exception as e:
                if retries < PRIVATE_CLOUD_SYNC_MAX_RETRIES:
                    batch['retries'] = retries + 1
                    batch['data'] = bytearray(chunk_data)
                    batch['queued_at'] = time.monotonic()  # reset age so next retry waits ~60s
                    pending[conv_id] = batch
                    logger.error(f"Private cloud batch upload failed (retry {retries + 1}): {e} {uid} {conv_id}")
                else:
                    logger.info(
                        f"Private cloud batch upload failed after {PRIVATE_CLOUD_SYNC_MAX_RETRIES} retries, dropping: {e} {uid} {conv_id}"
                    )
            del chunk_data

        while websocket_active or len(private_cloud_queue) > 0 or len(pending) > 0:
            await asyncio.sleep(PRIVATE_CLOUD_SYNC_PROCESS_INTERVAL)

            # Drain queue into pending batches
            if private_cloud_queue:
                chunks_to_process = private_cloud_queue.copy()
                private_cloud_queue.clear()
                for chunk_info in chunks_to_process:
                    _add_to_batch(chunk_info)

            if not pending:
                continue

            now = time.monotonic()
            batch_size_threshold = sample_rate * 2 * PRIVATE_CLOUD_CHUNK_DURATION

            # Determine which conversations to flush
            conv_ids_to_flush = []
            for conv_id, batch in pending.items():
                batch_age = now - batch['queued_at']
                is_shutdown = not websocket_active
                is_size_ready = len(batch['data']) >= batch_size_threshold
                is_age_ready = batch_age >= PRIVATE_CLOUD_BATCH_MAX_AGE
                if is_shutdown or is_size_ready or is_age_ready:
                    conv_ids_to_flush.append(conv_id)

            for conv_id in conv_ids_to_flush:
                await _flush_batch(conv_id)

    async def process_speaker_sample_queue():
        """Background task that processes speaker sample extraction requests."""
        nonlocal websocket_active

        while websocket_active or len(speaker_sample_queue) > 0:
            await asyncio.sleep(SPEAKER_SAMPLE_PROCESS_INTERVAL)

            if not speaker_sample_queue:
                continue

            current_time = time.time()

            # Separate ready and pending requests
            ready_requests = []
            pending_requests = []

            for request in list(speaker_sample_queue):
                if current_time - request['queued_at'] >= SPEAKER_SAMPLE_MIN_AGE:
                    ready_requests.append(request)
                else:
                    pending_requests.append(request)

            # Keep pending requests in queue (rebuild deque with pending only)
            speaker_sample_queue.clear()
            speaker_sample_queue.extend(pending_requests)

            # Process ready requests (fire and forget)
            for request in ready_requests:
                person_id = request['person_id']
                conv_id = request['conversation_id']
                segment_ids = request['segment_ids']

                try:
                    await extract_speaker_samples(
                        uid=uid,
                        person_id=person_id,
                        conversation_id=conv_id,
                        segment_ids=segment_ids,
                        sample_rate=sample_rate,
                    )
                except Exception as e:
                    logger.error(f"Error extracting speaker samples: {e} {uid} {conv_id}")

    async def process_transcript_queue():
        """Batched consumer for transcript events (realtime integrations + webhooks)."""
        nonlocal websocket_active

        while websocket_active or len(transcript_queue) > 0:
            await asyncio.sleep(TRANSCRIPT_QUEUE_FLUSH_INTERVAL)

            if not transcript_queue:
                continue

            # Process batch
            batch = list(transcript_queue)
            transcript_queue.clear()

            for item in batch:
                segments = item['segments']
                memory_id = item['memory_id']
                try:
                    await trigger_realtime_integrations(uid, segments, memory_id)
                    await realtime_transcript_webhook(uid, segments)
                except Exception as e:
                    logger.error(f"Error processing transcript batch: {e} {uid}")

    async def process_audio_bytes_queue():
        """Event-driven consumer for audio bytes triggers (app integrations + webhooks)."""
        nonlocal websocket_active

        while websocket_active or len(audio_bytes_queue) > 0:
            # Wait for signal or check periodically for shutdown
            try:
                await asyncio.wait_for(audio_bytes_event.wait(), timeout=1.0)
            except asyncio.TimeoutError:
                continue  # Check websocket_active and queue on timeout

            audio_bytes_event.clear()

            if not audio_bytes_queue:
                continue

            # Process all queued items
            batch = list(audio_bytes_queue)
            audio_bytes_queue.clear()

            for item in batch:
                try:
                    if item['type'] == 'app':
                        await trigger_realtime_audio_bytes(uid, item['sample_rate'], item['data'])
                    elif item['type'] == 'webhook':
                        await send_audio_bytes_developer_webhook(uid, item['sample_rate'], item['data'])
                except Exception as e:
                    logger.error(f"Error processing audio bytes: {e} {uid}")

    async def receive_tasks():
        nonlocal websocket_active
        nonlocal websocket_close_code
        nonlocal speaker_sample_queue
        nonlocal transcript_queue
        nonlocal audio_bytes_queue

        audiobuffer = bytearray()
        trigger_audiobuffer = bytearray()
        private_cloud_sync_buffer = bytearray()
        private_cloud_chunk_start_time = None
        current_conversation_id = None

        try:
            while websocket_active:
                data = await websocket.receive_bytes()
                header_type = struct.unpack('<I', data[:4])[0]

                # Heartbeat (data-frame keepalive from backend to reset GKE ILB idle timer)
                if header_type == 100:
                    continue

                # Conversation ID
                if header_type == 103:
                    new_conversation_id = bytes(data[4:]).decode("utf-8")
                    # Flush private cloud buffer for the old conversation before switching
                    if (
                        private_cloud_sync_enabled
                        and current_conversation_id
                        and current_conversation_id != new_conversation_id
                        and len(private_cloud_sync_buffer) > 0
                    ):
                        if len(private_cloud_queue) >= PRIVATE_CLOUD_QUEUE_MAX_SIZE:
                            logger.warning(
                                f"private_cloud_queue full ({len(private_cloud_queue)}/{PRIVATE_CLOUD_QUEUE_MAX_SIZE}), "
                                f"dropping oldest chunk to prevent OOM {uid}"
                            )
                        private_cloud_queue.append(
                            {
                                'data': bytes(private_cloud_sync_buffer),
                                'conversation_id': current_conversation_id,
                                'timestamp': private_cloud_chunk_start_time or time.time(),
                                'retries': 0,
                            }
                        )
                        logger.info(
                            f"Flushed private cloud buffer on conversation switch: {len(private_cloud_sync_buffer)} bytes {uid}"
                        )
                        private_cloud_sync_buffer = bytearray()
                        private_cloud_chunk_start_time = None
                    current_conversation_id = new_conversation_id
                    logger.info(f"Pusher received conversation_id: {current_conversation_id} {uid}")
                    continue

                # Transcript - queue for batched processing
                if header_type == 102:
                    res = json.loads(bytes(data[4:]).decode("utf-8"))
                    segments = res.get('segments')
                    memory_id = res.get('memory_id')
                    if len(transcript_queue) >= TRANSCRIPT_QUEUE_WARN_SIZE:
                        logger.warning(f"Warning: transcript_queue size {len(transcript_queue)} {uid}")
                    # Use memory_id if available, otherwise use current_conversation_id for conversations
                    conversation_or_memory_id = memory_id or current_conversation_id
                    transcript_queue.append({'segments': segments, 'memory_id': conversation_or_memory_id})
                    continue

                # Process conversation request
                if header_type == 104:
                    res = json.loads(bytes(data[4:]).decode("utf-8"))
                    conversation_id = res.get('conversation_id')
                    language = res.get('language', 'en')
                    byok_keys = res.get('byok_keys') or None
                    if conversation_id:
                        logger.info(f"Pusher received process_conversation request: {conversation_id} {uid}")
                        spawn(_process_conversation_task(uid, conversation_id, language, websocket, byok_keys))
                    continue

                # Speaker sample extraction request - queue for background processing
                if header_type == 105:
                    res = json.loads(bytes(data[4:]).decode("utf-8"))
                    person_id = res.get('person_id')
                    conv_id = res.get('conversation_id')
                    segment_ids = res.get('segment_ids', [])
                    if person_id and conv_id and segment_ids:
                        if len(speaker_sample_queue) >= SPEAKER_SAMPLE_QUEUE_WARN_SIZE:
                            logger.warning(f"Warning: speaker_sample_queue size {len(speaker_sample_queue)} {uid}")
                        logger.info(
                            f"Queued speaker sample request: person={person_id}, {len(segment_ids)} segments {uid}"
                        )
                        speaker_sample_queue.append(
                            {
                                'person_id': person_id,
                                'conversation_id': conv_id,
                                'segment_ids': segment_ids,
                                'queued_at': time.time(),
                            }
                        )
                    continue

                # Audio bytes
                if header_type == 101:
                    # Parse: header(4) | timestamp(8 bytes double) | audio_data
                    buffer_start_timestamp = struct.unpack("d", data[4:12])[0]
                    audio_data = data[12:]

                    # Only accumulate audio buffers if there's a consumer (app trigger or webhook)
                    # Without this guard, buffers grow ~16KB/s indefinitely for users with no audio apps
                    if has_audio_apps_enabled:
                        trigger_audiobuffer.extend(audio_data)
                    if audio_bytes_webhook_delay_seconds is not None:
                        audiobuffer.extend(audio_data)

                    # Private cloud sync - queue chunks for background processing
                    if private_cloud_sync_enabled and current_conversation_id:
                        if private_cloud_chunk_start_time is None:
                            # Use timestamp from first buffer of this 5-second chunk
                            private_cloud_chunk_start_time = buffer_start_timestamp

                        private_cloud_sync_buffer.extend(audio_data)
                        # Queue chunk every PRIVATE_CLOUD_CHUNK_DURATION seconds
                        if len(private_cloud_sync_buffer) >= sample_rate * 2 * PRIVATE_CLOUD_CHUNK_DURATION:
                            if len(private_cloud_queue) >= PRIVATE_CLOUD_QUEUE_MAX_SIZE:
                                logger.warning(
                                    f"private_cloud_queue full ({len(private_cloud_queue)}/{PRIVATE_CLOUD_QUEUE_MAX_SIZE}), "
                                    f"dropping oldest chunk to prevent OOM {uid}"
                                )
                            private_cloud_queue.append(
                                {
                                    'data': bytes(private_cloud_sync_buffer),
                                    'conversation_id': current_conversation_id,
                                    'timestamp': private_cloud_chunk_start_time,
                                    'retries': 0,
                                }
                            )
                            private_cloud_sync_buffer = bytearray()
                            private_cloud_chunk_start_time = None

                    # Queue audio bytes triggers for batched processing
                    if (
                        has_audio_apps_enabled
                        and len(trigger_audiobuffer) > sample_rate * audio_bytes_trigger_delay_seconds * 2
                    ):
                        if len(audio_bytes_queue) >= AUDIO_BYTES_QUEUE_WARN_SIZE:
                            logger.warning(f"Warning: audio_bytes_queue size {len(audio_bytes_queue)} {uid}")
                        audio_bytes_queue.append(
                            {
                                'type': 'app',
                                'sample_rate': sample_rate,
                                'data': trigger_audiobuffer.copy(),
                            }
                        )
                        audio_bytes_event.set()  # Wake consumer immediately
                        trigger_audiobuffer = bytearray()
                    if (
                        audio_bytes_webhook_delay_seconds is not None
                        and len(audiobuffer) > sample_rate * audio_bytes_webhook_delay_seconds * 2
                    ):
                        if len(audio_bytes_queue) >= AUDIO_BYTES_QUEUE_WARN_SIZE:
                            logger.warning(f"Warning: audio_bytes_queue size {len(audio_bytes_queue)} {uid}")
                        audio_bytes_queue.append(
                            {
                                'type': 'webhook',
                                'sample_rate': sample_rate,
                                'data': audiobuffer.copy(),
                            }
                        )
                        audio_bytes_event.set()  # Wake consumer immediately
                        audiobuffer = bytearray()
                    continue

        except WebSocketDisconnect:
            logger.info("WebSocket disconnected")
        except Exception as e:
            logger.error(f'Could not process audio: error {e}')
            websocket_close_code = 1011
        finally:
            # Flush any remaining private cloud sync buffer before shutdown
            if private_cloud_sync_enabled and current_conversation_id and len(private_cloud_sync_buffer) > 0:
                if len(private_cloud_queue) >= PRIVATE_CLOUD_QUEUE_MAX_SIZE:
                    logger.warning(
                        f"private_cloud_queue full ({len(private_cloud_queue)}/{PRIVATE_CLOUD_QUEUE_MAX_SIZE}), "
                        f"dropping oldest chunk to prevent OOM {uid}"
                    )
                private_cloud_queue.append(
                    {
                        'data': bytes(private_cloud_sync_buffer),
                        'conversation_id': current_conversation_id,
                        'timestamp': private_cloud_chunk_start_time or time.time(),
                        'retries': 0,
                    }
                )
                logger.info(f"Flushed final private cloud buffer: {len(private_cloud_sync_buffer)} bytes {uid}")
            websocket_active = False

    try:
        PUSHER_ACTIVE_WS_CONNECTIONS.inc()
        receive_task = asyncio.create_task(receive_tasks())
        speaker_sample_task = asyncio.create_task(process_speaker_sample_queue())
        private_cloud_task = asyncio.create_task(process_private_cloud_queue())
        transcript_task = asyncio.create_task(process_transcript_queue())
        audio_bytes_task = asyncio.create_task(process_audio_bytes_queue())
        await asyncio.gather(
            receive_task,
            speaker_sample_task,
            private_cloud_task,
            transcript_task,
            audio_bytes_task,
        )

    except Exception as e:
        logger.error(f"Error during WebSocket operation: {e}")
    finally:
        websocket_active = False

        # Cancel all tracked background tasks to prevent memory leaks
        tasks_to_cancel = list(bg_tasks)
        for task in tasks_to_cancel:
            task.cancel()
        if tasks_to_cancel:
            await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
        bg_tasks.clear()

        PUSHER_ACTIVE_WS_CONNECTIONS.dec()

        if websocket.client_state == WebSocketState.CONNECTED:
            try:
                await websocket.close(code=websocket_close_code)
            except Exception as e:
                logger.error(f"Error closing WebSocket: {e}")


@router.websocket("/v1/trigger/listen")
async def websocket_endpoint_trigger(
    websocket: WebSocket,
    uid: str,
    sample_rate: int = 8000,
):
    await _websocket_util_trigger(websocket, uid, sample_rate)