Output:
Chunk ID: 7774a8
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 3756
Output:
                    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:
                logger.exception(f"Streaming error: {e}")
                # Log streaming failure
                log_request_failed(
                    request_id=request_id,
                    status_code=500,
                    error_code="STREAMING_ERROR",
                    error_message=str(e),
                )
                metrics_request_failed(
                    status_code=500,
                    error_code="STREAMING_ERROR",
                    speaker=speaker,
                )
                clear_request_context()
                # Can't return error once streaming started
                # Best we can do is stop yielding
                return

        # Build response headers
        # Note: Some headers can't be set dynamically during streaming
        # TTFB and RTF will be logged but not in headers for streaming
        model_version = runtime.model_version if runtime else get_model_version()

        headers = {
            "X-Request-ID": request_id,
            "X-Model-Version": model_version,
            "X-Format": "wav",
            "X-Sample-Rate": str(req.sample_rate),
            "X-Stream": "true",
            "X-Chunking-Enabled": str(req.chunking).lower(),
            # Transfer-Encoding: chunked is set automatically by StreamingResponse
        }

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

        return StreamingResponse(
            audio_stream_generator(),
            media_type="audio/wav",
            headers=headers,
        )

    # === WebSocket Support (M5b) ===
    try:
        from veena3modal.api.websocket_handler import add_websocket_routes
        add_websocket_routes(app)
        logger.info("WebSocket support enabled at /v1/tts/ws")
    except ImportError as e:
        logger.warning(f"WebSocket support not available: {e}")

    return app


# 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.


sed -n '1,320p' veena3modal/api/websocket_handler.py


Output:
Chunk ID: 498ad9
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 2630
Output:
"""
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(
                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":


Summary: auto


I’ve identified the remaining gaps: streaming and chunking paths in runtime are not consistently attaching timing/`request_id`, and streaming pipeline logging has a few malformed blocks from earlier edits. I’m now reading the websocket end and long-text module to finish the wiring end-to-end, then I’ll patch all at once.


sed -n '320,700p' veena3modal/api/websocket_handler.py


Output:
Chunk ID: 732d3a
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 1585
Output:
                        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")


sed -n '1,420p' veena3modal/processing/long_text_processor.py


Output:
Chunk ID: 1f9e89
Wall time: 0.0527 seconds
Process exited with code 0
Original token count: 3571
Output:
"""
Long Text Processing Service

Handles text chunking and audio stitching for long text inputs.

Why This Approach:
- Text-based chunking: Preserves semantic meaning
- SNAC token sliding window: Separate layer for smooth audio
- Crossfade stitching: Seamless audio transitions between chunks
- Unlimited text: Can process any length text

