Output:
Chunk ID: 0468af
Wall time: 0.0509 seconds
Process exited with code 0
Original token count: 2410
Output:

    Raises:
        RuntimeError: If runtime not initialized or streaming pipeline unavailable
    """
    if not is_initialized():
        raise RuntimeError("TTS runtime not initialized")

    runtime = get_runtime()

    if runtime.streaming_pipeline is None:
        raise RuntimeError("Streaming pipeline not available")

    # Import audio utils for WAV header
    from veena3modal.audio.utils import create_wav_header

    # Metrics tracking
    metrics = {
        "ttfb_ms": 0,
        "chunks_sent": 0,
        "total_bytes": 0,
        "audio_duration_seconds": 0.0,
        "text_chunked": False,
    }

    start_time = time.time()
    first_chunk_time = None
    total_pcm_bytes = 0
    sample_rate = 16000  # BiCodec sample rate
    active_requests = _with_inflight_delta(1)
    if _timing_enabled():
        logger.info(
            "tts_runtime_stream_start",
            extra={
                "speaker": speaker,
                "text_length": len(text),
                "request_inflight": active_requests,
                "chunking": enable_chunking,
                "request_id": request_id,
            },
        )

    try:
        # Check if we need text chunking
        from veena3modal.processing.long_text_processor import LongTextProcessor
        long_processor = LongTextProcessor(pipeline=runtime.pipeline)
        needs_chunking = enable_chunking and long_processor.should_chunk(text)

        if needs_chunking:
            # Chunked streaming with voice consistency (global token caching)
            metrics["text_chunked"] = True
            async for audio_chunk, chunk_metrics in _stream_chunked_text(
                runtime,
                long_processor,
                text,
                speaker,
                temperature,
                top_k,
                top_p,
                max_tokens,
                repetition_penalty,
                seed,
                sample_rate,
                request_id,
            ):
                if first_chunk_time is None:
                    first_chunk_time = time.time()
                    metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
                    # Create streaming WAV header (size=0 for unknown length)
                    wav_header = create_wav_header(sample_rate=sample_rate, data_size=0)
                    audio_chunk = wav_header + audio_chunk

                total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
                metrics["chunks_sent"] += 1
                metrics["total_bytes"] = total_pcm_bytes + 44  # Include header
                metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)
                if chunk_metrics:
                    metrics.update(chunk_metrics)

                yield audio_chunk, metrics
        else:
            # Simple streaming (no text chunking)
            # OPTIMIZATION: If we have pre-cached global tokens for this speaker,
            # use continuation mode to skip the ~110ms global token pre-roll.
            # The model jumps straight to semantic token generation.
            cached_globals = runtime.speaker_global_cache.get(speaker)

            if cached_globals:
                stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
                    speaker=speaker,
                    text=text,
                    global_ids=cached_globals,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                    request_id=request_id,
                )
            else:
                stream_gen = runtime.streaming_pipeline.generate_speech_stream_indic(
                    speaker=speaker,
                    text=text,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                    request_id=request_id,
                )

            async for audio_chunk in stream_gen:
                if first_chunk_time is None:
                    first_chunk_time = time.time()
                    metrics["ttfb_ms"] = int((first_chunk_time - start_time) * 1000)
                    wav_header = create_wav_header(sample_rate=sample_rate, data_size=0)
                    audio_chunk = wav_header + audio_chunk

                total_pcm_bytes += len(audio_chunk) - (44 if metrics["chunks_sent"] == 0 else 0)
                metrics["chunks_sent"] += 1
                metrics["total_bytes"] = total_pcm_bytes + 44
                metrics["audio_duration_seconds"] = total_pcm_bytes / (sample_rate * 2)
                yield audio_chunk, metrics
    finally:
        active_requests = _with_inflight_delta(-1)
        if _timing_enabled():
            logger.info(
                "tts_runtime_stream_end",
                extra={
                    "speaker": speaker,
                    "request_id": request_id,
                    "request_inflight": active_requests,
                    "chunks_sent": metrics.get("chunks_sent", 0),
                },
            )

async def _stream_chunked_text(
    runtime: TTSRuntime,
    long_processor,
    text: str,
    speaker: str,
    temperature: float,
    top_k: int,
    top_p: float,
    max_tokens: int,
    repetition_penalty: float,
    seed: Optional[int],
    sample_rate: int,
    request_id: Optional[str] = None,
) -> AsyncGenerator[Tuple[bytes, Dict[str, Any]], None]:
    """
    Internal helper: stream chunked text with voice consistency.

    Uses global token caching from first chunk to maintain voice across chunks.
    """
    from veena3modal.audio.crossfade import crossfade_bytes_int16

    # Chunk text using the long text processor's chunker
    chunks = long_processor.chunk_text(text)

    if not chunks:
        return

    captured_globals: Optional[List[int]] = None
    previous_chunk_tail: Optional[bytes] = None
    chunk_metrics = {"chunks_processed": 0}

    for i, chunk_text in enumerate(chunks):
        chunk_metrics["chunks_processed"] = i + 1

        if i == 0:
            # First chunk: capture global tokens for voice consistency
            async for audio_bytes, global_ids in runtime.streaming_pipeline.generate_speech_stream_indic_first_chunk(
                speaker=speaker,
                text=chunk_text,
                temperature=temperature,
                top_k=top_k,
                top_p=top_p,
                max_tokens=max_tokens,
                repetition_penalty=repetition_penalty,
                seed=seed,
                request_id=request_id,
            ):
                if captured_globals is None and global_ids:
                    captured_globals = global_ids

                # Crossfade with previous tail (inter-chunk stitching)
                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                    previous_chunk_tail,
                    audio_bytes,
                    sample_rate_hz=sample_rate,
                    crossfade_ms=50,
                )

                if to_emit:
                    yield to_emit, chunk_metrics
        else:
            # Continuation chunks: use captured globals for voice consistency
            if captured_globals is None:
                logger.warning(f"No captured globals for chunk {i+1}, using regular streaming")
                async for audio_bytes in runtime.streaming_pipeline.generate_speech_stream_indic(
                    speaker=speaker,
                    text=chunk_text,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                    request_id=request_id,
                ):
                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        audio_bytes,
                        sample_rate_hz=sample_rate,
                        crossfade_ms=50,
                    )
                    if to_emit:
                        yield to_emit, chunk_metrics
            else:
                # Use continuation method with cached globals
                async for audio_bytes in runtime.streaming_pipeline.generate_speech_stream_indic_continuation(
                    speaker=speaker,
                    text=chunk_text,
                    global_ids=captured_globals,
                    temperature=temperature,
                    top_k=top_k,
                    top_p=top_p,
                    max_tokens=max_tokens,
                    repetition_penalty=repetition_penalty,
                    seed=seed,
                    request_id=request_id,
                ):
                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        audio_bytes,
                        sample_rate_hz=sample_rate,
                        crossfade_ms=50,
                    )
                    if to_emit:
                        yield to_emit, chunk_metrics

    # Flush remaining tail
    if previous_chunk_tail:
        yield previous_chunk_tail, chunk_metrics


# Collaboration Mode: Default

You are now in Default mode. Any previous instructions for other modes (e.g. Plan mode) are no longer active.

Your active mode changes only when new developer instructions with a different `<collaboration_mode>...</collaboration_mode>` change it; user requests or tool descriptions do not change mode by themselves. Known mode names are Default and Plan.

## request_user_input availability

The `request_user_input` tool is unavailable in Default mode. If you call it while in Default mode, it will return an error.

If a decision is necessary and cannot be discovered from local context, ask the user directly. However, in Default mode you should strongly prefer executing the user's request rather than stopping to ask questions.

Summary: auto

## Skills
A skill is a set of local instructions to follow that is stored in a `SKILL.md` file. Below is the list of skills that can be used. Each entry includes a name, description, and file path so you can open the source for full instructions when using a specific skill.
### Available skills
- skill-creator: Guide for creating effective skills. This skill should be used when users want to create a new skill (or update an existing skill) that extends Codex's capabilities with specialized knowledge, workflows, or tool integrations. (file: /home/ubuntu/.codex/skills/.system/skill-creator/SKILL.md)
- skill-installer: Install Codex skills into $CODEX_HOME/skills from a curated list or a GitHub repo path. Use when a user asks to list installable skills, install a curated skill, or install a skill from another repo (including private repos). (file: /home/ubuntu/.codex/skills/.system/skill-installer/SKILL.md)
### How to use skills
- Discovery: The list above is the skills available in this session (name + description + file path). Skill bodies live on disk at the listed paths.
- Trigger rules: If the user names a skill (with `$SkillName` or plain text) OR the task clearly matches a skill's description shown above, you must use that skill for that turn. Multiple mentions mean use them all. Do not carry skills across turns unless re-mentioned.
- Missing/blocked: If a named skill isn't in the list or the path can't be read, say so briefly and continue with the best fallback.
- How to use a skill (progressive disclosure):
  1) After deciding to use a skill, open its `SKILL.md`. Read only enough to follow the workflow.
  2) When `SKILL.md` references relative paths (e.g., `scripts/foo.py`), resolve them relative to the skill directory listed above first, and only consider other paths if needed.
  3) If `SKILL.md` points to extra folders such as `references/`, load only the specific files needed for the request; don't bulk-load everything.
  4) If `scripts/` exist, prefer running or patching them instead of retyping large code blocks.
  5) If `assets/` or templates exist, reuse them instead of recreating from scratch.
- Coordination and sequencing:
  - If multiple skills apply, choose the minimal set that covers the request and state the order you'll use them.
  - Announce which skill(s) you're using and why (one short line). If you skip an obvious skill, say why.
- Context hygiene:
  - Keep context small: summarize long sections instead of pasting them; only load extra files when needed.
  - Avoid deep reference-chasing: prefer opening only files directly linked from `SKILL.md` unless you're blocked.
  - When variants exist (frameworks, providers, domains), pick only the relevant reference file(s) and note that choice.
- Safety and fallback: If a skill can't be applied cleanly (missing files, unclear instructions), state the issue, pick the next-best approach, and continue.


I’ve identified a few concrete gaps: missing `request_id` propagation in some runtime/API calls and a `logger` reference bug risk in streaming pipeline. I’ll now patch these plus add a dedicated stress benchmark driver for the optimized local path and concurrency logs.


cd /home/ubuntu/veenaModal && sed -n '1,320p' veena3modal/api/fastapi_app.py && sed -n '1,280p' veena3modal/api/websocket_handler.py


Output:
Chunk ID: 09c0e4
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 5049
Output:
"""
FastAPI ASGI app factory for Modal deployment.

