Output:
Chunk ID: 380526
Wall time: 0.0506 seconds
Process exited with code 0
Original token count: 4331
Output:
        processed_token_count = 0  # Track how many vLLM tokens we've parsed
        total_audio_chunks = 0

        # BiCodec streaming configuration
        EXPECTED_GLOBAL_COUNT = 32  # BiCodec uses 32 fixed global tokens (per paper)
        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8, -120ms TTFB)
        # OPTIMIZATION: Increased from 24 to 48 to reduce O(n) re-decode calls by 2x.
        # At 50 TPS, 48 tokens = ~960ms audio per chunk. Cuts total decode calls roughly in half.
        DECODE_INTERVAL = 48
        CROSSFADE_MS = 50
        # OPTIMIZATION: Windowed decode - only decode last WINDOW_SIZE tokens instead of ALL.
        # BiCodec WaveGenerator has finite receptive field (~64 tokens based on kernel sizes).
        # Window must be >= receptive_field + DECODE_INTERVAL to avoid boundary artifacts.
        WINDOW_SIZE = 128  # 128 tokens covers receptive field + margin for crossfade

        # Generate tokens with vLLM (streaming)
        if request_id is None:
            import uuid
            request_id = f"bicodec-stream-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"

        # Reuse pre-warmed token parser (singleton, created once in __init__)
        token_parser = self.token_parser

        t_gen_start = time.time()
        batch_wall_ms: List[float] = []
        batch_token_counts: List[int] = []
        decode_calls = 0
        decode_wall_ms = 0.0
        decode_gpu_ms = 0.0
        parse_wall_ms = 0.0

        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Track timing
        t_first_token = None
        t_first_chunk_ready = None
        t_first_chunk_decoded = None
        iteration_count = 0
        last_decode_count = 0

        # Track audio samples EMITTED to the user (not just decoded)
        # BiCodec produces ~320 samples per semantic token at 16kHz
        # CRITICAL: This tracks samples that have been YIELDED, not held in tail
        total_samples_emitted_to_user = 0
        # Keep a tail to enable crossfading between consecutive emitted chunks
        previous_chunk_tail: Optional[bytes] = None

        # Stream tokens with BiCodec extraction
        async for request_output in results_generator:
            iteration_count += 1
            t_iter = time.time()

            # Track first token received
            if t_first_token is None:
                t_first_token = t_iter

            # Extract generated token IDs
            generated_ids = request_output.outputs[0].token_ids

            parse_start = time.time()
            # OPTIMIZATION: Only process NEW tokens (incremental parsing)
            # Old code: tokenizer.decode(ALL) + re.findall(ALL) = O(n²) per stream
            # New code: parse only new tokens = O(1) per token
            new_token_ids = generated_ids[processed_token_count:]
            processed_token_count = len(generated_ids)

            # Parse new tokens and append to buffers (in-place modification)
            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)
            parse_wall_ms += (time.time() - parse_start) * 1000
            batch_token_counts.append(len(new_token_ids))
            batch_wall_ms.append(0 if t_iter is None else (time.time() - t_iter) * 1000)

            # TRUE STREAMING with BiCodec two-phase generation:
            # Phase 1: Generate 32 global tokens (pre-roll)
            # Phase 2: Stream semantic tokens incrementally

            # Check if we have the 32 global tokens (pre-roll complete)
            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
                # Pre-roll complete! Now we can stream semantic tokens
                if t_first_chunk_ready is None and len(semantic_buffer) >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                    t_first_chunk_ready = time.time()

                # Stream semantic tokens incrementally
                # Only decode if we have enough NEW semantic tokens since last decode
                semantic_count = len(semantic_buffer)
                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                    # Check if we should decode (every DECODE_INTERVAL new semantic tokens)
                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
                        last_decode_count = semantic_count

                        # OPTIMIZATION: Windowed decode instead of full re-decode
                        # Only decode the last WINDOW_SIZE tokens (O(1) per chunk instead of O(n))
                        # For first few chunks when total tokens < WINDOW_SIZE, decode all
                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]

                        if len(semantic_buffer) <= WINDOW_SIZE:
                            # Early phase: decode all tokens (window not yet relevant)
                            decode_semantic = semantic_buffer
                        else:
                            # Windowed: only decode last WINDOW_SIZE tokens
                            decode_semantic = semantic_buffer[-WINDOW_SIZE:]

                        t_decode_start = time.time()
                        if emit_timing and torch.cuda.is_available():
                            torch.cuda.synchronize()
                            decode_start_event = torch.cuda.Event(enable_timing=True)
                            decode_end_event = torch.cuda.Event(enable_timing=True)
                            decode_start_event.record()
                        audio_bytes = await self.snac_decoder.decode_single_async(
                            semantic_ids=decode_semantic,
                            global_ids=window_global,
                            trim_warmup=False,
                            use_sliding_window=False,
                        )
                        t_decode_end = time.time()
                        decode_calls += 1
                        decode_wall_ms += (t_decode_end - t_decode_start) * 1000
                        if emit_timing and torch.cuda.is_available():
                            decode_end_event.record()
                            torch.cuda.synchronize()
                            decode_gpu_ms += decode_start_event.elapsed_time(decode_end_event)

                        if audio_bytes:
                            total_samples_decoded = len(audio_bytes) // 2

                            if len(semantic_buffer) <= WINDOW_SIZE:
                                # Early phase: extract new samples from full decode
                                if total_samples_decoded > total_samples_emitted_to_user:
                                    new_bytes_start = total_samples_emitted_to_user * 2
                                    new_audio_bytes = audio_bytes[new_bytes_start:]
                                else:
                                    new_audio_bytes = b""
                            else:
                                # Windowed phase: take only the tail portion of the window decode
                                # The last DECODE_INTERVAL tokens' worth of audio is the new content
                                # Each token produces ~320 samples at 16kHz/50TPS
                                new_samples_approx = DECODE_INTERVAL * 320
                                if total_samples_decoded > new_samples_approx:
                                    new_bytes_start = (total_samples_decoded - new_samples_approx) * 2
                                else:
                                    new_bytes_start = 0
                                new_audio_bytes = audio_bytes[new_bytes_start:]

                            if new_audio_bytes:
                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                                    previous_chunk_tail,
                                    new_audio_bytes,
                                    sample_rate_hz=16000,
                                    crossfade_ms=CROSSFADE_MS,
                                )

                                samples_emitted_in_this_chunk = len(to_emit) // 2
                                total_samples_emitted_to_user += samples_emitted_in_this_chunk

                                if to_emit:
                                    total_audio_chunks += 1
                                    if t_first_chunk_decoded is None:
                                        t_first_chunk_decoded = t_decode_end

                                    if emit_timing:
                                        logger.info(
                                            "tts_stream_decode_batch",
                                            extra={
                                                "speaker": speaker,
                                                "request_id": request_id,
                                                "batch_index": decode_calls,
                                                "decode_wall_ms": int((t_decode_end - t_decode_start) * 1000),
                                                "decode_gpu_ms": int(decode_start_event.elapsed_time(decode_end_event))
                                                if torch.cuda.is_available() else None,
                                                "semantic_count": len(semantic_buffer),
                                            },
                                        )

                                    yield to_emit

        # Final chunk: decode any remaining semantic tokens
        # Use the same 32 global tokens (they get pooled)
        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
            audio_bytes = self.snac_decoder.decode_streaming(
                semantic_ids=semantic_buffer,  # All semantic tokens
                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],  # Use first 32 global tokens
                use_sliding_window=False,  # Get all samples
                trim_warmup=False
            )
            if audio_bytes:
                # Only yield NEW samples
                total_samples_decoded = len(audio_bytes) // 2

                if total_samples_decoded > total_samples_emitted_to_user:
                    new_bytes_start = total_samples_emitted_to_user * 2
                    new_audio_bytes = audio_bytes[new_bytes_start:]

                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        new_audio_bytes,
                        sample_rate_hz=16000,
                        crossfade_ms=CROSSFADE_MS,
                    )

                    samples_emitted_in_final = len(to_emit) // 2
                    total_samples_emitted_to_user += samples_emitted_in_final

                    if to_emit:
                        total_audio_chunks += 1
                        yield to_emit

        # Flush any remaining tail that was held back for crossfade
        if previous_chunk_tail:
            total_audio_chunks += 1
            tail_samples = len(previous_chunk_tail) // 2
            total_samples_emitted_to_user += tail_samples
            yield previous_chunk_tail

        # Log completion summary
        t_end = time.time()
        audio_duration_s = total_samples_emitted_to_user / 16000  # 16kHz sample rate
        ttfb_ms = (t_first_chunk_decoded - t_start) * 1000 if t_first_chunk_decoded else 0
        print(f"🎵 Streaming complete: {audio_duration_s:.2f}s audio, TTFB: {ttfb_ms:.0f}ms, {total_audio_chunks} chunks")

    async def generate_speech_stream_indic_first_chunk(
        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 = DEFAULT_REPETITION_PENALTY,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[tuple, None]:
        """
        Generate speech for the FIRST chunk of multi-chunk text, capturing global tokens.

        This method is specifically for chunked generation: it yields (audio_bytes, global_ids)
        tuples instead of just audio_bytes. The caller captures global_ids from the first
        yield and passes them to generate_speech_stream_indic_continuation() for subsequent chunks.

        Use Case (chunked streaming with voice consistency):
            globals_captured = None
            async for audio_bytes, global_ids in pipeline.generate_speech_stream_indic_first_chunk(...):
                if globals_captured is None and global_ids:
                    globals_captured = global_ids  # Capture once
                yield audio_bytes
            # Now use globals_captured for subsequent chunks

        Args:
            speaker: Speaker name
            text: First text chunk to synthesize
            temperature: Sampling temperature
            top_k: Top-k sampling
            top_p: Nucleus sampling
            max_tokens: Max tokens to generate
            repetition_penalty: Prevent repetition
            seed: Random seed for reproducibility

        Yields:
            Tuple of (audio_bytes, global_ids)
            - audio_bytes: Raw PCM audio (int16, 16kHz)
            - global_ids: List of 32 global token IDs (populated after first decode, else empty)

        Thread Safety:
            This method is stateless and thread-safe. Each call creates its own
            request-scoped state. No global state is modified or shared.
        """
        import time
        t_start = time.time()

        # Build prompt using Indic prompt builder
        prompt = self.prompt_builder.build_prefix(speaker, text)

        # Configure sampling
        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
        sampling_params = SamplingParams(
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            min_tokens=DEFAULT_MIN_TOKENS,
            repetition_penalty=repetition_penalty,
            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
            skip_special_tokens=False,
            seed=seed,
        )

        # BiCodec token buffers
        semantic_buffer = []
        global_buffer = []
        processed_token_count = 0  # Track how many vLLM tokens we've parsed
        total_audio_chunks = 0

        # BiCodec streaming configuration
        EXPECTED_GLOBAL_COUNT = 32
        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
        CROSSFADE_MS = 50

        # Generate unique request ID
        import uuid
        request_id = f"bicodec-first-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"

        # Reuse pre-warmed token parser (singleton, created once in __init__)
        token_parser = self.token_parser

        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Track state
        last_decode_count = 0
        total_samples_emitted_to_user = 0
        previous_chunk_tail: Optional[bytes] = None
        captured_globals: List[int] = []  # Will be populated once we have 32 globals

        # Stream tokens with BiCodec extraction
        async for request_output in results_generator:
            generated_ids = request_output.outputs[0].token_ids

            # OPTIMIZATION: Incremental parsing (O(1) per new token)
            new_token_ids = generated_ids[processed_token_count:]
            processed_token_count = len(generated_ids)
            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer)

            # Capture global tokens once we have all 32
            if not captured_globals and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
                captured_globals = global_buffer[:EXPECTED_GLOBAL_COUNT]

            # Process if we have all global tokens
            if len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
                semantic_count = len(semantic_buffer)
                if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                    if semantic_count - last_decode_count >= DECODE_INTERVAL:
                        last_decode_count = semantic_count

                        all_semantic = semantic_buffer
                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]