Architecture:
1. Long text → IndicSentenceChunker → Text chunks
2. Each text chunk → TTS Model → Audio chunk
3. All audio chunks → Crossfade stitching → Final audio
"""

import asyncio
import numpy as np
import struct
from typing import List, Optional, AsyncGenerator
import logging

from veena3modal.processing.text_chunker import IndicSentenceChunker, crossfade_audio
from veena3modal.core.streaming_pipeline import Veena3SlidingWindowPipeline
from veena3modal.core.pipeline import SparkTTSPipeline

logger = logging.getLogger(__name__)

def create_wav_header(data_size: int, sample_rate: int = 24000, channels: int = 1, bits_per_sample: int = 16) -> bytes:
    """
    Create WAV file header for PCM data.

    Args:
        data_size: Size of PCM data in bytes
        sample_rate: Sample rate (Hz)
        channels: Number of channels
        bits_per_sample: Bits per sample

    Returns:
        WAV header bytes (44 bytes)
    """
    byte_rate = sample_rate * channels * bits_per_sample // 8
    block_align = channels * bits_per_sample // 8

    header = struct.pack('<4sI4s', b'RIFF', 36 + data_size, b'WAVE')
    header += struct.pack('<4sIHHIIHH', b'fmt ', 16, 1, channels, sample_rate, byte_rate, block_align, bits_per_sample)
    header += struct.pack('<4sI', b'data', data_size)

    return header

class LongTextProcessor:
    """
    Process long texts by chunking, generating audio, and stitching.

    Features:
    - Intelligent text chunking (respects sentence boundaries, Indic languages)
    - Parallel audio generation (if desired)
    - Crossfade stitching for seamless audio
    - Memory efficient (processes chunks sequentially by default)

    Why text chunking, not token chunking:
    - Text maintains semantic meaning
    - Natural sentence boundaries = better prosody
    - Simpler implementation
    - Model handles tokenization internally
    - SNAC token sliding window is separate (already implemented)

    Max Input Length Reference (post-normalization):
    "इसके निर्माण में मुख्य वास्तुकार उस्ताद अहमद लाहौरी के नेतृत्व में लगभग
    20,000 कारीगरों और शिल्पकारों ने दिन-रात मेहनत की थी। इमारत के सामने बना
    'चारबाग' शैली का उद्यान और पानी की नहरें इसकी सुंदरता में चार चाँद लगा देते हैं।"
    = ~232 characters after normalization (numbers expanded to words)
    """

    # Maximum safe input length for model (characters)
    # Based on reference Hindi sentence (~232 chars after normalization)
    # Model struggles with longer inputs, causes token repetition
    MAX_MODEL_INPUT_LENGTH = 230  # ~232 chars is the reference max

    # Threshold for when to enable chunking (characters)
    # Trigger chunking when text approaches max safe length
    # Using 95% of max to leave some buffer
    CHUNKING_THRESHOLD = 220  # ~95% of MAX_MODEL_INPUT_LENGTH

    # Chunk size for long texts (characters)
    # Use 75% of max safe length to ensure each chunk processes cleanly
    # This gives model headroom for complex multilingual content
    CHUNK_SIZE = 170  # ~75% of MAX_MODEL_INPUT_LENGTH

    # Crossfade duration (milliseconds)
    CROSSFADE_MS = 50  # 50ms smooth transition

    def __init__(
        self,
        pipeline: SparkTTSPipeline,
        streaming_pipeline: Optional[Veena3SlidingWindowPipeline] = None
    ):
        """
        Initialize long text processor.

        Args:
            pipeline: Non-streaming TTS pipeline
            streaming_pipeline: Optional streaming pipeline (for streaming mode)
        """
        self.pipeline = pipeline
        self.streaming_pipeline = streaming_pipeline
        self.chunker = IndicSentenceChunker(max_chunk_length=self.CHUNK_SIZE)

        logger.info(
            f"LongTextProcessor initialized",
            extra={
                "max_model_input_length": self.MAX_MODEL_INPUT_LENGTH,
                "chunking_threshold": self.CHUNKING_THRESHOLD,
                "chunk_size": self.CHUNK_SIZE,
                "crossfade_ms": self.CROSSFADE_MS,
            }
        )

    def should_chunk(self, text: str) -> bool:
        """Determine if text should be chunked."""
        return len(text) > self.CHUNKING_THRESHOLD

    def chunk_text(self, text: str) -> list:
        """
        Chunk text using the internal chunker.

        Delegates to self.chunker.chunk_text() for actual chunking logic.
        This method is used by _stream_chunked_text in tts_runtime.

        Args:
            text: Text to chunk

        Returns:
            List of text chunks
        """
        return self.chunker.chunk_text(text)

    async def generate_with_chunking(
        self,
        text: str,
        speaker: str,
        temperature: float = 0.4,
        top_k: int = 50,  # Added for Spark TTS
        top_p: float = 0.9,
        max_tokens: int = 4096,
        repetition_penalty: float = 1.05,  # Prevent token repetition
        seed: Optional[int] = None,
        sample_rate: int = 16000,  # BiCodec uses 16kHz, not 24kHz
    ) -> Optional[bytes]:
        """
        Generate audio for long text with chunking and stitching.

        Process:
        1. Chunk text at natural boundaries
        2. Generate audio for each chunk
        3. Stitch audio with crossfade

        Args:
            text: Long text to synthesize
            speaker: Speaker name (lipakshi, reet, etc.)
            temperature: Sampling temperature
            top_p: Nucleus sampling
            max_tokens: Max tokens per chunk
            seed: Random seed (same seed = same voice across chunks)
            sample_rate: Audio sample rate

        Returns:
            Stitched audio bytes (int16 PCM) or None if generation fails
        """
        # Detect language mix for logging
        lang_info = self.chunker.detect_language_mix(text)

        logger.info(
            f"Processing long text with chunking",
            extra={
                "text_length": len(text),
                "primary_language": lang_info['primary'],
                "lang_mix": lang_info,
                "speaker": speaker,
            }
        )

        # 1. Chunk text
        text_chunks = self.chunker.chunk_text(text)

        logger.info(
            f"Text chunked into {len(text_chunks)} chunks",
            extra={
                "num_chunks": len(text_chunks),
                "chunk_sizes": [len(chunk) for chunk in text_chunks],
            }
        )

        # 2. Generate audio for each chunk
        audio_chunks = []

        # Log summary at start
        logger.info(
            f"🎵 Generating {len(text_chunks)} chunks for {len(text)} chars",
            extra={
                "total_chunks": len(text_chunks),
                "total_chars": len(text),
                "speaker": speaker,
            }
        )

        for i, chunk in enumerate(text_chunks):
            # Calculate appropriate max_tokens for this chunk
            # BiCodec semantic tokens: ~50 tokens/second, ~5 words/second = ~10 tokens/word
            # Rough estimate: 5 chars/word, so ~2 tokens/char for audio generation
            # Add buffer for safety: multiply by 3 to handle variations
            chunk_max_tokens = min(
                int(len(chunk) * 3),  # ~2 tokens/char * 1.5 safety buffer = 3
                4096  # Cap at model's max context length
            )
            chunk_max_tokens = max(chunk_max_tokens, 100)  # Ensure minimum

            # Generate audio for this chunk (returns WAV format with header)
            audio_bytes = await self.pipeline.generate_speech_indic(
                speaker=speaker,
                text=chunk,
                temperature=temperature,
                top_k=top_k,  # Added for Spark TTS
                top_p=top_p,
                max_tokens=chunk_max_tokens,  # Use calculated limit, not global max_tokens
                repetition_penalty=repetition_penalty,  # Prevent token repetition
                seed=seed,  # Use same seed for consistent voice across chunks
            )

            if audio_bytes is None:
                logger.error(
                    f"❌ Failed to generate audio for chunk {i+1}/{len(text_chunks)}",
                    extra={
                        "chunk_index": i,
                        "chunk_length": len(chunk),
                        "chunk_text_preview": chunk[:150],
                        "total_chunks": len(text_chunks),
                    }
                )
                return None

            # Strip WAV header (44 bytes) to get raw PCM data for stitching
            pcm_data = audio_bytes[44:] if len(audio_bytes) > 44 else audio_bytes
            audio_chunks.append(pcm_data)

        # 3. Stitch audio with crossfade
        stitched_pcm = self.stitch_audio_chunks(
            audio_chunks,
            sample_rate=sample_rate,
            crossfade_ms=self.CROSSFADE_MS,
        )

        # 4. Add WAV header to stitched PCM data
        wav_header = create_wav_header(
            data_size=len(stitched_pcm),
            sample_rate=sample_rate,
            channels=1,
            bits_per_sample=16
        )
        final_audio = wav_header + stitched_pcm

        # Calculate total audio metrics
        total_audio_duration = len(stitched_pcm) / (sample_rate * 2)

        logger.info(
            f"✅ Generated {len(text_chunks)} chunks → {total_audio_duration:.1f}s audio",
            extra={
                "total_chunks": len(text_chunks),
                "audio_duration": total_audio_duration,
                "audio_size_mb": len(final_audio) / (1024 * 1024),
            }
        )

        return final_audio

    def stitch_audio_chunks(
        self,
        audio_chunks: List[bytes],
        sample_rate: int = 16000,  # BiCodec uses 16kHz, not 24kHz
        crossfade_ms: int = 50,
    ) -> bytes:
        """
        Stitch multiple audio chunks with crossfade transitions.

        Why crossfade:
        - Prevents pops/clicks at chunk boundaries
        - Smooth transitions between chunks
        - Professional audio quality

        Args:
            audio_chunks: List of audio bytes (int16 PCM)
            sample_rate: Sample rate (Hz)
            crossfade_ms: Crossfade duration (milliseconds)

        Returns:
            Stitched audio bytes
        """
        if not audio_chunks:
            return b''

        if len(audio_chunks) == 1:
            return audio_chunks[0]

        # Calculate crossfade samples
        crossfade_samples = int((crossfade_ms / 1000.0) * sample_rate)

        # Start with first chunk
        result = audio_chunks[0]

        # Stitch remaining chunks with crossfade
        for i in range(1, len(audio_chunks)):
            result = crossfade_audio(
                result,
                audio_chunks[i],
                crossfade_samples=crossfade_samples,
                sample_rate=sample_rate,
            )

            logger.debug(
                f"Crossfaded chunk {i+1}",
                extra={
                    "chunk_index": i,
                    "current_length_bytes": len(result),
                }
            )

        return result

    async def generate_with_chunking_streaming(
        self,
        text: str,
        speaker: str,
        temperature: float = 0.4,
        top_p: float = 0.9,
        max_tokens: int = 4096,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate audio for long text with streaming (experimental).

        Note: Streaming long texts is more complex:
        - Need to stream each chunk
        - Need to buffer for crossfade
        - Adds latency for stitching

        For now, recommend using non-streaming for long texts.

        Args:
            text: Long text to synthesize
            speaker: Speaker name
            temperature: Sampling temperature
            top_p: Nucleus sampling
            max_tokens: Max tokens per chunk
            repetition_penalty: Repetition penalty
            seed: Random seed

        Yields:
            Audio bytes (in chunks)
        """
        if self.streaming_pipeline is None:
            raise ValueError("Streaming pipeline not available")

        # Chunk text
        text_chunks = self.chunker.chunk_text(text)

        logger.info(
            f"Streaming long text with {len(text_chunks)} chunks",
            extra={"num_chunks": len(text_chunks)}
        )

        # For streaming, we process each chunk sequentially
        # Future optimization: overlap processing + streaming
        for i, chunk in enumerate(text_chunks):
            logger.info(f"Streaming chunk {i+1}/{len(text_chunks)}")

            # Stream this chunk
            async for audio_bytes in self.streaming_pipeline.generate_speech_stream_indic(
                speaker=speaker,
                text=chunk,
                temperature=temperature,
                top_p=top_p,
                max_tokens=max_tokens,
                seed=seed,
            ):
                yield audio_bytes

            # TODO: Add crossfade between chunks in streaming mode
            # This requires buffering the end of previous chunk
            # and beginning of next chunk