Phase 1: expose `/v1/tts/generate` and `/v1/tts/health` while reusing existing
Veena3 inference/streaming code from the Django project.

Phase 2+: extract shared code into `veena3modal/shared/` and remove Django coupling.
"""

# NOTE: Intentionally NOT using `from __future__ import annotations` here
# to ensure FastAPI recognizes Request type annotations properly

import time
import uuid
from typing import Optional

# Structured logging (JSON formatted, request-scoped)
from veena3modal.shared.logging import (
    get_logger,
    set_request_context,
    clear_request_context,
    log_request_received,
    log_first_audio_emitted,
    log_request_completed,
    log_request_failed,
)

# Prometheus metrics
from veena3modal.shared.metrics import (
    record_request_received as metrics_request_received,
    record_request_completed as metrics_request_completed,
    record_request_failed as metrics_request_failed,
    record_ttfb,
    record_rtf,
    record_audio_duration,
    record_chunks_sent,
    get_metrics_text,
    get_content_type as get_metrics_content_type,
)

# Auth and rate limiting
from veena3modal.api.auth import (
    get_api_validator,
    extract_api_key,
    hash_api_key,
)
from veena3modal.api.rate_limiter import get_rate_limiter
from veena3modal.api.error_handlers import (
    ErrorCode,
    create_error_response,
    get_error_status,
    is_gpu_fault,
    handle_gpu_fault,
    FeatureFlags,
)

logger = get_logger(__name__)

# Sentence storage (fire-and-forget, non-blocking)
# Imported at module level for singleton access
_sentence_store = None

def _get_sentence_store():
    """Lazy load sentence store singleton."""
    global _sentence_store
    if _sentence_store is None:
        from veena3modal.services.sentence_store import get_sentence_store
        _sentence_store = get_sentence_store()
    return _sentence_store

# Version info (will be set dynamically after model load)
_APP_VERSION = "0.1.0"
_MODEL_VERSION: Optional[str] = None
_startup_time: Optional[float] = None

def get_model_version() -> str:
    """Return the current model version, or 'not_loaded' if not initialized."""
    return _MODEL_VERSION or "not_loaded"

def set_model_version(version: str) -> None:
    """Set the model version after loading."""
    global _MODEL_VERSION
    _MODEL_VERSION = version

def create_app():
    """
    Factory function to create FastAPI app.

    Import inside factory so local tooling can import this module without FastAPI installed.
    """
    from fastapi import FastAPI, HTTPException, Request
    from fastapi.responses import JSONResponse, Response

    global _startup_time
    _startup_time = time.time()

    app = FastAPI(
        title="Veena3 TTS (Modal)",
        version=_APP_VERSION,
        description="High-quality multilingual Text-to-Speech API with true streaming support.",
    )

    @app.get("/v1/tts/health")
    async def tts_health():
        """
        Health check endpoint.

        Returns:
            - status: "healthy" | "degraded" | "unhealthy"
            - model_loaded: whether TTS model is initialized
            - model_version: version string of loaded model
            - uptime_seconds: time since app startup
            - gpu_available: whether GPU is detected (best-effort check)
        """
        # Check model status via runtime
        from veena3modal.services import tts_runtime
        model_loaded = tts_runtime.is_initialized()

        # Use runtime's model version if available
        runtime = tts_runtime.get_runtime()
        model_version = runtime.model_version if runtime else get_model_version()

        # Check GPU availability (best-effort, non-blocking)
        gpu_available = False
        try:
            import torch
            gpu_available = torch.cuda.is_available()
        except ImportError:
            pass  # torch not installed, skip GPU check

        # Compute uptime
        uptime = time.time() - _startup_time if _startup_time else 0

        # Determine overall status
        if model_loaded and gpu_available:
            status = "healthy"
        elif model_loaded or gpu_available:
            status = "degraded"
        else:
            status = "degraded"

        return JSONResponse(
            content={
                "status": status,
                "model_loaded": model_loaded,
                "model_version": model_version,
                "uptime_seconds": round(uptime, 2),
                "gpu_available": gpu_available,
                "app_version": _APP_VERSION,
            },
            headers={
                "X-Model-Version": model_version,
                "X-App-Version": _APP_VERSION,
            },
        )

    @app.get("/v1/tts/metrics")
    async def tts_metrics():
        """
        Prometheus metrics endpoint.

        Returns metrics in Prometheus text format for scraping.
        """
        return Response(
            content=get_metrics_text(),
            media_type=get_metrics_content_type(),
        )

    @app.post("/v1/tts/generate")
    async def tts_generate(request: Request):
        """
        Generate speech from text.

        Supports streaming (M4) and non-streaming (M3) modes.
        Currently implements non-streaming WAV generation only.
        """
        from veena3modal.api.schemas import TTSGenerateRequest, AudioFormat
        from veena3modal.services import tts_runtime

        request_id = str(uuid.uuid4())
        start_time = time.time()

        # === AUTH CHECK ===
        if FeatureFlags.is_auth_enabled():
            # Extract API key from headers
            headers_dict = dict(request.headers)
            api_key = extract_api_key(headers_dict)

            if not api_key:
                return JSONResponse(
                    status_code=401,
                    content=create_error_response(
                        code=ErrorCode.INVALID_API_KEY,
                        message="API key required. Use 'Authorization: Bearer <key>' or 'X-API-Key: <key>' header.",
                        request_id=request_id,
                    ),
                    headers={"X-Request-ID": request_id},
                )

            # Validate API key
            validator = get_api_validator()
            key_hash = hash_api_key(api_key)
            auth_result = validator.validate(key_hash)

            if not auth_result.is_valid:
                status_code = 401 if auth_result.error_code == "INVALID_API_KEY" else 403
                return JSONResponse(
                    status_code=status_code,
                    content=create_error_response(
                        code=ErrorCode(auth_result.error_code),
                        message=auth_result.error_message or "Authentication failed",
                        request_id=request_id,
                    ),
                    headers={"X-Request-ID": request_id},
                )

            # === RATE LIMIT CHECK ===
            if FeatureFlags.is_rate_limiting_enabled():
                rate_limiter = get_rate_limiter()
                allowed, remaining, reset_after = rate_limiter.check(key_hash)

                if not allowed:
                    headers = rate_limiter.get_headers(allowed, remaining, reset_after)
                    headers["X-Request-ID"] = request_id
                    return JSONResponse(
                        status_code=429,
                        content=create_error_response(
                            code=ErrorCode.RATE_LIMIT_EXCEEDED,
                            message=f"Rate limit exceeded. Retry after {int(reset_after) + 1} seconds.",
                            request_id=request_id,
                        ),
                        headers=headers,
                    )

        # Parse and validate request body
        try:
            body = await request.json()
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}")

        try:
            req = TTSGenerateRequest(**body)
        except Exception as e:
            error_msg = str(e)
            # Extract validation error details
            return JSONResponse(
                status_code=400,
                content={
                    "error": {
                        "code": "VALIDATION_ERROR",
                        "message": error_msg,
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Check if runtime is initialized
        if not tts_runtime.is_initialized():
            return JSONResponse(
                status_code=503,
                content={
                    "error": {
                        "code": "MODEL_NOT_LOADED",
                        "message": "TTS model not initialized. Please wait for warmup.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Get resolved speaker
        speaker = req.get_resolved_speaker()

        # Get normalized text (apply text normalization if enabled)
        normalizer_func = None
        if req.normalize:
            try:
                from veena3modal.processing.text_normalizer import normalize_text
                normalizer_func = lambda t: normalize_text(t, verbose=req.normalize_verbose)
            except ImportError:
                logger.warning("TextNormalizer not available, skipping text normalization")

        text = req.get_normalized_text(normalizer_func=normalizer_func)

        # Set request context for structured logging
        set_request_context(request_id=request_id)

        # Log request received (lifecycle event + metrics)
        log_request_received(
            request_id=request_id,
            text_length=len(text),
            speaker=speaker,
            stream=req.stream,
            format=req.format,
        )
        metrics_request_received(speaker=speaker, stream=req.stream, format=req.format)

        # Streaming mode (M4)
        if req.stream:
            return await _handle_streaming_request(req, request_id, text, speaker, start_time)

        # Format encoding will be applied after generation (non-streaming only)
        # For streaming, only WAV is supported
        target_format = req.format

        try:
            # Get output sample rate - use_enum_values=True means req.output is already a string
            output_sr = req.output if hasattr(req, 'output') else "16khz"

            # Generate speech (non-streaming)
            if req.chunking:
                audio_bytes, metrics = await tts_runtime.generate_speech_chunked(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
"""
WebSocket handler for TTS streaming.