# Collaboration Mode: Default

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

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

## request_user_input availability

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

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

Summary: auto

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


I’ve confirmed the key timing hooks are partially wired but a few paths are inconsistent; I’m now pulling the remaining sections to close gaps in streaming/entrypoint callsites and then I’ll apply one consolidated patch.


cd /home/ubuntu/veenaModal && sed -n '620,980p' veena3modal/core/streaming_pipeline.py && sed -n '1,260p' veena3modal/api/fastapi_app.py && sed -n '1,260p' veena3modal/api/websocket_handler.py


Output:
Chunk ID: 1bc55d
Wall time: 0.0507 seconds
Process exited with code 0
Original token count: 7511
Output:
                        window_global = global_buffer[:EXPECTED_GLOBAL_COUNT]

                        if self.snac_decoder.enable_batching:
                            audio_bytes = await self.snac_decoder.decode_single_async(
                                semantic_ids=all_semantic,
                                global_ids=window_global,
                                trim_warmup=False,
                                use_sliding_window=False
                            )
                        else:
                            audio_bytes = self.snac_decoder.decode_streaming(
                                semantic_ids=all_semantic,
                                global_ids=window_global,
                                use_sliding_window=False,
                                trim_warmup=False
                            )

                        if audio_bytes:
                            total_samples_decoded = len(audio_bytes) // 2

                            if total_samples_decoded > total_samples_emitted_to_user:
                                new_bytes_start = total_samples_emitted_to_user * 2
                                new_audio_bytes = audio_bytes[new_bytes_start:]

                                to_emit, previous_chunk_tail = crossfade_bytes_int16(
                                    previous_chunk_tail,
                                    new_audio_bytes,
                                    sample_rate_hz=16000,
                                    crossfade_ms=CROSSFADE_MS,
                                )

                                samples_emitted_in_this_chunk = len(to_emit) // 2
                                total_samples_emitted_to_user += samples_emitted_in_this_chunk

                                if to_emit:
                                    total_audio_chunks += 1
                                    # Yield tuple: (audio_bytes, global_ids)
                                    yield (to_emit, captured_globals)

        # Final chunk processing
        if len(semantic_buffer) > last_decode_count and len(global_buffer) >= EXPECTED_GLOBAL_COUNT:
            audio_bytes = self.snac_decoder.decode_streaming(
                semantic_ids=semantic_buffer,
                global_ids=global_buffer[:EXPECTED_GLOBAL_COUNT],
                use_sliding_window=False,
                trim_warmup=False
            )
            if audio_bytes:
                total_samples_decoded = len(audio_bytes) // 2

                if total_samples_decoded > total_samples_emitted_to_user:
                    new_bytes_start = total_samples_emitted_to_user * 2
                    new_audio_bytes = audio_bytes[new_bytes_start:]

                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        new_audio_bytes,
                        sample_rate_hz=16000,
                        crossfade_ms=CROSSFADE_MS,
                    )

                    samples_emitted_in_final = len(to_emit) // 2
                    total_samples_emitted_to_user += samples_emitted_in_final

                    if to_emit:
                        total_audio_chunks += 1
                        yield (to_emit, captured_globals)

        # Flush remaining tail
        if previous_chunk_tail:
            total_audio_chunks += 1
            tail_samples = len(previous_chunk_tail) // 2
            total_samples_emitted_to_user += tail_samples
            yield (previous_chunk_tail, captured_globals)

        t_end = time.time()
        audio_duration_s = total_samples_emitted_to_user / 16000
        print(f"🎵 First chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks, captured {len(captured_globals)} globals")

    async def generate_speech_stream_indic_continuation(
        self,
        speaker: str,
        text: str,
        global_ids: List[int],
        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 = DEFAULT_REPETITION_PENALTY,
        seed: Optional[int] = None,
    ) -> AsyncGenerator[bytes, None]:
        """
        Generate speech for CONTINUATION chunks using pre-captured global tokens.

        This is the key method for voice consistency in chunked generation:
        - Uses global_ids from the first chunk to maintain identical voice
        - Model generates only semantic tokens (skips global token generation)
        - Ensures no voice drift across chunks

        CRITICAL for production:
        - global_ids is request-scoped, passed explicitly (no shared state)
        - Thread-safe: each request has its own state
        - No global caching that could cause cross-request contamination

        Args:
            speaker: Speaker name (must match first chunk)
            text: Continuation text chunk to synthesize
            global_ids: 32 global token IDs captured from first chunk
            temperature: Sampling temperature
            top_k: Top-k sampling
            top_p: Nucleus sampling
            max_tokens: Max tokens to generate
            repetition_penalty: Prevent repetition
            seed: Random seed (use same as first chunk for consistency)

        Yields:
            Audio bytes (int16 PCM, 16kHz mono)

        Raises:
            ValueError: If global_ids doesn't contain exactly 32 tokens

        Thread Safety:
            Fully thread-safe. All state is request-scoped and passed explicitly.
        """
        import time
        t_start = time.time()

        # Validate global tokens
        EXPECTED_GLOBAL_COUNT = 32
        if len(global_ids) != EXPECTED_GLOBAL_COUNT:
            raise ValueError(
                f"Expected exactly {EXPECTED_GLOBAL_COUNT} global tokens, got {len(global_ids)}. "
                f"global_ids must be captured from first chunk generation."
            )

        # Build prompt WITH pre-filled global tokens
        # This tells the model to skip global generation and go straight to semantic tokens
        prompt = self.prompt_builder.build_prefix_with_globals(speaker, text, global_ids)

        # Configure sampling
        # OPTIMIZATION: Use Spark TTS stop token, not legacy SNAC CODE_END_TOKEN_ID
        sampling_params = SamplingParams(
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            max_tokens=max_tokens,
            min_tokens=DEFAULT_MIN_TOKENS,
            repetition_penalty=repetition_penalty,
            stop=TRAINING_STOP_TOKEN_IDS,  # "<|im_end|>" - matches non-streaming pipeline
            skip_special_tokens=False,
            seed=seed,
        )

        # BiCodec token buffers
        semantic_buffer = []
        global_buffer_unused = []  # Not used in continuation, but needed for parser API
        processed_token_count = 0  # Track how many vLLM tokens we've parsed
        total_audio_chunks = 0

        # BiCodec streaming configuration
        MIN_SEMANTIC_FOR_FIRST_CHUNK = 10  # OPTIMIZATION: Lowered from 16 (decoder min is 8)
        DECODE_INTERVAL = 48  # OPTIMIZATION: Increased from 24 to reduce O(n) re-decode calls
        CROSSFADE_MS = 50

        # Generate unique request ID
        import uuid
        request_id = f"bicodec-cont-{uuid.uuid4().hex[:8]}-{int(time.time() * 1000000)}"

        # Reuse pre-warmed token parser (singleton, created once in __init__)
        token_parser = self.token_parser

        results_generator = self.model.engine.generate(
            prompt=prompt,
            sampling_params=sampling_params,
            request_id=request_id,
        )

        # Track state
        last_decode_count = 0
        total_samples_emitted_to_user = 0
        previous_chunk_tail: Optional[bytes] = None

        # Stream tokens - model will generate ONLY semantic tokens (globals are pre-filled)
        async for request_output in results_generator:
            generated_ids = request_output.outputs[0].token_ids

            # OPTIMIZATION: Incremental parsing (O(1) per new token)
            # Note: In continuation, model only generates semantic tokens (globals pre-filled in prompt)
            new_token_ids = generated_ids[processed_token_count:]
            processed_token_count = len(generated_ids)
            token_parser.parse_incremental(new_token_ids, semantic_buffer, global_buffer_unused)

            semantic_count = len(semantic_buffer)
            if semantic_count >= MIN_SEMANTIC_FOR_FIRST_CHUNK:
                if semantic_count - last_decode_count >= DECODE_INTERVAL:
                    last_decode_count = semantic_count

                    # Use pre-captured global tokens for decoding
                    if self.snac_decoder.enable_batching:
                        audio_bytes = await self.snac_decoder.decode_single_async(
                            semantic_ids=semantic_buffer,
                            global_ids=global_ids,  # Use captured globals
                            trim_warmup=False,
                            use_sliding_window=False
                        )
                    else:
                        audio_bytes = self.snac_decoder.decode_streaming(
                            semantic_ids=semantic_buffer,
                            global_ids=global_ids,  # Use captured globals
                            use_sliding_window=False,
                            trim_warmup=False
                        )

                    if audio_bytes:
                        total_samples_decoded = len(audio_bytes) // 2

                        if total_samples_decoded > total_samples_emitted_to_user:
                            new_bytes_start = total_samples_emitted_to_user * 2
                            new_audio_bytes = audio_bytes[new_bytes_start:]

                            to_emit, previous_chunk_tail = crossfade_bytes_int16(
                                previous_chunk_tail,
                                new_audio_bytes,
                                sample_rate_hz=16000,
                                crossfade_ms=CROSSFADE_MS,
                            )

                            samples_emitted_in_this_chunk = len(to_emit) // 2
                            total_samples_emitted_to_user += samples_emitted_in_this_chunk

                            if to_emit:
                                total_audio_chunks += 1
                                yield to_emit

        # Final chunk processing
        if len(semantic_buffer) > last_decode_count:
            audio_bytes = self.snac_decoder.decode_streaming(
                semantic_ids=semantic_buffer,
                global_ids=global_ids,  # Use captured globals
                use_sliding_window=False,
                trim_warmup=False
            )
            if audio_bytes:
                total_samples_decoded = len(audio_bytes) // 2

                if total_samples_decoded > total_samples_emitted_to_user:
                    new_bytes_start = total_samples_emitted_to_user * 2
                    new_audio_bytes = audio_bytes[new_bytes_start:]

                    to_emit, previous_chunk_tail = crossfade_bytes_int16(
                        previous_chunk_tail,
                        new_audio_bytes,
                        sample_rate_hz=16000,
                        crossfade_ms=CROSSFADE_MS,
                    )

                    samples_emitted_in_final = len(to_emit) // 2
                    total_samples_emitted_to_user += samples_emitted_in_final

                    if to_emit:
                        total_audio_chunks += 1
                        yield to_emit

        # Flush remaining tail
        if previous_chunk_tail:
            total_audio_chunks += 1
            tail_samples = len(previous_chunk_tail) // 2
            total_samples_emitted_to_user += tail_samples
            yield previous_chunk_tail

        t_end = time.time()
        audio_duration_s = total_samples_emitted_to_user / 16000
        print(f"🎵 Continuation chunk complete: {audio_duration_s:.2f}s audio, {total_audio_chunks} chunks (using cached globals)")