sed -n '320,760p' veena3modal/core/pipeline.py


Output:
Chunk ID: 29302b
Wall time: 0.0517 seconds
Process exited with code 0
Original token count: 779
Output:

    async def generate_speech_indic(
        self,
        speaker: str,
        text: str,
        temperature: float = DEFAULT_TEMPERATURE,
        top_k: int = DEFAULT_TOP_K,
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = 1.05,
        seed: Optional[int] = None,
        return_timing: bool = False,
        request_id: Optional[str] = None,
    ):
        """
        Generate speech audio for Spark TTS (speaker-based).

        This method is kept for backward compatibility with existing API.
        It delegates to generate_speech() with the same implementation.

        Args:
            speaker: Speaker name (one of 12 predefined speakers)
            text: Text to synthesize with inline emotion tags
                Examples:
                - "Hello! Welcome."
                - "Hello <laugh> this is fun!" (will be normalized)
                - "नमस्ते! [excited] आज का दिन बहुत अच्छा है।"
            temperature: Sampling temperature
            top_k: Top-k sampling
            top_p: Nucleus sampling
            max_tokens: Max BiCodec tokens to generate
            seed: Random seed for reproducibility
            return_timing: If true, returns (audio_bytes, metrics_dict)

        Returns:
            Audio bytes (int16 PCM WAV, 16kHz mono) or None if failed.
            If return_timing is true, returns tuple(audio_bytes, metrics).
        """
        # Delegate to generate_speech() with the same implementation for Spark TTS
        return await self.generate_speech(
            speaker=speaker,
            text=text,
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            repetition_penalty=repetition_penalty,
            seed=seed,
            return_timing=return_timing,
            request_id=request_id,
        )

    def _extract_bicodec_tokens(self, generated_text: str) -> Tuple[List[int], List[int]]:
        """
        Extract BiCodec semantic and global tokens from generated text using regex.

        Spark TTS generates tokens in the format:
        - <|bicodec_semantic_{id}|>
        - <|bicodec_global_{id}|>

        Args:
            generated_text: Generated text containing BiCodec token markers

        Returns:
            Tuple of (semantic_ids, global_ids)
            Returns ([], []) if no tokens found
        """
        # Extract semantic tokens
        semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", generated_text)
        semantic_ids = [int(t) for t in semantic_matches]

        # Extract global tokens
        global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", generated_text)
        global_ids = [int(t) for t in global_matches]

        # Only log if there's an issue
        if not semantic_ids and not global_ids:
            logger.error(f"❌ No BiCodec tokens found! Generated text (first 1000 chars):\n{generated_text[:1000]}")

        return semantic_ids, global_ids


sed -n '1,360p' veena3modal/core/pipeline.py


Output:
Chunk ID: 4e5d8a
Wall time: 0.0514 seconds
Process exited with code 0
Original token count: 3336
Output:
"""
Spark TTS Generation Pipeline

End-to-end pipeline for TTS generation with BiCodec audio tokenizer.
Migrated from Veena3/Orpheus (SNAC) to Spark TTS (BiCodec).
"""

import asyncio
import logging
import os
import re
from statistics import mean, median
from typing import Optional, List, Dict, Any, Tuple

from vllm import SamplingParams
import torch

logger = logging.getLogger(__name__)

from veena3modal.core.constants import (
    TRAINING_STOP_TOKEN_IDS,
    DEFAULT_TEMPERATURE,
    DEFAULT_TOP_K,
    DEFAULT_TOP_P,
    DEFAULT_MAX_TOKENS,
    DEFAULT_SEED,
    AUDIO_SAMPLE_RATE,
)

def _detailed_timing_enabled() -> bool:
    """
    Return True when benchmark logs are enabled.

    Controlled by environment variable:
      VEENA3_TIMING_LOGS=true|1
    """
    return os.getenv("VEENA3_TIMING_LOGS", "").lower() in {"1", "true", "yes", "on"}