Provides bidirectional WebSocket communication for TTS:
- Client sends JSON request with text and parameters
- Server streams audio chunks as binary frames
- Client can send control messages (cancel, ping)

Protocol:
1. Client connects to /v1/tts/ws
2. Client sends JSON: {"text": "...", "speaker": "...", ...}
3. Server sends binary audio chunks
4. Server sends JSON: {"event": "complete", "metrics": {...}}
5. Connection closes

Control messages:
- {"event": "cancel"} - Stop generation
- {"event": "ping"} - Keep-alive (server responds with pong)
"""

import asyncio
import json
import time
import uuid
from typing import Optional, Dict, Any
from dataclasses import dataclass
from enum import Enum

from veena3modal.shared.logging import (
    get_logger,
    set_request_context,
    clear_request_context,
    log_request_received,
    log_first_audio_emitted,
    log_request_completed,
    log_request_failed,
)
from veena3modal.shared.metrics import (
    record_request_received as metrics_request_received,
    record_request_completed as metrics_request_completed,
    record_request_failed as metrics_request_failed,
    record_ttfb,
    record_rtf,
    record_audio_duration,
    record_chunks_sent,
)

logger = get_logger(__name__)

class WSMessageType(str, Enum):
    """WebSocket message types."""
    # Client -> Server
    REQUEST = "request"
    CANCEL = "cancel"
    PING = "ping"

    # Server -> Client
    AUDIO_CHUNK = "audio_chunk"  # Binary frame
    HEADER = "header"  # WAV header info
    PROGRESS = "progress"  # Progress update
    COMPLETE = "complete"  # Generation complete
    ERROR = "error"  # Error occurred
    PONG = "pong"  # Response to ping

@dataclass
class WSRequest:
    """Parsed WebSocket TTS request."""
    text: str
    speaker: str = "male1"
    temperature: float = 0.8
    top_k: int = 50
    top_p: float = 1.0
    max_tokens: int = 4096
    repetition_penalty: float = 1.05
    seed: Optional[int] = None
    chunking: bool = False
    normalize: bool = True
    format: str = "wav"
    sample_rate: int = 16000

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "WSRequest":
        """Parse request from JSON dict."""
        return cls(
            text=data.get("text", ""),
            speaker=data.get("speaker", "male1"),
            temperature=data.get("temperature", 0.8),
            top_k=data.get("top_k", 50),
            top_p=data.get("top_p", 1.0),
            max_tokens=data.get("max_tokens", 4096),
            repetition_penalty=data.get("repetition_penalty", 1.05),
            seed=data.get("seed"),
            chunking=data.get("chunking", False),
            normalize=data.get("normalize", True),
            format=data.get("format", "wav"),
            sample_rate=data.get("sample_rate", 16000),
        )

    def validate(self) -> Optional[str]:
        """Validate request, return error message if invalid."""
        if not self.text or not self.text.strip():
            return "Text is required and cannot be empty"
        if len(self.text) > 50000:
            return f"Text too long: {len(self.text)} chars (max 50000)"
        if self.format != "wav":
            return f"WebSocket streaming only supports WAV format (got {self.format})"
        return None

def create_error_message(code: str, message: str, request_id: str) -> str:
    """Create JSON error message."""
    return json.dumps({
        "event": WSMessageType.ERROR.value,
        "error": {
            "code": code,
            "message": message,
            "request_id": request_id,
        }
    })

def create_header_message(
    request_id: str,
    sample_rate: int,
    model_version: str,
) -> str:
    """Create JSON header message (sent before audio)."""
    return json.dumps({
        "event": WSMessageType.HEADER.value,
        "request_id": request_id,
        "sample_rate": sample_rate,
        "format": "wav",
        "model_version": model_version,
    })

def create_progress_message(
    chunks_sent: int,
    bytes_sent: int,
    elapsed_ms: int,
) -> str:
    """Create JSON progress message."""
    return json.dumps({
        "event": WSMessageType.PROGRESS.value,
        "chunks_sent": chunks_sent,
        "bytes_sent": bytes_sent,
        "elapsed_ms": elapsed_ms,
    })

def create_complete_message(
    request_id: str,
    metrics: Dict[str, Any],
) -> str:
    """Create JSON completion message."""
    return json.dumps({
        "event": WSMessageType.COMPLETE.value,
        "request_id": request_id,
        "metrics": metrics,
    })

async def handle_websocket_tts(websocket, api_key: Optional[str] = None):
    """
    Handle WebSocket TTS streaming connection.

    This is the main handler called from FastAPI WebSocket route.

    Args:
        websocket: FastAPI WebSocket instance
        api_key: Optional API key for authentication
    """
    from veena3modal.api.auth import get_api_validator, hash_api_key
    from veena3modal.api.rate_limiter import get_rate_limiter
    from veena3modal.api.error_handlers import FeatureFlags
    from veena3modal.api.schemas import resolve_speaker_name
    from veena3modal.services import tts_runtime

    request_id = str(uuid.uuid4())
    set_request_context(request_id=request_id)

    # Accept the WebSocket connection
    await websocket.accept()
    logger.info(f"WebSocket connection accepted: {request_id}")

    # Track cancellation
    cancelled = False

    try:
        # === AUTH CHECK (if enabled) ===
        if FeatureFlags.is_auth_enabled() and api_key:
            validator = get_api_validator()
            key_hash = hash_api_key(api_key)
            auth_result = validator.validate(key_hash)

            if not auth_result.is_valid:
                await websocket.send_text(create_error_message(
                    code=auth_result.error_code or "AUTH_FAILED",
                    message=auth_result.error_message or "Authentication failed",
                    request_id=request_id,
                ))
                await websocket.close(code=1008)  # Policy violation
                return

            # Rate limit check
            if FeatureFlags.is_rate_limiting_enabled():
                rate_limiter = get_rate_limiter()
                allowed, remaining, reset_after = rate_limiter.check(key_hash)
                if not allowed:
                    await websocket.send_text(create_error_message(
                        code="RATE_LIMIT_EXCEEDED",
                        message=f"Rate limit exceeded. Retry after {int(reset_after) + 1} seconds.",
                        request_id=request_id,
                    ))
                    await websocket.close(code=1008)
                    return

        # === WAIT FOR REQUEST ===
        try:
            # Wait for client to send request (with timeout)
            raw_message = await asyncio.wait_for(
                websocket.receive_text(),
                timeout=30.0
            )
            message = json.loads(raw_message)
        except asyncio.TimeoutError:
            await websocket.send_text(create_error_message(
                code="TIMEOUT",
                message="No request received within 30 seconds",
                request_id=request_id,
            ))
            await websocket.close(code=1000)
            return
        except json.JSONDecodeError as e:
            await websocket.send_text(create_error_message(
                code="INVALID_JSON",
                message=f"Invalid JSON: {e}",
                request_id=request_id,
            ))
            await websocket.close(code=1003)  # Unsupported data
            return

        # Handle control messages
        event_type = message.get("event", "request")
        if event_type == "ping":
            await websocket.send_text(json.dumps({"event": "pong"}))
            # Wait for actual request
            raw_message = await websocket.receive_text()
            message = json.loads(raw_message)

        # === PARSE REQUEST ===
        req = WSRequest.from_dict(message)
        validation_error = req.validate()
        if validation_error:
            await websocket.send_text(create_error_message(
                code="VALIDATION_ERROR",
                message=validation_error,
                request_id=request_id,
            ))
            await websocket.close(code=1003)
            return

        # Resolve speaker
        speaker = resolve_speaker_name(req.speaker)

        # Log request
        log_request_received(
            request_id=request_id,
            text_length=len(req.text),
            speaker=speaker,
            stream=True,
            format="wav",
        )
        metrics_request_received(speaker=speaker, stream=True, format="wav")

        # === CHECK MODEL ===
        if not tts_runtime.is_initialized():
            await websocket.send_text(create_error_message(


Summary: auto


cd /home/ubuntu/veenaModal && sed -n '280,620p' veena3modal/api/fastapi_app.py && sed -n '280,620p' veena3modal/api/websocket_handler.py


Output:
Chunk ID: b66dbf
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 5655
Output:
        normalizer_func = None
        if req.normalize:
            try:
                from veena3modal.processing.text_normalizer import normalize_text
                normalizer_func = lambda t: normalize_text(t, verbose=req.normalize_verbose)
            except ImportError:
                logger.warning("TextNormalizer not available, skipping text normalization")

        text = req.get_normalized_text(normalizer_func=normalizer_func)

        # Set request context for structured logging
        set_request_context(request_id=request_id)

        # Log request received (lifecycle event + metrics)
        log_request_received(
            request_id=request_id,
            text_length=len(text),
            speaker=speaker,
            stream=req.stream,
            format=req.format,
        )
        metrics_request_received(speaker=speaker, stream=req.stream, format=req.format)

        # Streaming mode (M4)
        if req.stream:
            return await _handle_streaming_request(req, request_id, text, speaker, start_time)

        # Format encoding will be applied after generation (non-streaming only)
        # For streaming, only WAV is supported
        target_format = req.format

        try:
            # Get output sample rate - use_enum_values=True means req.output is already a string
            output_sr = req.output if hasattr(req, 'output') else "16khz"

            # Generate speech (non-streaming)
            if req.chunking:
                audio_bytes, metrics = await tts_runtime.generate_speech_chunked(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    sample_rate=req.sample_rate,
                    output_sample_rate=output_sr,
                )
            else:
                audio_bytes, metrics = await tts_runtime.generate_speech(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    output_sample_rate=output_sr,
                )

            if audio_bytes is None:
                return JSONResponse(
                    status_code=500,
                    content={
                        "error": {
                            "code": "GENERATION_FAILED",
                            "message": "Audio generation failed. Check logs for details.",
                            "request_id": request_id,
                        }
                    },
                    headers={"X-Request-ID": request_id},
                )

            # Calculate RTF
            total_time = time.time() - start_time
            audio_duration = metrics.get("audio_duration_seconds", 0)
            rtf = total_time / audio_duration if audio_duration > 0 else 0

            # Build response headers
            runtime = tts_runtime.get_runtime()
            model_version = runtime.model_version if runtime else get_model_version()

            # Use actual output sample rate from metrics
            output_sample_rate = metrics.get("output_sample_rate", req.sample_rate)

            headers = {
                "X-Request-ID": request_id,
                "X-Model-Version": model_version,
                "X-Format": "wav",
                "X-Sample-Rate": str(output_sample_rate),
                "X-Stream": "false",
                "X-Audio-Bytes": str(len(audio_bytes)),
                "X-Audio-Seconds": f"{audio_duration:.2f}",
                "X-TTFB-ms": str(metrics.get("ttfb_ms", 0)),
                "X-RTF": f"{rtf:.3f}",
                "X-Text-Chunked": str(metrics.get("text_chunked", False)).lower(),
                "X-SR-Applied": str(metrics.get("sr_applied", False)).lower(),
            }

            if req.seed is not None:
                headers["X-Seed"] = str(req.seed)

            # Fire-and-forget sentence storage (non-blocking, after generation)
            sentence_store = _get_sentence_store()
            sentence_store.store_fire_and_forget(
                request_id=request_id,
                text=text,
                speaker=speaker,
                stream=False,
                format=req.format,
                temperature=req.temperature,
                top_k=req.top_k,
                top_p=req.top_p,
                max_tokens=req.max_tokens,
                repetition_penalty=req.repetition_penalty,
                seed=req.seed,
                text_chunked=metrics.get("text_chunked", False),
                ttfb_ms=metrics.get("ttfb_ms"),
                audio_duration_seconds=audio_duration,
            )

            # Log successful completion (lifecycle event + metrics)
            total_duration_ms = int(total_time * 1000)
            log_request_completed(
                request_id=request_id,
                status_code=200,
                total_duration_ms=total_duration_ms,
                audio_duration_seconds=audio_duration,
                rtf=rtf,
                chunks_sent=1,
            )
            metrics_request_completed(
                status_code=200,
                duration_seconds=total_time,
                speaker=speaker,
                stream=False,
            )
            record_ttfb(ttfb_seconds=metrics.get("ttfb_ms", 0) / 1000, speaker=speaker, stream=False)
            record_rtf(rtf=rtf, speaker=speaker)
            record_audio_duration(duration_seconds=audio_duration, speaker=speaker)

            clear_request_context()

            # Apply format encoding if not WAV
            media_type = "audio/wav"
            if target_format != AudioFormat.WAV.value:
                try:
                    from veena3modal.audio.encoder import (
                        OpusEncoder, MP3Encoder, MuLawEncoder, FLACEncoder
                    )

                    # Get sample rate from output (after potential SR)
                    sample_rate = metrics.get("output_sample_rate", 16000)

                    # Strip WAV header to get PCM data
                    pcm_data = audio_bytes[44:] if len(audio_bytes) > 44 else audio_bytes

                    if target_format == AudioFormat.OPUS.value:
                        encoder = OpusEncoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/opus"
                    elif target_format == AudioFormat.MP3.value:
                        encoder = MP3Encoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/mpeg"
                    elif target_format == AudioFormat.MULAW.value:
                        encoder = MuLawEncoder()
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/x-wav"
                    elif target_format == AudioFormat.FLAC.value:
                        encoder = FLACEncoder(sample_rate=sample_rate)
                        audio_bytes = encoder.encode(pcm_data, sample_rate)
                        media_type = "audio/flac"

                    headers["X-Audio-Bytes"] = str(len(audio_bytes))
                    logger.info(f"Encoded audio to {target_format}: {len(audio_bytes)} bytes")
                except ImportError as e:
                    logger.warning(f"Audio encoder not available: {e}")
                    # Fall back to WAV
                except Exception as e:
                    logger.error(f"Format encoding failed: {e}")
                    # Fall back to WAV

            return Response(
                content=audio_bytes,
                media_type=media_type,
                headers=headers,
            )

        except Exception as e:
            logger.exception(f"TTS generation error: {e}")

            # Log failure (lifecycle event + metrics)
            log_request_failed(
                request_id=request_id,
                status_code=500,
                error_code="INTERNAL_ERROR",
                error_message=str(e),
            )
            metrics_request_failed(
                status_code=500,
                error_code="INTERNAL_ERROR",
                speaker=speaker,
            )
            clear_request_context()

            return JSONResponse(
                status_code=500,
                content={
                    "error": {
                        "code": "INTERNAL_ERROR",
                        "message": str(e),
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

    async def _handle_streaming_request(req, request_id: str, text: str, speaker: str, start_time: float):
        """
        Handle streaming TTS request.

        Returns a StreamingResponse that yields WAV header + PCM chunks.
        True streaming: first bytes sent ASAP before full audio is generated.
        """
        from fastapi.responses import StreamingResponse
        from veena3modal.api.schemas import AudioFormat
        from veena3modal.services import tts_runtime

        # Check if streaming pipeline is available
        runtime = tts_runtime.get_runtime()
        if runtime is None or runtime.streaming_pipeline is None:
            return JSONResponse(
                status_code=503,
                content={
                    "error": {
                        "code": "STREAMING_UNAVAILABLE",
                        "message": "Streaming pipeline not initialized.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Non-WAV formats not implemented yet (M5)
        if req.format != AudioFormat.WAV.value:
            return JSONResponse(
                status_code=501,
                content={
                    "error": {
                        "code": "FORMAT_NOT_IMPLEMENTED",
                        "message": f"Streaming format '{req.format}' not yet implemented. Use format=wav.",
                        "request_id": request_id,
                    }
                },
                headers={"X-Request-ID": request_id},
            )

        # Set request context for structured logging
        set_request_context(request_id=request_id)

        # Create async generator for streaming response
        async def audio_stream_generator():
            """Yields audio chunks as they're generated."""
            final_metrics = {}
            first_chunk_sent = False
            chunks_count = 0
            try:
                async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    top_k=req.top_k,
                    top_p=req.top_p,
                    max_tokens=req.max_tokens,
                    repetition_penalty=req.repetition_penalty,
                    seed=req.seed,
                    enable_chunking=req.chunking,
                ):
                    final_metrics = metrics
                    chunks_count += 1

                    # Fire-and-forget sentence storage after first chunk (don't block TTFB)
                    if not first_chunk_sent:
                        first_chunk_sent = True
                        ttfb_ms = metrics.get("ttfb_ms", 0)

                        # Log first audio emitted (TTFB marker)
                        log_first_audio_emitted(
                            request_id=request_id,
                            ttfb_ms=ttfb_ms,
                            chunk_size_bytes=len(audio_chunk),
                        )
                        record_ttfb(ttfb_seconds=ttfb_ms / 1000, speaker=speaker, stream=True)

                        sentence_store = _get_sentence_store()
                        sentence_store.store_fire_and_forget(
                            request_id=request_id,
                            text=text,
                            speaker=speaker,
                            stream=True,
                            format=req.format,
                            temperature=req.temperature,
                            top_k=req.top_k,
                            top_p=req.top_p,
                            max_tokens=req.max_tokens,
                            repetition_penalty=req.repetition_penalty,
                            seed=req.seed,
                            text_chunked=metrics.get("text_chunked", False),
                            ttfb_ms=ttfb_ms,
                            # audio_duration_seconds not known yet for streaming
                        )

                    yield audio_chunk

                # Stream completed - log metrics
                total_time = time.time() - start_time
                audio_duration = final_metrics.get("audio_duration_seconds", 0)
                rtf = total_time / audio_duration if audio_duration > 0 else 0

                log_request_completed(
                    request_id=request_id,
                    status_code=200,
                    total_duration_ms=int(total_time * 1000),
                    audio_duration_seconds=audio_duration,
                    rtf=rtf,
                    chunks_sent=chunks_count,
                )
                metrics_request_completed(
                    status_code=200,
                    duration_seconds=total_time,
                    speaker=speaker,
                    stream=True,
                )
                record_rtf(rtf=rtf, speaker=speaker)
                record_audio_duration(duration_seconds=audio_duration, speaker=speaker)
                record_chunks_sent(chunks=chunks_count, speaker=speaker)
                clear_request_context()
            except Exception as e:
            await websocket.send_text(create_error_message(
                code="MODEL_NOT_LOADED",
                message="TTS model not initialized",
                request_id=request_id,
            ))
            await websocket.close(code=1011)  # Internal error
            return

        runtime = tts_runtime.get_runtime()
        model_version = runtime.model_version if runtime else "unknown"

        # === SEND HEADER INFO ===
        await websocket.send_text(create_header_message(
            request_id=request_id,
            sample_rate=req.sample_rate,
            model_version=model_version,
        ))

        # === STREAM AUDIO ===
        start_time = time.time()
        chunks_sent = 0
        total_bytes = 0
        first_chunk_time = None
        final_metrics = {}

        # Start a task to listen for cancel messages
        async def listen_for_cancel():
            nonlocal cancelled
            try:
                while not cancelled:
                    try:
                        msg = await asyncio.wait_for(
                            websocket.receive_text(),
                            timeout=0.1
                        )
                        data = json.loads(msg)
                        if data.get("event") == "cancel":
                            cancelled = True
                            logger.info(f"Client requested cancel: {request_id}")
                            break
                        elif data.get("event") == "ping":
                            await websocket.send_text(json.dumps({"event": "pong"}))
                    except asyncio.TimeoutError:
                        continue
                    except Exception:
                        break
            except Exception:
                pass

        cancel_task = asyncio.create_task(listen_for_cancel())

        try:
            # Apply text normalization if enabled
            text = req.text
            if req.normalize:
                try:
                    from veena3modal.processing.text_normalizer import normalize_text
                    text = normalize_text(text)
                except ImportError:
                    pass  # Normalizer not available

            # Stream audio chunks
            async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
                text=text,
                speaker=speaker,
                temperature=req.temperature,
                top_k=req.top_k,
                top_p=req.top_p,
                max_tokens=req.max_tokens,
                repetition_penalty=req.repetition_penalty,
                seed=req.seed,
                enable_chunking=req.chunking,
            ):
                if cancelled:
                    logger.info(f"Generation cancelled: {request_id}")
                    break

                final_metrics = metrics
                chunks_sent += 1
                total_bytes += len(audio_chunk)

                # Record TTFB on first chunk
                if first_chunk_time is None:
                    first_chunk_time = time.time()
                    ttfb_ms = int((first_chunk_time - start_time) * 1000)
                    log_first_audio_emitted(
                        request_id=request_id,
                        ttfb_ms=ttfb_ms,
                        chunk_size_bytes=len(audio_chunk),
                    )
                    record_ttfb(ttfb_seconds=ttfb_ms / 1000, speaker=speaker, stream=True)

                # Send audio as binary frame
                await websocket.send_bytes(audio_chunk)

                # Send progress every 10 chunks
                if chunks_sent % 10 == 0:
                    await websocket.send_text(create_progress_message(
                        chunks_sent=chunks_sent,
                        bytes_sent=total_bytes,
                        elapsed_ms=int((time.time() - start_time) * 1000),
                    ))

        finally:
            cancel_task.cancel()
            try:
                await cancel_task
            except asyncio.CancelledError:
                pass

        # === SEND COMPLETION ===
        total_time = time.time() - start_time
        audio_duration = final_metrics.get("audio_duration_seconds", 0)
        rtf = total_time / audio_duration if audio_duration > 0 else 0

        completion_metrics = {
            "request_id": request_id,
            "chunks_sent": chunks_sent,
            "total_bytes": total_bytes,
            "audio_duration_seconds": round(audio_duration, 2),
            "total_time_seconds": round(total_time, 2),
            "ttfb_ms": int((first_chunk_time - start_time) * 1000) if first_chunk_time else 0,
            "rtf": round(rtf, 3),
            "cancelled": cancelled,
        }

        await websocket.send_text(create_complete_message(
            request_id=request_id,
            metrics=completion_metrics,
        ))

        # Log completion
        if not cancelled:
            log_request_completed(
                request_id=request_id,
                status_code=200,
                total_duration_ms=int(total_time * 1000),
                audio_duration_seconds=audio_duration,
                rtf=rtf,
                chunks_sent=chunks_sent,
            )
            metrics_request_completed(
                status_code=200,
                duration_seconds=total_time,
                speaker=speaker,
                stream=True,
            )
            record_rtf(rtf=rtf, speaker=speaker)
            record_audio_duration(duration_seconds=audio_duration, speaker=speaker)
            record_chunks_sent(chunks=chunks_sent, speaker=speaker)

        # Close connection gracefully
        await websocket.close(code=1000)

    except Exception as e:
        logger.exception(f"WebSocket error: {e}")
        log_request_failed(
            request_id=request_id,
            status_code=500,
            error_code="WEBSOCKET_ERROR",
            error_message=str(e),
        )
        metrics_request_failed(
            status_code=500,
            error_code="WEBSOCKET_ERROR",
        )

        try:
            await websocket.send_text(create_error_message(
                code="INTERNAL_ERROR",
                message=str(e),
                request_id=request_id,
            ))
            await websocket.close(code=1011)
        except Exception:
            pass

    finally:
        clear_request_context()