"""
FastAPI ASGI app factory for Modal deployment.

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

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

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

import time
import uuid
from typing import Optional

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

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

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

logger = get_logger(__name__)

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

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

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

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

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

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

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

    global _startup_time
    _startup_time = time.time()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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


Summary: auto


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


Output:
Chunk ID: 456d52
Wall time: 0.0508 seconds
Process exited with code 0
Original token count: 5257
Output:
            )

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

        # Get resolved speaker
        speaker = req.get_resolved_speaker()

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

        text = req.get_normalized_text(normalizer_func=normalizer_func)

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

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

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

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

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

            # Generate speech (non-streaming)
            if req.chunking:
                audio_bytes, metrics = await tts_runtime.generate_speech_chunked(
                    text=text,
                    speaker=speaker,
                    temperature=req.temperature,
                    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,
                ):
                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":
                            await websocket.send_text(json.dumps({"event": "pong"}))
                    except asyncio.TimeoutError:
                        continue
                    except Exception:
                        break
            except Exception:
                pass

        cancel_task = asyncio.create_task(listen_for_cancel())

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    finally:
        clear_request_context()

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

    Call this from create_app() to enable WebSocket support.

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

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

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

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

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

        await handle_websocket_tts(websocket, api_key)

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


cd /home/ubuntu/veenaModal && nl -ba veena3modal/services/tts_runtime.py | sed -n '1,260p' && nl -ba veena3modal/services/tts_runtime.py | sed -n '260,560p' && nl -ba veena3modal/core/streaming_pipeline.py | sed -n '1,240p' && nl -ba veena3modal/core/streaming_pipeline.py | sed -n '240,720p' && nl -ba veena3modal/api/fastapi_app.py | sed -n '1,380p' && nl -ba veena3modal/api/fastapi_app.py | sed -n '380,520p' && nl -ba veena3modal/api/websocket_handler.py | sed -n '1,260p' && nl -ba veena3modal/api/websocket_handler.py | sed -n '260,540p'