class SparkTTSPipeline:
    """
    End-to-end TTS pipeline for Spark TTS with BiCodec.

    Replaces SNAC-based token extraction with BiCodec regex-based extraction.
    """

    def __init__(
        self,
        model,
        prompt_builder,
        bicodec_decoder,
    ):
        """
        Initialize pipeline.

        Args:
            model: SparkTTS Model instance
            prompt_builder: IndicPromptBuilder instance (with Spark TTS format)
            bicodec_decoder: BiCodecDecoder instance
        """
        self.model = model
        self.prompt_builder = prompt_builder
        self.bicodec_decoder = bicodec_decoder

        print(f"🚀 SparkTTSPipeline initialized")

    async def generate_speech(
        self,
        speaker: str,
        text: str,
        temperature: float = DEFAULT_TEMPERATURE,
        top_k: int = DEFAULT_TOP_K,
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = 1.05,
        seed: Optional[int] = None,
        return_timing: bool = False,
        request_id: Optional[str] = None,
    ) -> Optional[bytes] | tuple[Optional[bytes], Dict[str, Any]]:
        """
        Generate speech audio (non-streaming) using Spark TTS.

        NOTE: This method signature changed from description-based to speaker-based
        to align with Spark TTS architecture.

        Args:
            speaker: Speaker name (one of 12 predefined speakers)
            text: Text to synthesize (with optional [emotion] tags)
            temperature: Sampling temperature (default: 0.8 for Spark TTS)
            top_k: Top-k sampling (default: 50)
            top_p: Nucleus sampling (default: 1.0)
            max_tokens: Max BiCodec tokens to generate (default: 2048)
            seed: Random seed for reproducibility
            return_timing: If true, returns (audio_bytes, metrics_dict)

        Returns:
            Audio bytes (int16 PCM WAV, 16kHz mono) or None if failed.
            If return_timing is true, returns tuple(audio_bytes, metrics).
        """
        # Build prompt using Spark TTS format
        prompt = self.prompt_builder.build_prefix(speaker, text)

        # Configure sampling for Spark TTS
        sampling_params = SamplingParams(
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            repetition_penalty=repetition_penalty,
            stop=TRAINING_STOP_TOKEN_IDS,
            skip_special_tokens=False,
            seed=seed,
        )

        # Use vLLM engine to generate
        if request_id is None:
            import uuid
            request_id = f"req-{uuid.uuid4().hex[:12]}"

        emit_timing = return_timing or _detailed_timing_enabled()
        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Collect all outputs and track generation timing
        final_output = None
        token_batches: List[int] = []
        batch_wall_ms: List[float] = []
        batch_gpu_ms: List[float] = []
        total_token_ids = 0
        first_token_ms: Optional[float] = None
        t_start = asyncio.get_event_loop().time()

        iterator = results_generator.__aiter__()
        while True:
            wall_start = asyncio.get_event_loop().time()
            start_event = None
            end_event = None
            if torch.cuda.is_available():
                torch.cuda.synchronize()
                start_event = torch.cuda.Event(enable_timing=True)
                end_event = torch.cuda.Event(enable_timing=True)
                start_event.record()

            try:
                request_output = await iterator.__anext__()
            except StopAsyncIteration:
                break

            wall_end = asyncio.get_event_loop().time()
            batch_ms = (wall_end - wall_start) * 1000
            batch_wall_ms.append(batch_ms)

            batch_gpu_ms_value: Optional[float] = None
            if start_event is not None and end_event is not None:
                end_event.record()
                torch.cuda.synchronize()
                batch_gpu_ms_value = start_event.elapsed_time(end_event)
            if batch_gpu_ms_value is not None:
                batch_gpu_ms.append(batch_gpu_ms_value)

            total_generated = len(request_output.outputs[0].token_ids)
            new_tokens = max(0, total_generated - total_token_ids)
            total_token_ids = total_generated
            token_batches.append(new_tokens)

            if first_token_ms is None:
                first_token_ms = (wall_end - t_start) * 1000

            if emit_timing:
                logger.info(
                    "tts_llm_batch",
                    extra={
                        "speaker": speaker,
                        "request_id": request_id,
                        "batch_index": len(token_batches),
                        "batch_tokens": new_tokens,
                        "llm_batch_wall_ms": round(batch_ms, 3),
                        "llm_batch_gpu_ms": round(batch_gpu_ms_value, 3)
                        if batch_gpu_ms_value is not None
                        else None,
                    },
                )
            final_output = request_output

        if final_output is None:
            if emit_timing:
                logger.warning(
                    "tts_llm_no_output",
                    extra={"speaker": speaker, "request_id": request_id},
                )
            print(f"❌ Generation failed: no output")
            if return_timing:
                return None, {
                    "request_id": request_id,
                    "ttfb_ms": 0,
                    "token_total": 0,
                    "batch_count": 0,
                }
            return None

        # Extract generated text (contains BiCodec tokens)
        parse_start = asyncio.get_event_loop().time()
        generated_text = final_output.outputs[0].text
        generated_ids = final_output.outputs[0].token_ids

        # Extract BiCodec tokens using regex
        semantic_ids, global_ids = self._extract_bicodec_tokens(generated_text)

        parse_ms = (asyncio.get_event_loop().time() - parse_start) * 1000

        if not semantic_ids or not global_ids:
            if emit_timing:
                logger.error(
                    "tts_bicodec_parse_failed",
                    extra={
                        "speaker": speaker,
                        "request_id": request_id,
                        "semantic_count": len(semantic_ids),
                        "global_count": len(global_ids),
                        "token_total": len(generated_ids),
                    },
                )
            logger.error(
                "❌ No BiCodec tokens found in output",
                extra={
                    "speaker": speaker,
                    "text_length": len(text),
                    "text_preview": text[:100],
                    "generated_tokens": len(generated_ids),
                    "semantic_tokens": len(semantic_ids),
                    "global_tokens": len(global_ids),
                    "generated_text_preview": generated_text[:500],
                }
            )
            if return_timing:
                return None, {
                    "request_id": request_id,
                    "ttfb_ms": int(first_token_ms or 0),
                    "token_total": len(generated_ids),
                    "batch_count": len(token_batches),
                }
            return None

        # Validate tokens
        if not self.bicodec_decoder.validate_tokens(semantic_ids, global_ids):
            logger.error("❌ BiCodec token validation failed", extra={
                "speaker": speaker,
                "text_length": len(text),
                "semantic_tokens": len(semantic_ids),
                "global_tokens": len(global_ids),
            })
            if return_timing:
                return None, {
                    "request_id": request_id,
                    "ttfb_ms": int(first_token_ms or 0),
                    "token_total": len(generated_ids),
                    "batch_count": len(token_batches),
                }
            return None

        # Decode to audio
        decode_start = asyncio.get_event_loop().time()
        audio_bytes = self.bicodec_decoder.decode_to_bytes(semantic_ids, global_ids)
        decode_ms = (asyncio.get_event_loop().time() - decode_start) * 1000

        if audio_bytes is None:
            if emit_timing:
                logger.error("tts_decode_failed", extra={"speaker": speaker, "request_id": request_id})
            logger.error("❌ BiCodec decode failed", extra={
                "speaker": speaker,
                "text_length": len(text),
            })
            if return_timing:
                return None, {
                    "request_id": request_id,
                    "ttfb_ms": int(first_token_ms or 0),
                    "token_total": len(generated_ids),
                    "batch_count": len(token_batches),
                }
            return None

        total_ms = (asyncio.get_event_loop().time() - t_start) * 1000

        timing_metrics: Dict[str, Any] = {
            "request_id": request_id,
            "speaker": speaker,
            "text_length": len(text),
            "batch_count": len(token_batches),
            "token_total": len(generated_ids),
            "token_batches": token_batches,
            "ttfb_ms": int(first_token_ms or 0),
            "llm_wall_ms": int(sum(batch_wall_ms)),
            "llm_avg_batch_ms": int(mean(batch_wall_ms)) if batch_wall_ms else 0,
            "llm_gpu_batch_ms_total": int(sum(batch_gpu_ms)),
            "llm_gpu_batch_ms_min": int(min(batch_gpu_ms)) if batch_gpu_ms else 0,
            "llm_gpu_batch_ms_max": int(max(batch_gpu_ms)) if batch_gpu_ms else 0,
            "llm_gpu_batch_ms_p50": int(median(batch_gpu_ms)) if batch_gpu_ms else 0,
            "decode_cpu_ms": int(decode_ms),
            "parse_ms": int(parse_ms),
            "total_ms": int(total_ms),
        }
        if token_batches:
            timing_metrics["tokens_per_batch_min"] = min(token_batches)
            timing_metrics["tokens_per_batch_max"] = max(token_batches)
            timing_metrics["tokens_per_batch_p50"] = int(median(token_batches))

        if emit_timing:
            logger.info("tts_llm_complete", extra=timing_metrics)

        from veena3modal.audio.utils import add_wav_header
        wav_bytes = add_wav_header(audio_bytes, sample_rate=AUDIO_SAMPLE_RATE)

        if return_timing:
            return wav_bytes, timing_metrics
        return wav_bytes

    async def generate_speech_indic(
        self,
        speaker: str,
        text: str,
        temperature: float = DEFAULT_TEMPERATURE,
        top_k: int = DEFAULT_TOP_K,
        top_p: float = DEFAULT_TOP_P,
        max_tokens: int = DEFAULT_MAX_TOKENS,
        repetition_penalty: float = 1.05,
        seed: Optional[int] = None,
        return_timing: bool = False,
        request_id: Optional[str] = None,
    ):
        """
        Generate speech audio for Spark TTS (speaker-based).

        This method is kept for backward compatibility with existing API.
        It delegates to generate_speech() with the same implementation.

        Args:
            speaker: Speaker name (one of 12 predefined speakers)
            text: Text to synthesize with inline emotion tags
                Examples:
                - "Hello! Welcome."
                - "Hello <laugh> this is fun!" (will be normalized)
                - "नमस्ते! [excited] आज का दिन बहुत अच्छा है।"
            temperature: Sampling temperature
            top_k: Top-k sampling
            top_p: Nucleus sampling
            max_tokens: Max BiCodec tokens to generate
            seed: Random seed for reproducibility
            return_timing: If true, returns (audio_bytes, metrics_dict)

        Returns:
            Audio bytes (int16 PCM WAV, 16kHz mono) or None if failed.
            If return_timing is true, returns tuple(audio_bytes, metrics).
        """
        # Delegate to generate_speech() with the same implementation for Spark TTS
        return await self.generate_speech(
            speaker=speaker,


Output:
Chunk ID: f59a94
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 2574
Output:
"""
Modal entrypoint for Veena3 TTS service.

Deploys a GPU-backed ASGI app with:
- True streaming TTS
- Autoscaling (0 to N containers)
- Per-container concurrency (vLLM batching)
- Memory snapshots for fast cold starts

Usage:
    # Deploy to Modal
    modal deploy veena3modal/app.py

    # Serve locally (for testing)
    modal serve veena3modal/app.py

    # Run a single function
    modal run veena3modal/app.py::tts_api
"""

from __future__ import annotations

import os
import modal

# Get the repo root for adding local sources
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# === Image Definition ===
# Full dependencies for TTS inference

image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install(
        "ffmpeg",  # Audio encoding (opus, mp3, flac, mulaw)
        "git",     # For pip git installs
    )
    .pip_install(
        # Core ML/AI
        "torch>=2.0.0",
        "torchaudio>=2.0.0",
        "transformers>=4.35.0",
        "vllm>=0.4.0",

        # FastAPI
        "fastapi[standard]>=0.100.0",
        "uvicorn[standard]>=0.23.0",

        # Audio processing
        "numpy>=1.24.0",
        "scipy>=1.10.0",
        "librosa>=0.10.0",
        "soundfile>=0.12.0",

        # Utils
        "pydantic>=2.0.0",
        "prometheus_client>=0.17.0",
        "supabase>=2.0.0",

        # vLLM dependencies
        "einops>=0.6.0",
        "einx>=0.2.0",

        # SparkTTS dependencies (from external/sparktts/requirements.txt)
        "omegaconf>=2.3.0",
        "safetensors>=0.5.0",
        "soxr>=0.5.0",

        # AP-BWE (super-resolution) dependencies
        "matplotlib>=3.8.0",
        "natsort>=8.0.0",
        "joblib>=1.0.0",
    )
    .run_commands(
        # Ensure CUDA is properly configured
        "pip install --no-cache-dir nvidia-ml-py3 || true"
    )
    # Environment settings
    .env({
        "AUTH_BYPASS_MODE": "true",  # Disable auth for testing
        # sparktts is at /root/external/sparktts, AP-BWE at /root/external/AP-BWE
        "PYTHONPATH": "/root/external/sparktts:/root/external/AP-BWE:/root",
        "_IMAGE_BUILD_VERSION": "2025-12-25-sr-chunking-support",  # Add SR support to generate_speech_chunked
    })
    # Add sparktts package - create proper structure by copying files
    # The original has symlinks that don't work with Modal mounts
    .run_commands(
        # Create sparktts package directory
        "mkdir -p /root/external/sparktts/sparktts/models",
        "mkdir -p /root/external/sparktts/sparktts/utils",
    )
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "sparktts", "sparktts"),
        remote_path="/root/external/sparktts/sparktts",
        copy=True,
    )
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "sparktts", "models"),
        remote_path="/root/external/sparktts/sparktts/models",  # Put models inside sparktts package
        copy=True,
    )
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "sparktts", "utils"),
        remote_path="/root/external/sparktts/sparktts/utils",  # Put utils inside sparktts package
        copy=True,
    )
    # Add AP-BWE for super-resolution (16kHz -> 48kHz)
    .add_local_dir(
        local_path=os.path.join(REPO_ROOT, "external", "AP-BWE"),
        remote_path="/root/external/AP-BWE",
        copy=True,
    )
    # Non-copy local mounts (MUST BE LAST - no build steps after this!)
    .add_local_python_source("veena3modal")
)