def add_websocket_routes(app):
    """
    Add WebSocket routes to FastAPI app.

    Call this from create_app() to enable WebSocket support.

    Args:
        app: FastAPI application instance
    """
    from fastapi import WebSocket, Query

    @app.websocket("/v1/tts/ws")
    async def websocket_tts(
        websocket: WebSocket,
        api_key: Optional[str] = Query(None, alias="api_key"),
    ):
        """
        WebSocket endpoint for TTS streaming.

        Connect with optional api_key query parameter:
        ws://host/v1/tts/ws?api_key=your-key

        Then send JSON request:
        {"text": "Hello world", "speaker": "male1"}

        Receive:
        - JSON header message
        - Binary audio chunks
        - JSON progress messages (every 10 chunks)
        - JSON completion message
        """
        # Also check Authorization header
        auth_header = websocket.headers.get("authorization", "")
        if auth_header.startswith("Bearer "):
            api_key = auth_header[7:]

        await handle_websocket_tts(websocket, api_key)

    logger.info("WebSocket route /v1/tts/ws registered")


cd /home/ubuntu/veenaModal && rg -n "def _handle_streaming_request|audio_stream_generator|await generate_speech_streaming|except Exception as e" veena3modal/api/fastapi_app.py


Output:
Chunk ID: 28e0a2
Wall time: 0.0501 seconds
Process exited with code 0
Original token count: 97
Output:
242:        except Exception as e:
247:        except Exception as e:
461:                except Exception as e:
471:        except Exception as e:
500:    async def _handle_streaming_request(req, request_id: str, text: str, speaker: str, start_time: float):
544:        async def audio_stream_generator():
620:            except Exception as e:
658:            audio_stream_generator(),


