runtime/turn_manager.py

import logging
from typing import Any
from models.event import Event, EventType
from models.session_state import TurnState
from queues import queue_names

logger = logging.getLogger("s2s_server")

class TurnManager:
    def __init__(self, context: Any):
        self.context = context
        self.state = TurnState.IDLE

        # Subscribe handlers
        self.context.event_bus.subscribe(EventType.USER_SPEECH_START, self.on_user_speech_start)
        self.context.event_bus.subscribe(EventType.USER_SPEECH_END, self.on_user_speech_end)
        self.context.event_bus.subscribe(EventType.STABLE_TRANSCRIPT, self.on_stable_transcript)
        self.context.event_bus.subscribe(EventType.LLM_STARTED, self.on_llm_started)
        self.context.event_bus.subscribe(EventType.BARGE_IN, self.on_barge_in)
        self.context.event_bus.subscribe(EventType.CANCEL_ALL, self.on_cancel_all)
        self.context.event_bus.subscribe(EventType.TTS_FINISHED, self.on_tts_finished)
        self.context.event_bus.subscribe(EventType.TTS_STARTED, self.on_tts_started)

    async def on_user_speech_start(self, event: Event):
        import uuid
        self.context.turn_id = str(uuid.uuid4())[:8]
        self.context.logger.info(
            f"[{self.context.session_id}] TurnManager: User speech start detected (state: {self.state.name}, turn_id: {self.context.turn_id})"
        )

        # Interruption check
        if self.context.session_config.allow_interruptions and self.state in (
            TurnState.THINKING,
            TurnState.SPEAKING,
        ):
            await self.context.event_bus.publish(EventType.BARGE_IN)

        self.state = TurnState.USER_SPEAKING
        await self.context.websocket.send_json({"type": "status", "message": "Listening..."})

    async def on_user_speech_end(self, event: Event):
        self.context.logger.info(f"[{self.context.session_id}] TurnManager: User speech end detected.")
        full_audio = event.payload
        self.context.queue_registry.get(queue_names.TRANSCRIPT).put_nowait(full_audio)

    async def on_stable_transcript(self, event: Event):
        text = event.payload
        self.context.logger.info(
            f"[{self.context.session_id}] TurnManager: Stable transcript received: '{text}'"
        )
        await self.context.websocket.send_json({"type": "transcription", "text": text})

        # Feed the text to the LLM queue
        self.context.queue_registry.get(queue_names.LLM).put_nowait(text)

    async def on_llm_started(self, event: Event):
        self.state = TurnState.THINKING
        await self.context.websocket.send_json({"type": "status", "message": "Thinking..."})

    async def on_tts_started(self, event: Event):
        self.state = TurnState.SPEAKING

    async def on_barge_in(self, event: Event):
        self.context.logger.info(f"[{self.context.session_id}] TurnManager: Barge-in triggered.")
        self.state = TurnState.INTERRUPTED
        await self.context.websocket.send_json({"type": "interrupt"})
        await self.context.event_bus.publish(EventType.CANCEL_ALL)

    async def on_cancel_all(self, event: Event):
        self.context.logger.info(f"[{self.context.session_id}] TurnManager: Cancelling active response.")
        # Cancel tasks in task supervisor
        await self.context.task_supervisor.cancel_tasks("llm_generation", "tts_generation")
        # Perform pipeline queue and accumulator flushes
        self.context.cleanup_manager.perform_cleanup()

    async def on_tts_finished(self, event: Event):
        tts_queue = self.context.queue_registry.get(queue_names.TTS)
        if tts_queue.empty() and self.state != TurnState.USER_SPEAKING:
            self.context.logger.info(f"[{self.context.session_id}] TurnManager: Turn complete. Returning to IDLE.")
            self.state = TurnState.IDLE
            await self.context.websocket.send_json({"type": "done"})
            await self.context.event_bus.publish(EventType.TURN_COMPLETED)