# === Volume for Model Weights ===
# Models should be uploaded to this volume before deployment

try:
    model_volume = modal.Volume.from_name("veena3-models", create_if_missing=True)
except Exception:
    model_volume = None  # Will fail at runtime if volume not available

# === Secrets (Optional) ===
# Secrets are optional - the app will run in bypass mode without them
# Create secrets in Modal dashboard: https://modal.com/secrets
# Or use modal.Secret.from_dotenv() for local testing

# === App Definition ===

app = modal.App(
    name="veena3-tts",
    image=image,
)

@app.cls(
    # === GPU Selection ===
    # L40S: Best cost/performance ratio for TTS inference
    # - 48GB VRAM, $1.50/hr (vs A100-80GB at $4.50/hr)
    # - Handles ~12 req/s with <2s p95 latency per container
    # - Memory-bound workload, so faster GPUs have diminishing returns
    gpu="L40S",

    # === CPU Allocation (OPTIMIZATION Dec 2025) ===
    # Explicitly reserve 4 vCPUs for the Python driver process
    # Prevents CPU throttling during vLLM scheduler + streaming loop
    # Modal default is only 0.125 cores which can bottleneck streaming
    cpu=4.0,

    volumes={"/models": model_volume} if model_volume else {},
    secrets=[modal.Secret.from_name("veena3-secrets")],  # Supabase credentials

    # === Autoscaling Configuration ===
    # See .cursor/scaling.md for detailed analysis
    #
    # LOAD TEST RESULTS (Dec 2025):
    # - Cold boot: ~49s (memory snapshot captures imports, not GPU model)
    # - MAX SUSTAINABLE: 280 req/s @ 250 concurrent users (P95=1417ms, 100% success)
    # - Breaking point: 300+ users (429 errors start)
    # - Per container: ~14 req/s (280/20 containers)
    #
    min_containers=1,        # Keep 1 warm 24/7 for instant response
    max_containers=50,       # ~280 req/s capacity at P95 < 3s
    buffer_containers=1,     # REQUIRED: 49s cold boot too slow for burst
    scaledown_window=300,    # 5 min idle before scaledown (user changed from 120)

    # === Timeouts ===
    timeout=120,             # 2 min max per request (reduced from 10 min)
    startup_timeout=600,     # 10 min for model loading (reduced from 20 min)

    # === Memory Optimization ===
    # Memory snapshot for faster cold boots
    # Note: GPU snapshot is experimental (alpha) - disabled for stability
    enable_memory_snapshot=True,
)
# === Per-Container Concurrency ===
# vLLM continuous batching handles concurrent requests efficiently
# Increased from 12 to 20 to maximize GPU utilization while keeping P95 < 3s
@modal.concurrent(max_inputs=20, target_inputs=15)
class TTSService:
    """
    TTS service class with model lifecycle management.

    Uses @modal.enter for model loading (once per container).
    Memory snapshot captures imports and non-GPU state for faster restarts.
    """

    @modal.enter()
    def load_model(self):
        """
        Load TTS model when container starts.

        Called once per container lifecycle.
        Memory snapshot captures import state but not GPU memory.
        """
        import os
        import logging

        logging.basicConfig(level=logging.INFO)
        logger = logging.getLogger(__name__)

        logger.info("🚀 Starting TTS model loading...")

        # Set model paths from Volume mount or env vars
        model_path = os.environ.get(
            "MODEL_PATH",
            "/models/spark_tts_4speaker"
        )
        sr_path = os.environ.get(
            "AP_BWE_CHECKPOINT_DIR",
            "/models/ap_bwe/16kto48k"
        )

        # Check if model exists
        if not os.path.exists(model_path):
            logger.warning(f"Model path not found: {model_path}")
            logger.warning("TTS will return 503 until model is available")
            return

        try:
            from veena3modal.services.tts_runtime import initialize_runtime

            # Initialize with GPU memory optimization
            self.runtime = initialize_runtime(
                model_path=model_path,
                sr_checkpoint_dir=sr_path if os.path.exists(sr_path) else None,
                device="cuda",
                gpu_memory_utilization=0.85,
                enable_sr=os.path.exists(sr_path),
            )

            logger.info(f"✅ TTS model loaded: {self.runtime.model_version}")

        except Exception as e:
            logger.error(f"❌ Model loading failed: {e}")
            import traceback
            traceback.print_exc()

    @modal.asgi_app()
    def serve(self):
        """
        Return FastAPI ASGI app for serving requests.
        """
        from veena3modal.api.fastapi_app import create_app
        return create_app()