cd /home/ubuntu/veenaModal && nl -ba veena3modal/api/fastapi_app.py | sed -n '520,690p'


Output:
Chunk ID: a45be7
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 1866
Output:
   520	                        "request_id": request_id,
   521	                    }
   522	                },
   523	                headers={"X-Request-ID": request_id},
   524	            )
   525
   526	        # Non-WAV formats not implemented yet (M5)
   527	        if req.format != AudioFormat.WAV.value:
   528	            return JSONResponse(
   529	                status_code=501,
   530	                content={
   531	                    "error": {
   532	                        "code": "FORMAT_NOT_IMPLEMENTED",
   533	                        "message": f"Streaming format '{req.format}' not yet implemented. Use format=wav.",
   534	                        "request_id": request_id,
   535	                    }
   536	                },
   537	                headers={"X-Request-ID": request_id},
   538	            )
   539
   540	        # Set request context for structured logging
   541	        set_request_context(request_id=request_id)
   542
   543	        # Create async generator for streaming response
   544	        async def audio_stream_generator():
   545	            """Yields audio chunks as they're generated."""
   546	            final_metrics = {}
   547	            first_chunk_sent = False
   548	            chunks_count = 0
   549	            try:
   550	                async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
   551	                    text=text,
   552	                    speaker=speaker,
   553	                    temperature=req.temperature,
   554	                    top_k=req.top_k,
   555	                    top_p=req.top_p,
   556	                    max_tokens=req.max_tokens,
   557	                    repetition_penalty=req.repetition_penalty,
   558	                    seed=req.seed,
   559	                    enable_chunking=req.chunking,
   560	                ):
   561	                    final_metrics = metrics
   562	                    chunks_count += 1
   563
   564	                    # Fire-and-forget sentence storage after first chunk (don't block TTFB)
   565	                    if not first_chunk_sent:
   566	                        first_chunk_sent = True
   567	                        ttfb_ms = metrics.get("ttfb_ms", 0)
   568
   569	                        # Log first audio emitted (TTFB marker)
   570	                        log_first_audio_emitted(
   571	                            request_id=request_id,
   572	                            ttfb_ms=ttfb_ms,
   573	                            chunk_size_bytes=len(audio_chunk),
   574	                        )
   575	                        record_ttfb(ttfb_seconds=ttfb_ms / 1000, speaker=speaker, stream=True)
   576
   577	                        sentence_store = _get_sentence_store()
   578	                        sentence_store.store_fire_and_forget(
   579	                            request_id=request_id,
   580	                            text=text,
   581	                            speaker=speaker,
   582	                            stream=True,
   583	                            format=req.format,
   584	                            temperature=req.temperature,
   585	                            top_k=req.top_k,
   586	                            top_p=req.top_p,
   587	                            max_tokens=req.max_tokens,
   588	                            repetition_penalty=req.repetition_penalty,
   589	                            seed=req.seed,
   590	                            text_chunked=metrics.get("text_chunked", False),
   591	                            ttfb_ms=ttfb_ms,
   592	                            # audio_duration_seconds not known yet for streaming
   593	                        )
   594
   595	                    yield audio_chunk
   596
   597	                # Stream completed - log metrics
   598	                total_time = time.time() - start_time
   599	                audio_duration = final_metrics.get("audio_duration_seconds", 0)
   600	                rtf = total_time / audio_duration if audio_duration > 0 else 0
   601
   602	                log_request_completed(
   603	                    request_id=request_id,
   604	                    status_code=200,
   605	                    total_duration_ms=int(total_time * 1000),
   606	                    audio_duration_seconds=audio_duration,
   607	                    rtf=rtf,
   608	                    chunks_sent=chunks_count,
   609	                )
   610	                metrics_request_completed(
   611	                    status_code=200,
   612	                    duration_seconds=total_time,
   613	                    speaker=speaker,
   614	                    stream=True,
   615	                )
   616	                record_rtf(rtf=rtf, speaker=speaker)
   617	                record_audio_duration(duration_seconds=audio_duration, speaker=speaker)
   618	                record_chunks_sent(chunks=chunks_count, speaker=speaker)
   619	                clear_request_context()
   620	            except Exception as e:
   621	                logger.exception(f"Streaming error: {e}")
   622	                # Log streaming failure
   623	                log_request_failed(
   624	                    request_id=request_id,
   625	                    status_code=500,
   626	                    error_code="STREAMING_ERROR",
   627	                    error_message=str(e),
   628	                )
   629	                metrics_request_failed(
   630	                    status_code=500,
   631	                    error_code="STREAMING_ERROR",
   632	                    speaker=speaker,
   633	                )
   634	                clear_request_context()
   635	                # Can't return error once streaming started
   636	                # Best we can do is stop yielding
   637	                return
   638
   639	        # Build response headers
   640	        # Note: Some headers can't be set dynamically during streaming
   641	        # TTFB and RTF will be logged but not in headers for streaming
   642	        model_version = runtime.model_version if runtime else get_model_version()
   643
   644	        headers = {
   645	            "X-Request-ID": request_id,
   646	            "X-Model-Version": model_version,
   647	            "X-Format": "wav",
   648	            "X-Sample-Rate": str(req.sample_rate),
   649	            "X-Stream": "true",
   650	            "X-Chunking-Enabled": str(req.chunking).lower(),
   651	            # Transfer-Encoding: chunked is set automatically by StreamingResponse
   652	        }
   653
   654	        if req.seed is not None:
   655	            headers["X-Seed"] = str(req.seed)
   656
   657	        return StreamingResponse(
   658	            audio_stream_generator(),
   659	            media_type="audio/wav",
   660	            headers=headers,
   661	        )
   662
   663	    # === WebSocket Support (M5b) ===
   664	    try:
   665	        from veena3modal.api.websocket_handler import add_websocket_routes
   666	        add_websocket_routes(app)
   667	        logger.info("WebSocket support enabled at /v1/tts/ws")
   668	    except ImportError as e:
   669	        logger.warning(f"WebSocket support not available: {e}")
   670
   671	    return app
   672
   673


