first

speech_to_speech/api/openai_realtime/pipeline_unit.py

@@ -0,0 +1,70 @@
+import asyncio
+from queue import Queue
+from threading import Event
+from typing import Any, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from speech_to_speech.api.openai_realtime.service import RealtimeService
+from speech_to_speech.pipeline.cancel_scope import CancelScope
+
+
+class SessionState(BaseModel):
+    """Per-websocket ephemeral state.
+
+    Created when the route handler claims a PipelineUnit on `ws.accept()`; dropped
+    when the websocket disconnects. Holding the websocket reference, the service
+    session id, and any send-loop scratch (pending_output_item) here ensures these
+    fields share one lifecycle — a stale value can't outlive its session.
+
+    `drained` is set by the send loop when SESSION_END travels through the handler
+    chain back to the output queue; the route handler awaits it before clearing
+    `PipelineUnit.session`, so a new client cannot claim the unit until in-flight
+    work from this session has fully reset.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    websocket: Any
+    session_id: str = ""
+    pending_output_item: Any = None
+    drained: asyncio.Event = Field(default_factory=asyncio.Event)
+    # Wall-clock time when the ws disconnected (route handler released its claim).
+    # `None` while the client is still active. Used by /v1/pool to surface stuck
+    # units (handlers haven't finished propagating SESSION_END).
+    released_at: Optional[float] = None
+    # Wall-clock time when the drain wait gave up and quarantined the unit
+    # (SESSION_END_QUARANTINE_TIMEOUT_S elapsed). The unit stays unclaimable —
+    # its handlers may still emit this session's output — until SESSION_END
+    # actually drains. Reported as "stuck" by /v1/pool.
+    quarantined_at: Optional[float] = None
+
+
+class PipelineUnit(BaseModel):
+    """One isolated realtime pipeline.
+
+    Each unit owns its queues, events, RealtimeService, and the chain of handler
+    instances (VAD, STT, transcription notifier, LM, LM output processor, TTS).
+    Lives inside the pool managed by RealtimeServer; the websocket route handler
+    claims a free unit (`session is None`) on `accept` and releases it on disconnect
+    by setting `session` back to None.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    index: int
+    service: RealtimeService
+    cancel_scope: CancelScope
+    should_listen: Event
+    response_playing: Event
+    input_queue: Queue
+    output_queue: Queue
+    text_output_queue: Queue
+    text_prompt_queue: Queue
+    handlers: list[Any]
+    delay_filler_threshold_s: float = 0.0
+    delay_filler_short_clips: list[bytes] = Field(default_factory=list)
+    delay_filler_full_clips: list[bytes] = Field(default_factory=list)
+    http_tts_reserved: bool = False
+
+    session: Optional[SessionState] = None