# === Standalone Function (Alternative Deployment) ===
# Use this if you don't need @modal.enter lifecycle
# Note: Prefer TTSService class for production (better model lifecycle management)

@app.function(
    gpu="L40S",  # Best cost/performance for TTS
    cpu=4.0,     # Explicit CPU allocation (OPTIMIZATION Dec 2025)
    volumes={"/models": model_volume} if model_volume else {},
    secrets=[modal.Secret.from_name("veena3-secrets")],
    min_containers=0,        # Scale to zero (no min for standalone)
    max_containers=5,        # Lower cap for standalone
    buffer_containers=1,
    scaledown_window=180,    # 3 min
    timeout=120,             # 2 min max
    startup_timeout=600,     # 10 min for model load
)
@modal.asgi_app()
def tts_api():
    """
    Standalone ASGI app function.

    Note: Model will be loaded on first request (higher TTFB).
    For production, use TTSService class with @modal.enter.
    """
    from veena3modal.api.fastapi_app import create_app
    return create_app()

# === Health Check Function ===

@app.function(
    # No GPU needed for health check
)
def health_check():
    """
    Simple health check that can be called from Modal dashboard.
    """
    return {
        "status": "ok",
        "service": "veena3-tts",
        "version": "0.1.0",
    }

# === Debug Function ===

@app.function(
    gpu="L40S",
    volumes={"/models": model_volume} if model_volume else {},
    timeout=300,
)
def debug_imports():
    """
    Debug function to test import paths and file structure.
    Run with: modal run veena3modal/app.py::debug_imports
    """
    import os
    import sys

    results = {
        "pythonpath": os.environ.get("PYTHONPATH", "NOT_SET"),
        "sys_path": sys.path[:10],  # First 10 entries
        "files": {},
        "import_tests": {},
    }

    # Check key directories
    dirs_to_check = [
        "/root",
        "/root/external",
        "/root/external/sparktts",


sed -n '320,760p' scripts/stress_test_local.py


Output:
Chunk ID: cd4057
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 3222
Output:
                        speaker=speaker,
                        text_length=len(text),
                        stream=True,
                    )
        else:
            # Non-streaming: standard request
            response = await client.post(
                url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=timeout,
            )
            latency = (time.time() - start) * 1000

            # Server-reported TTFB
            server_ttfb = 0.0
            ttfb_header = response.headers.get("X-TTFB-ms")
            if ttfb_header:
                try:
                    server_ttfb = float(ttfb_header)
                except ValueError:
                    pass

            # Audio duration from header or estimate
            audio_seconds = 0.0
            audio_sec_header = response.headers.get("X-Audio-Seconds")
            if audio_sec_header:
                try:
                    audio_seconds = float(audio_sec_header)
                except ValueError:
                    pass
            elif response.status_code == 200:
                pcm_bytes = max(0, len(response.content) - 44)
                audio_seconds = pcm_bytes / (16000 * 2)

            if response.status_code == 200:
                return RequestResult(
                    success=True,
                    status_code=200,
                    latency_ms=latency,
                    ttfb_ms=latency,  # Non-streaming: TTFB = total latency
                    server_ttfb_ms=server_ttfb,
                    audio_bytes=len(response.content),
                    audio_seconds=audio_seconds,
                    speaker=speaker,
                    text_length=len(text),
                    stream=False,
                )
            else:
                error_text = response.text[:200] if response.text else f"HTTP {response.status_code}"
                return RequestResult(
                    success=False,
                    status_code=response.status_code,
                    latency_ms=latency,
                    error=error_text,
                    speaker=speaker,
                    text_length=len(text),
                    stream=False,
                )

    except httpx.TimeoutException:
        return RequestResult(
            success=False, status_code=0,
            latency_ms=(time.time() - start) * 1000,
            error="TIMEOUT",
            speaker=speaker, text_length=len(text), stream=stream,
        )
    except Exception as e:
        return RequestResult(
            success=False, status_code=0,
            latency_ms=(time.time() - start) * 1000,
            error=str(e)[:200],
            speaker=speaker, text_length=len(text), stream=stream,
        )

# === Load Runner ===

async def run_level(
    base_url: str,
    num_requests: int,
    concurrency: int,
    text_category: str = "short",
    stream: bool = False,
    gpu_monitor: Optional[GPUMonitor] = None,
) -> LevelResult:
    """Run a single load level and collect results."""
    mode = "streaming" if stream else "non-streaming"
    level_name = f"{concurrency}c-{text_category}-{mode}"

    level = LevelResult(
        level_name=level_name,
        concurrency=concurrency,
        total_requests=num_requests,
        mode=mode,
        text_category=text_category,
    )

    semaphore = asyncio.Semaphore(concurrency)

    # Determine text for each request
    if text_category == "mixed":
        texts = list(TEST_TEXTS.values())
    else:
        texts = [TEST_TEXTS.get(text_category, TEST_TEXTS["short"])]

    async def limited_request(i: int, client: httpx.AsyncClient) -> RequestResult:
        async with semaphore:
            text = texts[i % len(texts)]
            speaker = SPEAKERS[i % len(SPEAKERS)]
            return await make_request(client, base_url, text, speaker, stream=stream)

    # Start GPU monitoring for this level
    if gpu_monitor:
        gpu_monitor.snapshots.clear()

    start_time = time.time()

    # Use connection pooling for realistic load
    limits = httpx.Limits(
        max_connections=concurrency + 10,
        max_keepalive_connections=concurrency,
    )
    async with httpx.AsyncClient(limits=limits) as client:
        tasks = [limited_request(i, client) for i in range(num_requests)]
        results = await asyncio.gather(*tasks, return_exceptions=True)

    level.wall_time_seconds = time.time() - start_time

    for r in results:
        if isinstance(r, Exception):
            level.results.append(RequestResult(
                success=False, status_code=0, latency_ms=0,
                error=str(r)[:200], stream=stream,
            ))
        elif isinstance(r, RequestResult):
            level.results.append(r)

    if gpu_monitor:
        level.gpu_summary = gpu_monitor.summary()

    return level

# === Reporting ===