cd /home/ubuntu/veenaModal && rg -n "generate_speech_streaming\(|request_id|chunked" veena3modal/api/websocket_handler.py && nl -ba veena3modal/api/websocket_handler.py | sed -n '220,520p'


Output:
Chunk ID: da8477
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 3299
Output:
112:def create_error_message(code: str, message: str, request_id: str) -> str:
119:            "request_id": request_id,
125:    request_id: str,
132:        "request_id": request_id,
154:    request_id: str,
160:        "request_id": request_id,
181:    request_id = str(uuid.uuid4())
182:    set_request_context(request_id=request_id)
186:    logger.info(f"WebSocket connection accepted: {request_id}")
202:                    request_id=request_id,
215:                        request_id=request_id,
232:                request_id=request_id,
240:                request_id=request_id,
260:                request_id=request_id,
270:            request_id=request_id,
283:                request_id=request_id,
293:            request_id=request_id,
318:                            logger.info(f"Client requested cancel: {request_id}")
342:            async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
354:                    logger.info(f"Generation cancelled: {request_id}")
366:                        request_id=request_id,
396:            "request_id": request_id,
407:            request_id=request_id,
414:                request_id=request_id,
437:            request_id=request_id,
451:                request_id=request_id,
   220	        # === WAIT FOR REQUEST ===
   221	        try:
   222	            # Wait for client to send request (with timeout)
   223	            raw_message = await asyncio.wait_for(
   224	                websocket.receive_text(),
   225	                timeout=30.0
   226	            )
   227	            message = json.loads(raw_message)
   228	        except asyncio.TimeoutError:
   229	            await websocket.send_text(create_error_message(
   230	                code="TIMEOUT",
   231	                message="No request received within 30 seconds",
   232	                request_id=request_id,
   233	            ))
   234	            await websocket.close(code=1000)
   235	            return
   236	        except json.JSONDecodeError as e:
   237	            await websocket.send_text(create_error_message(
   238	                code="INVALID_JSON",
   239	                message=f"Invalid JSON: {e}",
   240	                request_id=request_id,
   241	            ))
   242	            await websocket.close(code=1003)  # Unsupported data
   243	            return
   244
   245	        # Handle control messages
   246	        event_type = message.get("event", "request")
   247	        if event_type == "ping":
   248	            await websocket.send_text(json.dumps({"event": "pong"}))
   249	            # Wait for actual request
   250	            raw_message = await websocket.receive_text()
   251	            message = json.loads(raw_message)
   252
   253	        # === PARSE REQUEST ===
   254	        req = WSRequest.from_dict(message)
   255	        validation_error = req.validate()
   256	        if validation_error:
   257	            await websocket.send_text(create_error_message(
   258	                code="VALIDATION_ERROR",
   259	                message=validation_error,
   260	                request_id=request_id,
   261	            ))
   262	            await websocket.close(code=1003)
   263	            return
   264
   265	        # Resolve speaker
   266	        speaker = resolve_speaker_name(req.speaker)
   267
   268	        # Log request
   269	        log_request_received(
   270	            request_id=request_id,
   271	            text_length=len(req.text),
   272	            speaker=speaker,
   273	            stream=True,
   274	            format="wav",
   275	        )
   276	        metrics_request_received(speaker=speaker, stream=True, format="wav")
   277
   278	        # === CHECK MODEL ===
   279	        if not tts_runtime.is_initialized():
   280	            await websocket.send_text(create_error_message(
   281	                code="MODEL_NOT_LOADED",
   282	                message="TTS model not initialized",
   283	                request_id=request_id,
   284	            ))
   285	            await websocket.close(code=1011)  # Internal error
   286	            return
   287
   288	        runtime = tts_runtime.get_runtime()
   289	        model_version = runtime.model_version if runtime else "unknown"
   290
   291	        # === SEND HEADER INFO ===
   292	        await websocket.send_text(create_header_message(
   293	            request_id=request_id,
   294	            sample_rate=req.sample_rate,
   295	            model_version=model_version,
   296	        ))
   297
   298	        # === STREAM AUDIO ===
   299	        start_time = time.time()
   300	        chunks_sent = 0
   301	        total_bytes = 0
   302	        first_chunk_time = None
   303	        final_metrics = {}
   304
   305	        # Start a task to listen for cancel messages
   306	        async def listen_for_cancel():
   307	            nonlocal cancelled
   308	            try:
   309	                while not cancelled:
   310	                    try:
   311	                        msg = await asyncio.wait_for(
   312	                            websocket.receive_text(),
   313	                            timeout=0.1
   314	                        )
   315	                        data = json.loads(msg)
   316	                        if data.get("event") == "cancel":
   317	                            cancelled = True
   318	                            logger.info(f"Client requested cancel: {request_id}")
   319	                            break
   320	                        elif data.get("event") == "ping":
   321	                            await websocket.send_text(json.dumps({"event": "pong"}))
   322	                    except asyncio.TimeoutError:
   323	                        continue
   324	                    except Exception:
   325	                        break
   326	            except Exception:
   327	                pass
   328
   329	        cancel_task = asyncio.create_task(listen_for_cancel())
   330
   331	        try:
   332	            # Apply text normalization if enabled
   333	            text = req.text
   334	            if req.normalize:
   335	                try:
   336	                    from veena3modal.processing.text_normalizer import normalize_text
   337	                    text = normalize_text(text)
   338	                except ImportError:
   339	                    pass  # Normalizer not available
   340
   341	            # Stream audio chunks
   342	            async for audio_chunk, metrics in tts_runtime.generate_speech_streaming(
   343	                text=text,
   344	                speaker=speaker,
   345	                temperature=req.temperature,
   346	                top_k=req.top_k,
   347	                top_p=req.top_p,
   348	                max_tokens=req.max_tokens,
   349	                repetition_penalty=req.repetition_penalty,
   350	                seed=req.seed,
   351	                enable_chunking=req.chunking,
   352	            ):
   353	                if cancelled:
   354	                    logger.info(f"Generation cancelled: {request_id}")
   355	                    break
   356
   357	                final_metrics = metrics
   358	                chunks_sent += 1
   359	                total_bytes += len(audio_chunk)
   360
   361	                # Record TTFB on first chunk
   362	                if first_chunk_time is None:
   363	                    first_chunk_time = time.time()
   364	                    ttfb_ms = int((first_chunk_time - start_time) * 1000)
   365	                    log_first_audio_emitted(
   366	                        request_id=request_id,
   367	                        ttfb_ms=ttfb_ms,
   368	                        chunk_size_bytes=len(audio_chunk),
   369	                    )
   370	                    record_ttfb(ttfb_seconds=ttfb_ms / 1000, speaker=speaker, stream=True)
   371
   372	                # Send audio as binary frame
   373	                await websocket.send_bytes(audio_chunk)
   374
   375	                # Send progress every 10 chunks
   376	                if chunks_sent % 10 == 0:
   377	                    await websocket.send_text(create_progress_message(
   378	                        chunks_sent=chunks_sent,
   379	                        bytes_sent=total_bytes,
   380	                        elapsed_ms=int((time.time() - start_time) * 1000),
   381	                    ))
   382
   383	        finally:
   384	            cancel_task.cancel()
   385	            try:
   386	                await cancel_task
   387	            except asyncio.CancelledError:
   388	                pass
   389
   390	        # === SEND COMPLETION ===
   391	        total_time = time.time() - start_time
   392	        audio_duration = final_metrics.get("audio_duration_seconds", 0)
   393	        rtf = total_time / audio_duration if audio_duration > 0 else 0
   394
   395	        completion_metrics = {
   396	            "request_id": request_id,
   397	            "chunks_sent": chunks_sent,
   398	            "total_bytes": total_bytes,
   399	            "audio_duration_seconds": round(audio_duration, 2),
   400	            "total_time_seconds": round(total_time, 2),
   401	            "ttfb_ms": int((first_chunk_time - start_time) * 1000) if first_chunk_time else 0,
   402	            "rtf": round(rtf, 3),
   403	            "cancelled": cancelled,
   404	        }
   405
   406	        await websocket.send_text(create_complete_message(
   407	            request_id=request_id,
   408	            metrics=completion_metrics,
   409	        ))
   410
   411	        # Log completion
   412	        if not cancelled:
   413	            log_request_completed(
   414	                request_id=request_id,
   415	                status_code=200,
   416	                total_duration_ms=int(total_time * 1000),
   417	                audio_duration_seconds=audio_duration,
   418	                rtf=rtf,
   419	                chunks_sent=chunks_sent,
   420	            )
   421	            metrics_request_completed(
   422	                status_code=200,
   423	                duration_seconds=total_time,
   424	                speaker=speaker,
   425	                stream=True,
   426	            )
   427	            record_rtf(rtf=rtf, speaker=speaker)
   428	            record_audio_duration(duration_seconds=audio_duration, speaker=speaker)
   429	            record_chunks_sent(chunks=chunks_sent, speaker=speaker)
   430
   431	        # Close connection gracefully
   432	        await websocket.close(code=1000)
   433
   434	    except Exception as e:
   435	        logger.exception(f"WebSocket error: {e}")
   436	        log_request_failed(
   437	            request_id=request_id,
   438	            status_code=500,
   439	            error_code="WEBSOCKET_ERROR",
   440	            error_message=str(e),
   441	        )
   442	        metrics_request_failed(
   443	            status_code=500,
   444	            error_code="WEBSOCKET_ERROR",
   445	        )
   446
   447	        try:
   448	            await websocket.send_text(create_error_message(
   449	                code="INTERNAL_ERROR",
   450	                message=str(e),
   451	                request_id=request_id,
   452	            ))
   453	            await websocket.close(code=1011)
   454	        except Exception:
   455	            pass
   456
   457	    finally:
   458	        clear_request_context()
   459
   460
   461	def add_websocket_routes(app):
   462	    """
   463	    Add WebSocket routes to FastAPI app.
   464
   465	    Call this from create_app() to enable WebSocket support.
   466
   467	    Args:
   468	        app: FastAPI application instance
   469	    """
   470	    from fastapi import WebSocket, Query
   471
   472	    @app.websocket("/v1/tts/ws")
   473	    async def websocket_tts(
   474	        websocket: WebSocket,
   475	        api_key: Optional[str] = Query(None, alias="api_key"),
   476	    ):
   477	        """
   478	        WebSocket endpoint for TTS streaming.
   479
   480	        Connect with optional api_key query parameter:
   481	        ws://host/v1/tts/ws?api_key=your-key
   482
   483	        Then send JSON request:
   484	        {"text": "Hello world", "speaker": "male1"}
   485
   486	        Receive:
   487	        - JSON header message
   488	        - Binary audio chunks
   489	        - JSON progress messages (every 10 chunks)
   490	        - JSON completion message
   491	        """
   492	        # Also check Authorization header
   493	        auth_header = websocket.headers.get("authorization", "")
   494	        if auth_header.startswith("Bearer "):
   495	            api_key = auth_header[7:]
   496
   497	        await handle_websocket_tts(websocket, api_key)
   498
   499	    logger.info("WebSocket route /v1/tts/ws registered")
   500


cd /home/ubuntu/veenaModal && sed -n '1,320p' scripts/profile_pipeline.py && sed -n '1,320p' scripts/stress_test_local.py && sed -n '1,340p' scripts/stress_test_production.py