def print_level_report(level: LevelResult):
    """Print a compact report for one load level."""
    lat = level.latency_stats
    ttfb = level.ttfb_stats

    print(f"\n{'=' * 72}")
    print(f"  {level.level_name.upper()}")
    print(f"  {level.total_requests} requests, {level.concurrency} concurrent, {level.mode}")
    print(f"{'=' * 72}")

    print(f"  Success:     {len(level.successes)}/{level.total_requests} ({level.success_rate:.0%})")
    print(f"  Wall time:   {level.wall_time_seconds:.2f}s")
    print(f"  Throughput:  {level.throughput_rps:.2f} req/s")
    print(f"  Audio total: {level.audio_duration_total:.1f}s")
    print(f"  Eff. RTF:    {level.effective_rtf:.3f}")

    if lat:
        print(f"\n  Latency (ms):")
        print(f"    min={lat['min']:.0f}  avg={lat['avg']:.0f}  "
              f"p50={lat['p50']:.0f}  p95={lat['p95']:.0f}  "
              f"p99={lat['p99']:.0f}  max={lat['max']:.0f}")

    if ttfb and level.mode == "streaming":
        print(f"  TTFB (ms):")
        print(f"    min={ttfb['min']:.0f}  avg={ttfb['avg']:.0f}  "
              f"p50={ttfb['p50']:.0f}  p95={ttfb['p95']:.0f}  max={ttfb['max']:.0f}")

    if level.gpu_summary and "memory_mb" in level.gpu_summary:
        gpu = level.gpu_summary
        mem = gpu["memory_mb"]
        util = gpu["gpu_util_pct"]
        print(f"\n  GPU:")
        print(f"    Memory: {mem['avg']:.0f}MB avg, {mem['max']:.0f}MB peak / {mem['total']:.0f}MB")
        print(f"    Util:   {util['avg']:.0f}% avg, {util['max']:.0f}% peak")
        print(f"    Temp:   {gpu['temperature_c']['avg']:.0f}C avg, {gpu['temperature_c']['max']:.0f}C peak")

    errors = level.error_summary()
    if errors:
        print(f"\n  Errors:")
        for err, count in sorted(errors.items(), key=lambda x: -x[1]):
            print(f"    [{count}x] {err}")

def print_summary_table(all_levels: List[LevelResult]):
    """Print a compact comparison table across all levels."""
    print(f"\n{'=' * 100}")
    print(f"  STRESS TEST SUMMARY")
    print(f"{'=' * 100}")

    header = f"{'Level':<28} {'OK%':>5} {'RPS':>7} {'p50':>7} {'p95':>7} {'p99':>7} {'RTF':>6} {'GPU%':>5} {'GPUmem':>8}"
    print(header)
    print("-" * 100)

    for level in all_levels:
        lat = level.latency_stats
        gpu_util_avg = level.gpu_summary.get("gpu_util_pct", {}).get("avg", 0)
        gpu_mem_max = level.gpu_summary.get("memory_mb", {}).get("max", 0)

        p50 = f"{lat.get('p50', 0):.0f}" if lat else "---"
        p95 = f"{lat.get('p95', 0):.0f}" if lat else "---"
        p99 = f"{lat.get('p99', 0):.0f}" if lat else "---"

        print(
            f"{level.level_name:<28} "
            f"{level.success_rate:>4.0%} "
            f"{level.throughput_rps:>7.2f} "
            f"{p50:>7} "
            f"{p95:>7} "
            f"{p99:>7} "
            f"{level.effective_rtf:>6.3f} "
            f"{gpu_util_avg:>4.0f}% "
            f"{gpu_mem_max:>6.0f}MB"
        )

    print("=" * 100)

# === Main ===

async def run_stress_test(args):
    """Run the full stress test suite."""
    base_url = args.url

    # Verify server is up
    print(f"Checking server at {base_url}...")
    try:
        async with httpx.AsyncClient() as client:
            resp = await client.get(f"{base_url}/v1/tts/health", timeout=10)
            health = resp.json()
            print(f"  Status: {health.get('status')}")
            print(f"  Model:  {health.get('model_version')}")
            print(f"  GPU:    {health.get('gpu_available')}")
            if health.get("status") != "healthy":
                print("WARNING: Server not fully healthy, results may be unreliable")
    except Exception as e:
        print(f"ERROR: Cannot reach server at {base_url}: {e}")
        sys.exit(1)

    # Parse concurrency levels
    levels = [int(x) for x in args.levels.split(",")]

    # Determine test modes
    test_modes = []
    if not args.stream_only:
        test_modes.append(False)  # non-streaming
    if not args.no_stream:
        test_modes.append(True)  # streaming

    # GPU monitor
    gpu_monitor = GPUMonitor(interval_seconds=0.5)
    gpu_monitor.start()

    all_results: List[LevelResult] = []

    # Warmup
    print(f"\nWarming up (2 requests)...")
    warmup = await run_level(base_url, num_requests=2, concurrency=1, text_category="short", stream=False)
    if warmup.success_rate < 1.0:
        print(f"  WARNING: Warmup had failures ({warmup.success_rate:.0%} success)")
    else:
        lat = warmup.latency_stats
        print(f"  OK ({lat.get('avg', 0):.0f}ms avg)")

    # Run each level
    for stream in test_modes:
        mode_label = "STREAMING" if stream else "NON-STREAMING"
        print(f"\n{'#' * 72}")
        print(f"  {mode_label} TESTS")
        print(f"{'#' * 72}")

        for conc in levels:
            # Scale requests: enough to get meaningful stats, cap at reasonable total
            num_requests = max(conc * 3, 10)  # At least 3x concurrency or 10
            num_requests = min(num_requests, 200)  # Cap at 200

            text_cat = args.text if args.text != "mixed" else "mixed"

            print(f"\n>>> {conc} concurrent, {num_requests} total requests, text={text_cat}, stream={stream}")

            level_result = await run_level(
                base_url=base_url,
                num_requests=num_requests,
                concurrency=conc,
                text_category=text_cat,
                stream=stream,
                gpu_monitor=gpu_monitor,
            )

            print_level_report(level_result)
            all_results.append(level_result)

            # Brief cooldown between levels
            if conc < levels[-1]:
                print(f"\n  Cooling down 3s...")
                await asyncio.sleep(3)

    gpu_monitor.stop()

    # Summary
    print_summary_table(all_results)

    # Save results to JSON
    output_path = args.output
    if output_path:
        save_results(all_results, output_path)
        print(f"\nResults saved to {output_path}")

def save_results(all_results: List[LevelResult], path: str):
    """Save results as JSON for later analysis."""
    data = []
    for level in all_results:
        data.append({
            "level": level.level_name,
            "concurrency": level.concurrency,
            "total_requests": level.total_requests,
            "mode": level.mode,
            "text_category": level.text_category,
            "success_rate": level.success_rate,
            "throughput_rps": level.throughput_rps,
            "wall_time_s": level.wall_time_seconds,
            "latency_ms": level.latency_stats,
            "ttfb_ms": level.ttfb_stats,
            "effective_rtf": level.effective_rtf,
            "audio_total_seconds": level.audio_duration_total,
            "errors": level.error_summary(),
            "gpu": level.gpu_summary,
        })

    with open(path, "w") as f:
        json.dump(data, f, indent=2, default=str)

def main():
    parser = argparse.ArgumentParser(
        description="Veena3 TTS Local Stress Test",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--url", default=DEFAULT_URL, help=f"Server URL (default: {DEFAULT_URL})")
    parser.add_argument("--levels", default="1,5,10,20,50", help="Comma-separated concurrency levels (default: 1,5,10,20,50)")
    parser.add_argument("--quick", action="store_true", help="Quick mode: 1,5,10 only")
    parser.add_argument("--text", default="short", choices=["short", "medium", "long", "mixed"], help="Text category (default: short)")
    parser.add_argument("--stream-only", action="store_true", help="Only test streaming")
    parser.add_argument("--no-stream", action="store_true", help="Only test non-streaming")
    parser.add_argument("--output", default="stress_test_results.json", help="Output JSON path")

    args = parser.parse_args()

    if args.quick:
        args.levels = "1,5,10"

    asyncio.run(run_stress_test(args))

if __name__ == "__main__":
    main()


sed -n '1,260p' veena3modal/shared/logging.py


Output:
Chunk ID: e9de49
Wall time: 0.0512 seconds
Process exited with code 0
Original token count: 1755
Output:
"""
Structured logging for Veena3 TTS service.

Features:
- JSON-formatted logs for structured queries
- Request ID tracking via context
- Lifecycle event helpers (request_received, first_audio_emitted, etc.)
- NO PII in logs (text content stored in Supabase, not logs)

Usage:
    from veena3modal.shared.logging import get_logger, log_event, set_request_context

    logger = get_logger(__name__)
    set_request_context(request_id="req-123")
    log_event(logger, "request_received", text_length=100, speaker="lipakshi")
"""

import logging
import json
import sys
import traceback
from datetime import datetime, timezone
from typing import Any, Dict, Optional
from contextvars import ContextVar

# Context variable for request-scoped data
_request_context: ContextVar[Dict[str, Any]] = ContextVar('request_context', default={})

def set_request_context(request_id: Optional[str] = None, **kwargs) -> None:
    """
    Set request context for current async task.

    Args:
        request_id: Unique request identifier
        **kwargs: Additional context fields (user_id, etc.)
    """
    ctx = {"request_id": request_id, **kwargs}
    _request_context.set(ctx)

def get_request_context() -> Dict[str, Any]:
    """Get current request context."""
    return _request_context.get()

def clear_request_context() -> None:
    """Clear request context after request completes."""
    _request_context.set({})

class JSONFormatter(logging.Formatter):
    """
    JSON log formatter for structured logging.

    Output format:
    {
        "timestamp": "2024-01-01T12:00:00.000Z",
        "level": "INFO",
        "logger": "veena3modal.api",
        "message": "Request received",
        "request_id": "req-123",
        ...extra fields...
    }
    """

    # Fields to exclude from extra (already handled separately)
    RESERVED_ATTRS = {
        'args', 'asctime', 'created', 'exc_info', 'exc_text', 'filename',
        'funcName', 'levelname', 'levelno', 'lineno', 'module', 'msecs',
        'message', 'msg', 'name', 'pathname', 'process', 'processName',
        'relativeCreated', 'stack_info', 'thread', 'threadName',
        'taskName',  # Python 3.12+
    }

    def format(self, record: logging.LogRecord) -> str:
        """Format log record as JSON."""
        # Build base log entry
        log_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
        }

        # Add request context if available
        ctx = get_request_context()
        if ctx.get("request_id"):
            log_entry["request_id"] = ctx["request_id"]

        # Add extra fields from record
        for key, value in record.__dict__.items():
            if key not in self.RESERVED_ATTRS and not key.startswith('_'):
                try:
                    # Ensure value is JSON serializable
                    json.dumps(value)
                    log_entry[key] = value
                except (TypeError, ValueError):
                    log_entry[key] = str(value)

        # Add exception info if present
        if record.exc_info:
            log_entry["exception"] = ''.join(traceback.format_exception(*record.exc_info))

        return json.dumps(log_entry, default=str)

def get_logger(name: str) -> logging.Logger:
    """
    Get a logger configured for structured JSON output.

    Args:
        name: Logger name (typically __name__)

    Returns:
        Configured logger instance
    """
    logger = logging.getLogger(name)

    # Only configure if no handlers exist (avoid duplicate setup)
    if not logger.handlers:
        handler = logging.StreamHandler(sys.stdout)
        handler.setFormatter(JSONFormatter())
        logger.addHandler(handler)
        logger.setLevel(logging.INFO)

    return logger

def log_event(
    logger: logging.Logger,
    event_type: str,
    level: int = logging.INFO,
    request_id: Optional[str] = None,
    **kwargs
) -> None:
    """
    Log a structured event.

    Args:
        logger: Logger instance
        event_type: Event name (e.g., "request_received", "first_audio_emitted")
        level: Log level (default: INFO)
        request_id: Override request_id from context
        **kwargs: Event-specific fields
    """
    extra = {"event": event_type, **kwargs}

    if request_id:
        extra["request_id"] = request_id

    logger.log(level, event_type, extra=extra)

def create_lifecycle_event(
    event_type: str,
    request_id: str,
    **kwargs
) -> Dict[str, Any]:
    """
    Create a lifecycle event dictionary.

    Lifecycle events:
    - request_received: TTS request started
    - auth_validated: API key validated
    - generation_started: Model inference began
    - first_audio_emitted: First audio chunk sent (TTFB)
    - request_completed: Request finished successfully
    - request_failed: Request failed with error

    Args:
        event_type: Event name
        request_id: Request identifier
        **kwargs: Event-specific fields

    Returns:
        Event dictionary ready for logging
    """
    return {
        "event": event_type,
        "request_id": request_id,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        **kwargs
    }

# Pre-configured logger for the TTS service
tts_logger = get_logger("veena3modal.tts")

def log_request_received(
    request_id: str,
    text_length: int,
    speaker: str,
    stream: bool,
    format: str,
    **kwargs
) -> None:
    """Log request_received lifecycle event."""
    log_event(
        tts_logger,
        "request_received",
        request_id=request_id,
        text_length=text_length,
        speaker=speaker,
        stream=stream,
        format=format,
        **kwargs
    )

def log_first_audio_emitted(
    request_id: str,
    ttfb_ms: int,
    chunk_size_bytes: int,
    **kwargs
) -> None:
    """Log first_audio_emitted lifecycle event (TTFB marker)."""
    log_event(
        tts_logger,
        "first_audio_emitted",
        request_id=request_id,
        ttfb_ms=ttfb_ms,
        chunk_size_bytes=chunk_size_bytes,
        **kwargs
    )

def log_request_completed(
    request_id: str,
    status_code: int,
    total_duration_ms: int,
    audio_duration_seconds: float,
    rtf: float,
    chunks_sent: int = 1,
    **kwargs
) -> None:
    """Log request_completed lifecycle event."""
    log_event(
        tts_logger,
        "request_completed",
        request_id=request_id,
        status_code=status_code,
        total_duration_ms=total_duration_ms,
        audio_duration_seconds=audio_duration_seconds,
        rtf=rtf,
        chunks_sent=chunks_sent,
        **kwargs
    )

def log_request_failed(
    request_id: str,
    status_code: int,
    error_code: str,
    error_message: str,
    **kwargs
) -> None:
    """Log request_failed lifecycle event."""
    log_event(


rg -n "VEENA3_TIMING_LOGS|tts_stream|tts_runtime|llm_batch" scripts veena3modal | head -n 200