From fd370b31c0f99ab7a2b8c7d8517d56db941154f9 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Tue, 28 Jul 2026 17:15:29 -0400 Subject: [PATCH 1/9] feat: port data streams to the livekit-ffi rust implementation Replaces the hand-rolled Python data streams (client-side header/chunk/trailer framing over SendStreamHeader/Chunk/TrailerRequest) with the FFI's native implementation (data streams v2: single-packet streams + deflate): - Writers open handle-based FFI streams (text/byte_stream_open/write/close); chunking, stream ids, and compression now happen in rust - send_text/send_file are one-shot FFI calls and gain a compress option - Incoming streams arrive via text/byte_stream_opened room events; readers consume per-handle reader events after a read_incremental request, and unhandled topics dispose the owned reader handle - New public API: StreamError (raised when a stream terminates abnormally, previously indistinguishable from clean EOS) and RoomOptions.data_stream/DataStreamOptions.max_payload_byte_length - Drops the use_legacy_client_implementation override, so data streams v2 support is advertised to other clients The consumer-facing API is unchanged. Behavior changes: read errors now raise StreamError, stream ids are rust-generated unless supplied, sender_identity is honored consistently (was ignored for chunks/trailers), and trailers of targeted streams are no longer broadcast room-wide. --- livekit-rtc/livekit/rtc/__init__.py | 4 + livekit-rtc/livekit/rtc/data_stream.py | 449 ++++++++++++++--------- livekit-rtc/livekit/rtc/participant.py | 87 +++-- livekit-rtc/livekit/rtc/room.py | 151 ++++---- tests/rtc/test_e2e_data_streams.py | 487 +++++++++++++++++++++++++ 5 files changed, 900 insertions(+), 278 deletions(-) create mode 100644 tests/rtc/test_e2e_data_streams.py diff --git a/livekit-rtc/livekit/rtc/__init__.py b/livekit-rtc/livekit/rtc/__init__.py index fabc1a5f..3ed48bcb 100644 --- a/livekit-rtc/livekit/rtc/__init__.py +++ b/livekit-rtc/livekit/rtc/__init__.py @@ -66,6 +66,7 @@ from .room import ( ConnectError, DataPacket, + DataStreamOptions, Room, RoomOptions, RtcConfiguration, @@ -123,6 +124,7 @@ TextStreamWriter, ByteStreamWriter, ByteStreamReader, + StreamError, ) from .data_track import ( LocalDataTrack, @@ -172,6 +174,7 @@ "ConnectError", "Room", "RoomOptions", + "DataStreamOptions", "RtcConfiguration", "SimulateScenarioKind", "SipDTMF", @@ -217,6 +220,7 @@ "TextStreamWriter", "ByteStreamReader", "ByteStreamWriter", + "StreamError", "AudioProcessingModule", "PlatformAudio", "PlatformAudioSource", diff --git a/livekit-rtc/livekit/rtc/data_stream.py b/livekit-rtc/livekit/rtc/data_stream.py index bf62a217..50177646 100644 --- a/livekit-rtc/livekit/rtc/data_stream.py +++ b/livekit-rtc/livekit/rtc/data_stream.py @@ -15,16 +15,13 @@ from __future__ import annotations import asyncio -import uuid import datetime from collections.abc import Callable from dataclasses import dataclass from typing import AsyncIterator, Optional, Dict, List -from ._proto.room_pb2 import DataStream as proto_DataStream +from ._proto import data_stream_pb2 as proto_data_stream from ._proto import ffi_pb2 as proto_ffi -from ._proto import room_pb2 as proto_room from ._ffi_client import FfiClient -from ._utils import split_utf8 from typing import TYPE_CHECKING @@ -32,6 +29,19 @@ from .participant import LocalParticipant STREAM_CHUNK_SIZE = 15_000 +"""Deprecated: chunking now happens inside the FFI; kept for compatibility.""" + +_DISCONNECT_ERROR = "Disconnected while receiving" + + +class StreamError(ConnectionError): + """Raised when a data stream operation fails or an incoming stream + terminates abnormally (e.g. aborted by the sender or exceeding the + room's maximum payload length).""" + + def __init__(self, description: str) -> None: + super().__init__(description) + self.description = description @dataclass @@ -49,40 +59,87 @@ class TextStreamInfo(BaseStreamInfo): attachments: List[str] +@dataclass +class ByteStreamInfo(BaseStreamInfo): + name: str + + +def _text_stream_info_from_proto(info: proto_data_stream.TextStreamInfo) -> TextStreamInfo: + return TextStreamInfo( + stream_id=info.stream_id, + mime_type=info.mime_type, + topic=info.topic, + timestamp=info.timestamp, + size=info.total_length if info.HasField("total_length") else None, + attributes=dict(info.attributes), + attachments=list(info.attached_stream_ids), + ) + + +def _byte_stream_info_from_proto(info: proto_data_stream.ByteStreamInfo) -> ByteStreamInfo: + return ByteStreamInfo( + stream_id=info.stream_id, + mime_type=info.mime_type, + topic=info.topic, + timestamp=info.timestamp, + size=info.total_length if info.HasField("total_length") else None, + attributes=dict(info.attributes), + name=info.name, + ) + + class TextStreamReader: def __init__( self, - header: proto_DataStream.Header, + owned_info: proto_data_stream.OwnedTextStreamReader, + *, + on_close: Optional[Callable[[], None]] = None, ) -> None: - self._header = header - self._info = TextStreamInfo( - stream_id=header.stream_id, - mime_type=header.mime_type, - topic=header.topic, - timestamp=header.timestamp, - size=header.total_length, - attributes=dict(header.attributes), - attachments=list(header.text_header.attached_stream_ids), + self._info = _text_stream_info_from_proto(owned_info.info) + # the read_incremental request below consumes the FFI handle, so it + # must be kept as a raw id and never wrapped in an FfiHandle + handle_id = owned_info.handle.id + self._reader_handle = handle_id + self._on_close = on_close + self._closed = False + self._error: Optional[StreamError] = None + # subscribe before read_incremental so no reader event can be missed + self._queue = FfiClient.instance.queue.subscribe( + filter_fn=lambda e: ( + e.WhichOneof("message") == "text_stream_reader_event" + and e.text_stream_reader_event.reader_handle == handle_id + ), ) - self._queue: asyncio.Queue[proto_DataStream.Chunk | None] = asyncio.Queue() - - async def _on_chunk_update(self, chunk: proto_DataStream.Chunk) -> None: - await self._queue.put(chunk) - - async def _on_stream_close(self, trailer: proto_DataStream.Trailer) -> None: - self.info.attributes = self.info.attributes or {} - self.info.attributes.update(trailer.attributes) - await self._queue.put(None) + req = proto_ffi.FfiRequest() + req.text_read_incremental.reader_handle = handle_id + FfiClient.instance.request(req) def __aiter__(self) -> AsyncIterator[str]: return self async def __anext__(self) -> str: - item = await self._queue.get() - if item is None: - raise StopAsyncIteration - decodedStr = item.content.decode() - return decodedStr + while True: + if self._closed: + if self._error is not None: + raise self._error + raise StopAsyncIteration + event: proto_ffi.FfiEvent = await self._queue.get() + stream_event = event.text_stream_reader_event + detail = stream_event.WhichOneof("detail") + if detail == "chunk_received": + if not stream_event.chunk_received.content: + continue + return stream_event.chunk_received.content + elif detail == "eos": + eos = stream_event.eos + self._info.attributes = self._info.attributes or {} + self._info.attributes.update(eos.attributes) + if eos.HasField("error"): + self._error = StreamError(eos.error.description) + self._close() + if self._error is not None: + raise self._error + raise StopAsyncIteration @property def info(self) -> TextStreamInfo: @@ -94,48 +151,97 @@ async def read_all(self) -> str: final_string += chunk return final_string + def _close(self) -> None: + if not self._closed: + self._closed = True + FfiClient.instance.queue.unsubscribe(self._queue) + if self._on_close is not None: + self._on_close() -@dataclass -class ByteStreamInfo(BaseStreamInfo): - name: str + def _signal_disconnect(self) -> None: + """Injects a synthetic EOS-with-error event so pending reads wake up + and raise StreamError when the room disconnects mid-stream.""" + if self._closed: + return + event = proto_ffi.FfiEvent() + event.text_stream_reader_event.reader_handle = self._reader_handle + event.text_stream_reader_event.eos.error.description = _DISCONNECT_ERROR + self._queue.put_nowait(event) class ByteStreamReader: - def __init__(self, header: proto_DataStream.Header, capacity: int = 0) -> None: - self._header = header - self._info = ByteStreamInfo( - stream_id=header.stream_id, - mime_type=header.mime_type, - topic=header.topic, - timestamp=header.timestamp, - size=header.total_length, - attributes=dict(header.attributes), - name=header.byte_header.name, + def __init__( + self, + owned_info: proto_data_stream.OwnedByteStreamReader, + capacity: int = 0, + *, + on_close: Optional[Callable[[], None]] = None, + ) -> None: + # capacity is ignored: chunk delivery is push-based from the FFI + self._info = _byte_stream_info_from_proto(owned_info.info) + handle_id = owned_info.handle.id + self._reader_handle = handle_id + self._on_close = on_close + self._closed = False + self._error: Optional[StreamError] = None + self._queue = FfiClient.instance.queue.subscribe( + filter_fn=lambda e: ( + e.WhichOneof("message") == "byte_stream_reader_event" + and e.byte_stream_reader_event.reader_handle == handle_id + ), ) - self._queue: asyncio.Queue[proto_DataStream.Chunk | None] = asyncio.Queue(capacity) - - async def _on_chunk_update(self, chunk: proto_DataStream.Chunk) -> None: - await self._queue.put(chunk) - - async def _on_stream_close(self, trailer: proto_DataStream.Trailer) -> None: - self.info.attributes = self.info.attributes or {} - self.info.attributes.update(trailer.attributes) - await self._queue.put(None) + req = proto_ffi.FfiRequest() + req.byte_read_incremental.reader_handle = handle_id + FfiClient.instance.request(req) def __aiter__(self) -> AsyncIterator[bytes]: return self async def __anext__(self) -> bytes: - item = await self._queue.get() - if item is None: - raise StopAsyncIteration - - return item.content + while True: + if self._closed: + if self._error is not None: + raise self._error + raise StopAsyncIteration + event: proto_ffi.FfiEvent = await self._queue.get() + stream_event = event.byte_stream_reader_event + detail = stream_event.WhichOneof("detail") + if detail == "chunk_received": + if not stream_event.chunk_received.content: + continue + return stream_event.chunk_received.content + elif detail == "eos": + eos = stream_event.eos + self._info.attributes = self._info.attributes or {} + self._info.attributes.update(eos.attributes) + if eos.HasField("error"): + self._error = StreamError(eos.error.description) + self._close() + if self._error is not None: + raise self._error + raise StopAsyncIteration @property def info(self) -> ByteStreamInfo: return self._info + def _close(self) -> None: + if not self._closed: + self._closed = True + FfiClient.instance.queue.unsubscribe(self._queue) + if self._on_close is not None: + self._on_close() + + def _signal_disconnect(self) -> None: + """Injects a synthetic EOS-with-error event so pending reads wake up + and raise StreamError when the room disconnects mid-stream.""" + if self._closed: + return + event = proto_ffi.FfiEvent() + event.byte_stream_reader_event.reader_handle = self._reader_handle + event.byte_stream_reader_event.eos.error.description = _DISCONNECT_ERROR + self._queue.put_nowait(event) + class BaseStreamWriter: def __init__( @@ -150,100 +256,46 @@ def __init__( sender_identity: str | None = None, ): self._local_participant = local_participant - if stream_id is None: - stream_id = str(uuid.uuid4()) - timestamp = int(datetime.datetime.now().timestamp() * 1000) - self._header = proto_DataStream.Header( - stream_id=stream_id, - timestamp=timestamp, - mime_type=mime_type, - topic=topic, - attributes=attributes, - total_length=total_size, - ) - self._next_chunk_index: int = 0 - self._destination_identities = destination_identities + self._total_size = total_size self._sender_identity = sender_identity or self._local_participant.identity + # the writer handle is assigned by the FFI when the stream is opened; + # the close request consumes it, so it is kept as a raw id + self._writer_handle: Optional[int] = None + self._write_lock = asyncio.Lock() self._closed = False - async def _send_header(self) -> None: - req = proto_ffi.FfiRequest( - send_stream_header=proto_room.SendStreamHeaderRequest( - header=self._header, - local_participant_handle=self._local_participant._ffi_handle.handle, - destination_identities=self._destination_identities, - sender_identity=self._sender_identity, - ) - ) - - queue = FfiClient.instance.queue.subscribe() - try: - resp = FfiClient.instance.request(req) - cb: proto_ffi.FfiEvent = await queue.wait_for( - lambda e: e.send_stream_header.async_id == resp.send_stream_header.async_id - ) - finally: - FfiClient.instance.queue.unsubscribe(queue) - - if cb.send_stream_header.error: - raise ConnectionError(cb.send_stream_header.error) - - async def _send_chunk(self, chunk: proto_DataStream.Chunk) -> None: - if self._closed: - raise RuntimeError(f"Cannot send chunk after stream is closed: {chunk}") - req = proto_ffi.FfiRequest( - send_stream_chunk=proto_room.SendStreamChunkRequest( - chunk=chunk, - local_participant_handle=self._local_participant._ffi_handle.handle, - sender_identity=self._local_participant.identity, - destination_identities=self._destination_identities, - ) - ) + def _provisional_timestamp(self) -> int: + return int(datetime.datetime.now().timestamp() * 1000) + async def _wait_for_callback( + self, req: proto_ffi.FfiRequest, callback_field: str, response_field: str + ) -> proto_ffi.FfiEvent: queue = FfiClient.instance.queue.subscribe() try: resp = FfiClient.instance.request(req) + async_id = getattr(resp, response_field).async_id cb: proto_ffi.FfiEvent = await queue.wait_for( - lambda e: e.send_stream_chunk.async_id == resp.send_stream_chunk.async_id + lambda e: getattr(e, callback_field).async_id == async_id ) finally: FfiClient.instance.queue.unsubscribe(queue) - if cb.send_stream_chunk.error: - raise ConnectionError(cb.send_stream_chunk.error) - - async def _send_trailer(self, trailer: proto_DataStream.Trailer) -> None: - req = proto_ffi.FfiRequest( - send_stream_trailer=proto_room.SendStreamTrailerRequest( - trailer=trailer, - local_participant_handle=self._local_participant._ffi_handle.handle, - sender_identity=self._local_participant.identity, - ) - ) - - queue = FfiClient.instance.queue.subscribe() - try: - resp = FfiClient.instance.request(req) - cb: proto_ffi.FfiEvent = await queue.wait_for( - lambda e: e.send_stream_trailer.async_id == resp.send_stream_trailer.async_id - ) - finally: - FfiClient.instance.queue.unsubscribe(queue) - - if cb.send_stream_trailer.error: - raise ConnectionError(cb.send_stream_trailer.error) + if getattr(cb, callback_field).HasField("error"): + raise StreamError(getattr(cb, callback_field).error.description) + return cb async def aclose( self, *, reason: str = "", attributes: Optional[Dict[str, str]] = None ) -> None: if self._closed: raise RuntimeError("Stream already closed") + if self._writer_handle is None: + raise RuntimeError("Stream is not open") self._closed = True - await self._send_trailer( - trailer=proto_DataStream.Trailer( - stream_id=self._header.stream_id, reason=reason, attributes=attributes - ) - ) + await self._send_close(reason=reason, attributes=attributes) + + async def _send_close(self, *, reason: str, attributes: Optional[Dict[str, str]]) -> None: + raise NotImplementedError class TextStreamWriter(BaseStreamWriter): @@ -269,32 +321,59 @@ def __init__( destination_identities=destination_identities, sender_identity=sender_identity, ) - self._header.text_header.operation_type = proto_DataStream.OperationType.CREATE + options = proto_data_stream.StreamTextOptions(topic=topic) + if attributes: + options.attributes.update(attributes) + if destination_identities: + options.destination_identities.extend(destination_identities) + if stream_id is not None: + options.id = stream_id if reply_to_id: - self._header.text_header.reply_to_stream_id = reply_to_id + options.reply_to_stream_id = reply_to_id + options.sender_identity = self._sender_identity + self._options = options + # provisional info, replaced by the FFI-provided info once opened self._info = TextStreamInfo( - stream_id=self._header.stream_id, - mime_type=self._header.mime_type, - topic=self._header.topic, - timestamp=self._header.timestamp, - size=self._header.total_length, - attributes=dict(self._header.attributes), - attachments=list(self._header.text_header.attached_stream_ids), + stream_id=stream_id or "", + mime_type="text/plain", + topic=topic, + timestamp=self._provisional_timestamp(), + size=total_size, + attributes=dict(attributes) if attributes else {}, + attachments=[], ) - self._write_lock = asyncio.Lock() + + async def _send_header(self) -> None: + req = proto_ffi.FfiRequest() + req.text_stream_open.local_participant_handle = self._local_participant._ffi_handle.handle + req.text_stream_open.options.CopyFrom(self._options) + cb = await self._wait_for_callback(req, "text_stream_open", "text_stream_open") + owned = cb.text_stream_open.writer + self._writer_handle = owned.handle.id + self._info = _text_stream_info_from_proto(owned.info) + if self._info.size is None and self._total_size is not None: + # total_size is not transmitted for text streams; keep it local + self._info.size = self._total_size async def write(self, text: str) -> None: async with self._write_lock: - for chunk in split_utf8(text, STREAM_CHUNK_SIZE): - content = chunk - chunk_index = self._next_chunk_index - self._next_chunk_index += 1 - chunk_msg = proto_DataStream.Chunk( - stream_id=self._header.stream_id, - chunk_index=chunk_index, - content=content, - ) - await self._send_chunk(chunk_msg) + if self._closed: + raise RuntimeError("Cannot write after stream is closed") + if self._writer_handle is None: + raise RuntimeError("Stream is not open") + req = proto_ffi.FfiRequest() + req.text_stream_write.writer_handle = self._writer_handle + req.text_stream_write.text = text + await self._wait_for_callback(req, "text_stream_writer_write", "text_stream_write") + + async def _send_close(self, *, reason: str, attributes: Optional[Dict[str, str]]) -> None: + assert self._writer_handle is not None + req = proto_ffi.FfiRequest() + req.text_stream_close.writer_handle = self._writer_handle + req.text_stream_close.reason = reason + if attributes: + req.text_stream_close.attributes.update(attributes) + await self._wait_for_callback(req, "text_stream_writer_close", "text_stream_close") @property def info(self) -> TextStreamInfo: @@ -323,32 +402,60 @@ def __init__( mime_type=mime_type, destination_identities=destination_identities, ) - self._header.byte_header.name = name + options = proto_data_stream.StreamByteOptions( + topic=topic, + name=name, + mime_type=mime_type, + ) + if attributes: + options.attributes.update(attributes) + if destination_identities: + options.destination_identities.extend(destination_identities) + if stream_id is not None: + options.id = stream_id + if total_size is not None: + options.total_length = total_size + options.sender_identity = self._sender_identity + self._options = options + # provisional info, replaced by the FFI-provided info once opened self._info = ByteStreamInfo( - stream_id=self._header.stream_id, - mime_type=self._header.mime_type, - topic=self._header.topic, - timestamp=self._header.timestamp, - size=self._header.total_length, - attributes=dict(self._header.attributes), - name=self._header.byte_header.name, + stream_id=stream_id or "", + mime_type=mime_type, + topic=topic, + timestamp=self._provisional_timestamp(), + size=total_size, + attributes=dict(attributes) if attributes else {}, + name=name, ) - self._write_lock = asyncio.Lock() + + async def _send_header(self) -> None: + req = proto_ffi.FfiRequest() + req.byte_stream_open.local_participant_handle = self._local_participant._ffi_handle.handle + req.byte_stream_open.options.CopyFrom(self._options) + cb = await self._wait_for_callback(req, "byte_stream_open", "byte_stream_open") + owned = cb.byte_stream_open.writer + self._writer_handle = owned.handle.id + self._info = _byte_stream_info_from_proto(owned.info) async def write(self, data: bytes) -> None: async with self._write_lock: - chunked_data = [ - data[i : i + STREAM_CHUNK_SIZE] for i in range(0, len(data), STREAM_CHUNK_SIZE) - ] - - for chunk in chunked_data: - chunk_msg = proto_DataStream.Chunk( - stream_id=self._header.stream_id, - chunk_index=self._next_chunk_index, - content=chunk, - ) - await self._send_chunk(chunk_msg) - self._next_chunk_index += 1 + if self._closed: + raise RuntimeError("Cannot write after stream is closed") + if self._writer_handle is None: + raise RuntimeError("Stream is not open") + req = proto_ffi.FfiRequest() + req.byte_stream_write.writer_handle = self._writer_handle + req.byte_stream_write.bytes = data + await self._wait_for_callback(req, "byte_stream_writer_write", "byte_stream_write") + + async def _send_close(self, *, reason: str, attributes: Optional[Dict[str, str]]) -> None: + assert self._writer_handle is not None + req = proto_ffi.FfiRequest() + req.byte_stream_close.writer_handle = self._writer_handle + req.byte_stream_close.reason = reason + if attributes: + req.byte_stream_close.attributes.update(attributes) + await self._wait_for_callback(req, "byte_stream_writer_close", "byte_stream_close") @property def info(self) -> ByteStreamInfo: diff --git a/livekit-rtc/livekit/rtc/participant.py b/livekit-rtc/livekit/rtc/participant.py index 2a53526e..6e6cd556 100644 --- a/livekit-rtc/livekit/rtc/participant.py +++ b/livekit-rtc/livekit/rtc/participant.py @@ -20,7 +20,6 @@ import enum import os import mimetypes -import aiofiles from typing import List, Union, Callable, Dict, Awaitable, Optional, Mapping, cast, TypeVar from abc import abstractmethod, ABC @@ -58,7 +57,9 @@ TextStreamInfo, ByteStreamWriter, ByteStreamInfo, - STREAM_CHUNK_SIZE, + StreamError, + _text_stream_info_from_proto, + _byte_stream_info_from_proto, ) from .data_track import LocalDataTrack from ._proto import data_track_pb2 as proto_data_track @@ -677,20 +678,36 @@ async def send_text( topic: str = "", attributes: Optional[Dict[str, str]] = None, reply_to_id: str | None = None, + compress: bool | None = None, ) -> TextStreamInfo: - total_size = len(text.encode()) - writer = await self.stream_text( - destination_identities=destination_identities, - topic=topic, - attributes=attributes, - reply_to_id=reply_to_id, - total_size=total_size, - ) + req = proto_ffi.FfiRequest() + req.send_text.local_participant_handle = self._ffi_handle.handle + req.send_text.text = text + options = req.send_text.options + options.topic = topic + if attributes: + options.attributes.update(attributes) + if destination_identities: + options.destination_identities.extend(destination_identities) + if reply_to_id: + options.reply_to_stream_id = reply_to_id + if compress is not None: + options.compress = compress + options.sender_identity = self.identity + + queue = FfiClient.instance.queue.subscribe() + try: + resp = FfiClient.instance.request(req) + cb: proto_ffi.FfiEvent = await queue.wait_for( + lambda e: e.send_text.async_id == resp.send_text.async_id + ) + finally: + FfiClient.instance.queue.unsubscribe(queue) - await writer.write(text) - await writer.aclose() + if cb.send_text.HasField("error"): + raise StreamError(cb.send_text.error.description) - return writer.info + return _text_stream_info_from_proto(cb.send_text.info) async def stream_bytes( self, @@ -730,29 +747,43 @@ async def send_file( destination_identities: Optional[List[str]] = None, attributes: Optional[Dict[str, str]] = None, stream_id: str | None = None, + compress: bool | None = None, ) -> ByteStreamInfo: - file_size = os.path.getsize(file_path) file_name = os.path.basename(file_path) mime_type, _ = mimetypes.guess_type(file_path) if mime_type is None: mime_type = "application/octet-stream" # Fallback MIME type for unknown files - writer: ByteStreamWriter = await self.stream_bytes( - name=file_name, - total_size=file_size, - mime_type=mime_type, - attributes=attributes, - stream_id=stream_id, - destination_identities=destination_identities, - topic=topic, - ) + req = proto_ffi.FfiRequest() + req.send_file.local_participant_handle = self._ffi_handle.handle + req.send_file.file_path = file_path + options = req.send_file.options + options.topic = topic + options.name = file_name + options.mime_type = mime_type + if attributes: + options.attributes.update(attributes) + if destination_identities: + options.destination_identities.extend(destination_identities) + if stream_id is not None: + options.id = stream_id + if compress is not None: + options.compress = compress + options.sender_identity = self.identity + + queue = FfiClient.instance.queue.subscribe() + try: + resp = FfiClient.instance.request(req) + cb: proto_ffi.FfiEvent = await queue.wait_for( + lambda e: e.send_file.async_id == resp.send_file.async_id + ) + finally: + FfiClient.instance.queue.unsubscribe(queue) - async with aiofiles.open(file_path, "rb") as f: - while bytes := await f.read(STREAM_CHUNK_SIZE): - await writer.write(bytes) - await writer.aclose() + if cb.send_file.HasField("error"): + raise StreamError(cb.send_file.error.description) - return writer.info + return _byte_stream_info_from_proto(cb.send_file.info) async def publish_data_track( self, diff --git a/livekit-rtc/livekit/rtc/room.py b/livekit-rtc/livekit/rtc/room.py index 83cb06ee..c288e682 100644 --- a/livekit-rtc/livekit/rtc/room.py +++ b/livekit-rtc/livekit/rtc/room.py @@ -107,6 +107,14 @@ class RtcConfiguration: the SFU.""" +@dataclass +class DataStreamOptions: + max_payload_byte_length: int | None = None + """Maximum decompressed payload size in bytes accepted for a single incoming + data stream; oversized streams terminate with StreamError on the receiver. + When None, uses the FFI default.""" + + @dataclass class RoomOptions: auto_subscribe: bool = True @@ -122,6 +130,8 @@ class RoomOptions: """Timeout in seconds for each signal connection attempt. When None, uses the default (5s).""" single_peer_connection: bool | None = None """Use a single peer connection for both publish and subscribe. When None, uses the default (false).""" + data_stream: DataStreamOptions | None = None + """Options for incoming data streams. When None, uses the FFI defaults.""" @dataclass @@ -174,15 +184,15 @@ def __init__( self._room_queue = BroadcastQueue[proto_ffi.FfiEvent]() self._info = proto_room.RoomInfo() self._rpc_invocation_tasks: set[asyncio.Task] = set() - self._data_stream_tasks: set[asyncio.Task] = set() self._remote_participants: Dict[str, RemoteParticipant] = {} self._connection_state = ConnectionState.CONN_DISCONNECTED self._first_sid_future = asyncio.Future[str]() self._local_participant: LocalParticipant | None = None - self._text_stream_readers: Dict[str, TextStreamReader] = {} - self._byte_stream_readers: Dict[str, ByteStreamReader] = {} + # active incoming stream readers, keyed by FFI reader handle id + self._text_stream_readers: Dict[int, TextStreamReader] = {} + self._byte_stream_readers: Dict[int, ByteStreamReader] = {} self._text_stream_handlers: Dict[str, TextStreamHandler] = {} self._byte_stream_handlers: Dict[str, ByteStreamHandler] = {} @@ -467,10 +477,13 @@ def on_participant_connected(participant): req.connect.options.auto_subscribe = options.auto_subscribe req.connect.options.dynacast = options.dynacast - # The Python SDK still implements data streams in Python on top of raw FFI - # packets, so always advertise only legacy (v1) data stream support to other - # clients. - req.connect.options.data_stream.use_legacy_client_implementation = True + if ( + options.data_stream is not None + and options.data_stream.max_payload_byte_length is not None + ): + req.connect.options.data_stream.max_payload_byte_length = ( + options.data_stream.max_payload_byte_length + ) if options.connect_timeout is not None: req.connect.options.connect_timeout_ms = int(options.connect_timeout * 1000) @@ -672,7 +685,7 @@ async def disconnect( return await self._drain_rpc_invocation_tasks() - await self._drain_data_stream_tasks() + self._error_stream_readers() req = proto_ffi.FfiRequest() req.disconnect.room_handle = self._ffi_handle.handle # type: ignore @@ -721,7 +734,7 @@ async def _listen_task(self) -> None: # Clean up any pending RPC invocation tasks await self._drain_rpc_invocation_tasks() - await self._drain_data_stream_tasks() + self._error_stream_readers() def _on_rpc_method_invocation(self, rpc_invocation: RpcMethodInvocationEvent) -> None: if self._local_participant is None: @@ -1109,23 +1122,10 @@ def _on_room_event(self, event: proto_room.RoomEvent) -> None: self.emit("reconnecting") elif which == "reconnected": self.emit("reconnected") - elif which == "stream_header_received": - self._handle_stream_header( - event.stream_header_received.header, - event.stream_header_received.participant_identity, - ) - elif which == "stream_chunk_received": - task = asyncio.create_task(self._handle_stream_chunk(event.stream_chunk_received.chunk)) - self._data_stream_tasks.add(task) - task.add_done_callback(self._data_stream_tasks.discard) - - elif which == "stream_trailer_received": - task = asyncio.create_task( - self._handle_stream_trailer(event.stream_trailer_received.trailer) - ) - self._data_stream_tasks.add(task) - task.add_done_callback(self._data_stream_tasks.discard) - + elif which == "text_stream_opened": + self._handle_text_stream_opened(event.text_stream_opened) + elif which == "byte_stream_opened": + self._handle_byte_stream_opened(event.byte_stream_opened) elif which == "room_updated": self._info = event.room_updated self._resolve_first_sid(self._info.sid) @@ -1153,57 +1153,56 @@ def _on_room_event(self, event: proto_room.RoomEvent) -> None: elif which == "data_track_unpublished": self.emit("data_track_unpublished", event.data_track_unpublished.sid) - def _handle_stream_header( - self, header: proto_room.DataStream.Header, participant_identity: str - ) -> None: - stream_type = header.WhichOneof("content_header") - if stream_type == "text_header": - text_stream_handler = self._text_stream_handlers.get(header.topic) - if text_stream_handler is None: - logging.info( - "ignoring text stream with topic '%s', no callback attached", - header.topic, - ) - return - - text_reader = TextStreamReader(header) - self._text_stream_readers[header.stream_id] = text_reader - text_stream_handler(text_reader, participant_identity) - elif stream_type == "byte_header": - byte_stream_handler = self._byte_stream_handlers.get(header.topic) - if byte_stream_handler is None: - logging.info( - "ignoring byte stream with topic '%s', no callback attached", - header.topic, - ) - return - - byte_reader = ByteStreamReader(header) - self._byte_stream_readers[header.stream_id] = byte_reader - byte_stream_handler(byte_reader, participant_identity) - else: - logging.warning("received unknown header type, %s", stream_type) - pass + def _handle_text_stream_opened(self, opened: proto_room.TextStreamOpened) -> None: + topic = opened.reader.info.topic + text_stream_handler = self._text_stream_handlers.get(topic) + handle_id = opened.reader.handle.id + if text_stream_handler is None: + logging.info( + "ignoring text stream with topic '%s', no callback attached", + topic, + ) + # the reader is never consumed by read_incremental, so the owned + # handle (and its buffered chunks) must be disposed here + FfiHandle(handle_id).dispose() + return - async def _handle_stream_chunk(self, chunk: proto_room.DataStream.Chunk) -> None: - text_reader = self._text_stream_readers.get(chunk.stream_id) - file_reader = self._byte_stream_readers.get(chunk.stream_id) + def on_close() -> None: + self._text_stream_readers.pop(handle_id, None) + + text_reader = TextStreamReader(opened.reader, on_close=on_close) + self._text_stream_readers[handle_id] = text_reader + text_stream_handler(text_reader, opened.participant_identity) + + def _handle_byte_stream_opened(self, opened: proto_room.ByteStreamOpened) -> None: + topic = opened.reader.info.topic + byte_stream_handler = self._byte_stream_handlers.get(topic) + handle_id = opened.reader.handle.id + if byte_stream_handler is None: + logging.info( + "ignoring byte stream with topic '%s', no callback attached", + topic, + ) + # the reader is never consumed by read_incremental, so the owned + # handle (and its buffered chunks) must be disposed here + FfiHandle(handle_id).dispose() + return - if text_reader: - await text_reader._on_chunk_update(chunk) - elif file_reader: - await file_reader._on_chunk_update(chunk) + def on_close() -> None: + self._byte_stream_readers.pop(handle_id, None) - async def _handle_stream_trailer(self, trailer: proto_room.DataStream.Trailer) -> None: - text_reader = self._text_stream_readers.get(trailer.stream_id) - file_reader = self._byte_stream_readers.get(trailer.stream_id) + byte_reader = ByteStreamReader(opened.reader, on_close=on_close) + self._byte_stream_readers[handle_id] = byte_reader + byte_stream_handler(byte_reader, opened.participant_identity) - if text_reader: - await text_reader._on_stream_close(trailer) - self._text_stream_readers.pop(trailer.stream_id) - elif file_reader: - await file_reader._on_stream_close(trailer) - self._byte_stream_readers.pop(trailer.stream_id) + def _error_stream_readers(self) -> None: + """Wakes up any pending stream reads with a StreamError on disconnect.""" + for text_reader in self._text_stream_readers.values(): + text_reader._signal_disconnect() + for byte_reader in self._byte_stream_readers.values(): + byte_reader._signal_disconnect() + self._text_stream_readers.clear() + self._byte_stream_readers.clear() async def _drain_rpc_invocation_tasks(self) -> None: if self._rpc_invocation_tasks: @@ -1211,12 +1210,6 @@ async def _drain_rpc_invocation_tasks(self) -> None: task.cancel() await asyncio.gather(*self._rpc_invocation_tasks, return_exceptions=True) - async def _drain_data_stream_tasks(self) -> None: - if self._data_stream_tasks: - for task in self._data_stream_tasks: - task.cancel() - await asyncio.gather(*self._data_stream_tasks, return_exceptions=True) - def _retrieve_remote_participant(self, identity: str) -> Optional[RemoteParticipant]: """Retrieve a remote participant by identity""" return self._remote_participants.get(identity, None) diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py new file mode 100644 index 00000000..7c6f57fb --- /dev/null +++ b/tests/rtc/test_e2e_data_streams.py @@ -0,0 +1,487 @@ +""" +End-to-end tests for data streams (text/byte streams over the FFI-backed +rust implementation). Mirrors the node SDK's e2e_data_streams tests, plus +coverage of receiver-side payload caps, abnormal termination, and trailer +attributes. + +Requirements: +- LIVEKIT_URL: LiveKit server URL +- LIVEKIT_API_KEY: API key for authentication +- LIVEKIT_API_SECRET: API secret for authentication + +Tests will be skipped if these environment variables are not set. + +Usage: + pytest test_e2e_data_streams.py -v +""" + +import asyncio +import os +import time +import uuid +from typing import Any, List, Optional, Tuple + +import pytest + +from livekit import api, rtc + + +def skip_if_no_credentials() -> Any: + required_vars = ["LIVEKIT_URL", "LIVEKIT_API_KEY", "LIVEKIT_API_SECRET"] + missing = [var for var in required_vars if not os.getenv(var)] + return pytest.mark.skipif( + bool(missing), reason=f"Missing environment variables: {', '.join(missing)}" + ) + + +def create_token(identity: str, room_name: str) -> str: + return ( + api.AccessToken() + .with_identity(identity) + .with_name(identity) + .with_grants(api.VideoGrants(room_join=True, room=room_name)) + .to_jwt() + ) + + +def unique_room_name(base: str) -> str: + return f"{base}-{uuid.uuid4().hex[:8]}" + + +def pseudo_random_text(length: int, seed: int = 0x5EED) -> str: + """Deterministic pseudo-random lowercase text (mulberry32 PRNG). + + Random lowercase carries ~4.7 bits of entropy per 8-bit byte, so deflate + compresses it well under its raw size — exercising the chunked-compressed + wire path. + """ + a = seed & 0xFFFFFFFF + chars: List[str] = [] + for _ in range(length): + a = (a + 0x6D2B79F5) & 0xFFFFFFFF + t = a + t = (t ^ (t >> 15)) * (t | 1) & 0xFFFFFFFF + t = (t + ((t ^ (t >> 7)) * (t | 61) & 0xFFFFFFFF)) & 0xFFFFFFFF + t = t ^ (t >> 14) + chars.append(chr(97 + t % 26)) + return "".join(chars) + + +async def connect_rooms( + room_name: str, + *, + receiver_options: Optional[rtc.RoomOptions] = None, +) -> Tuple[rtc.Room, rtc.Room]: + """Connects a (receiver, sender) room pair and waits for mutual visibility.""" + url = os.getenv("LIVEKIT_URL") + assert url is not None + + receiver = rtc.Room() + sender = rtc.Room() + await receiver.connect( + url, create_token("receiver", room_name), receiver_options or rtc.RoomOptions() + ) + await sender.connect(url, create_token("sender", room_name)) + + deadline = asyncio.get_event_loop().time() + 5.0 + while asyncio.get_event_loop().time() < deadline: + if len(receiver.remote_participants) == 1 and len(sender.remote_participants) == 1: + break + await asyncio.sleep(0.05) + else: + raise AssertionError("participants did not become visible to each other") + + return receiver, sender + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_send_text_small() -> None: + """Small text via send_text: info assertions + round trip.""" + receiver, sender = await connect_rooms(unique_room_name("ds-text-small")) + try: + text_to_send = "some-text" + received: asyncio.Future[str] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, participant_identity: str) -> None: + assert participant_identity == "sender" + + async def read() -> None: + received.set_result(await reader.read_all()) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("some-topic", handler) + + info = await sender.local_participant.send_text(text_to_send, topic="some-topic") + assert info.stream_id + assert abs(info.timestamp - time.time() * 1000) <= 2_000 + assert info.size == len(text_to_send.encode()) + assert info.mime_type == "text/plain" + assert info.topic == "some-topic" + + assert await asyncio.wait_for(received, timeout=5.0) == text_to_send + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_stream_bytes_small() -> None: + """Small byte stream via the incremental writer.""" + receiver, sender = await connect_rooms(unique_room_name("ds-bytes-small")) + try: + bytes_to_send = b"\xfa" * 16 + received: asyncio.Future[bytes] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.ByteStreamReader, participant_identity: str) -> None: + assert participant_identity == "sender" + + async def read() -> None: + chunks = [chunk async for chunk in reader] + received.set_result(b"".join(chunks)) + + asyncio.create_task(read()) + + receiver.register_byte_stream_handler("some-topic", handler) + + writer = await sender.local_participant.stream_bytes( + "test-bytes", topic="some-topic", total_size=len(bytes_to_send) + ) + assert writer.info.stream_id + assert abs(writer.info.timestamp - time.time() * 1000) <= 2_000 + assert writer.info.mime_type == "application/octet-stream" + assert writer.info.topic == "some-topic" + assert writer.info.name == "test-bytes" + await writer.write(bytes_to_send) + await writer.aclose() + + assert await asyncio.wait_for(received, timeout=5.0) == bytes_to_send + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_send_text_large_compressible() -> None: + """~50KB of pseudo-random lowercase: too big to inline and compresses well, + exercising the chunked + deflate-raw wire path (rust compresses, rust + decompresses, python reads).""" + receiver, sender = await connect_rooms(unique_room_name("ds-text-large")) + try: + text = pseudo_random_text(50_000) + received: asyncio.Future[str] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + async def read() -> None: + received.set_result(await reader.read_all()) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("large-text", handler) + + info = await sender.local_participant.send_text(text, topic="large-text") + assert info.size == len(text.encode()) + + assert await asyncio.wait_for(received, timeout=10.0) == text + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_send_text_compress_false() -> None: + """compress=False round-trips the payload uncompressed.""" + receiver, sender = await connect_rooms(unique_room_name("ds-text-nocompress")) + try: + text = pseudo_random_text(50_000) + received: asyncio.Future[str] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + async def read() -> None: + received.set_result(await reader.read_all()) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("nocompress-text", handler) + + await sender.local_participant.send_text(text, topic="nocompress-text", compress=False) + + assert await asyncio.wait_for(received, timeout=10.0) == text + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_stream_bytes_large_incompressible() -> None: + """Large uniform-random payload spanning many chunks (uncompressed path).""" + receiver, sender = await connect_rooms(unique_room_name("ds-bytes-random")) + try: + payload = os.urandom(1_000_000) + received: asyncio.Future[bytes] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.ByteStreamReader, _identity: str) -> None: + async def read() -> None: + chunks = [chunk async for chunk in reader] + received.set_result(b"".join(chunks)) + + asyncio.create_task(read()) + + receiver.register_byte_stream_handler("random-bytes", handler) + + writer = await sender.local_participant.stream_bytes( + "random-bytes", topic="random-bytes", total_size=len(payload) + ) + write_chunk_size = 64 * 1024 + for offset in range(0, len(payload), write_chunk_size): + await writer.write(payload[offset : offset + write_chunk_size]) + await writer.aclose() + + result = await asyncio.wait_for(received, timeout=20.0) + assert len(result) == len(payload) + assert result == payload + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_stream_bytes_large_patterned() -> None: + """50KB patterned payload via the byte-stream writer.""" + receiver, sender = await connect_rooms(unique_room_name("ds-bytes-patterned")) + try: + payload = bytes(i % 251 for i in range(50_000)) + received: asyncio.Future[bytes] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.ByteStreamReader, _identity: str) -> None: + async def read() -> None: + chunks = [chunk async for chunk in reader] + received.set_result(b"".join(chunks)) + + asyncio.create_task(read()) + + receiver.register_byte_stream_handler("patterned-bytes", handler) + + writer = await sender.local_participant.stream_bytes( + "patterned", topic="patterned-bytes", total_size=len(payload) + ) + await writer.write(payload) + await writer.aclose() + + assert await asyncio.wait_for(received, timeout=10.0) == payload + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_stream_text_incremental_unicode() -> None: + """Multiple writes with multi-byte characters: UTF-8-aware chunking must + not split codepoints.""" + receiver, sender = await connect_rooms(unique_room_name("ds-text-unicode")) + try: + pieces = ["héllo wörld — ", "日本語のテキスト、", "🌍🚀 emoji tail"] + expected = "".join(pieces) + received: asyncio.Future[str] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + async def read() -> None: + received.set_result(await reader.read_all()) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("unicode-text", handler) + + writer = await sender.local_participant.stream_text(topic="unicode-text") + for piece in pieces: + await writer.write(piece) + await writer.aclose() + + assert await asyncio.wait_for(received, timeout=5.0) == expected + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_receiver_rejects_oversized_payload() -> None: + """A receiver with a small max_payload_byte_length must reject an + oversized stream with StreamError rather than delivering truncated data.""" + receiver, sender = await connect_rooms( + unique_room_name("ds-payload-cap"), + receiver_options=rtc.RoomOptions( + data_stream=rtc.DataStreamOptions(max_payload_byte_length=1_000) + ), + ) + try: + text = pseudo_random_text(50_000) + result: asyncio.Future[BaseException] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + async def read() -> None: + try: + data = await reader.read_all() + result.set_exception( + AssertionError(f"expected StreamError, read {len(data)} chars") + ) + except rtc.StreamError as e: + result.set_result(e) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("capped-topic", handler) + + await sender.local_participant.send_text(text, topic="capped-topic") + + error = await asyncio.wait_for(result, timeout=10.0) + assert "maximum size" in str(error) + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_abnormal_close_raises_on_receiver() -> None: + """Closing a stream with a reason mid-transfer must raise StreamError on + the receiving reader.""" + receiver, sender = await connect_rooms(unique_room_name("ds-abnormal-close")) + try: + result: asyncio.Future[BaseException] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + async def read() -> None: + try: + data = await reader.read_all() + result.set_exception( + AssertionError(f"expected StreamError, read {len(data)} chars") + ) + except rtc.StreamError as e: + result.set_result(e) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("abort-topic", handler) + + writer = await sender.local_participant.stream_text(topic="abort-topic") + await writer.write("partial data") + await writer.aclose(reason="cancelled by test") + + error = await asyncio.wait_for(result, timeout=5.0) + assert "cancelled by test" in str(error) + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_trailer_attributes_merged_into_info() -> None: + """Attributes passed to aclose() must appear in the receiving reader's + info after a clean EOS.""" + receiver, sender = await connect_rooms(unique_room_name("ds-trailer-attrs")) + try: + received: asyncio.Future[Tuple[str, rtc.TextStreamInfo]] = ( + asyncio.get_event_loop().create_future() + ) + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + async def read() -> None: + text = await reader.read_all() + received.set_result((text, reader.info)) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("attrs-topic", handler) + + writer = await sender.local_participant.stream_text( + topic="attrs-topic", attributes={"initial": "yes"} + ) + await writer.write("hello") + await writer.aclose(attributes={"result": "ok", "count": "1"}) + + text, info = await asyncio.wait_for(received, timeout=5.0) + assert text == "hello" + assert info.attributes is not None + assert info.attributes.get("initial") == "yes" + assert info.attributes.get("result") == "ok" + assert info.attributes.get("count") == "1" + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_receiver_disconnect_mid_stream() -> None: + """Disconnecting the receiving room mid-stream must surface StreamError to + the reader instead of hanging.""" + receiver, sender = await connect_rooms(unique_room_name("ds-disconnect")) + try: + result: asyncio.Future[BaseException] = asyncio.get_event_loop().create_future() + handler_called = asyncio.Event() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + handler_called.set() + + async def read() -> None: + try: + await reader.read_all() + result.set_exception(AssertionError("expected StreamError, got clean EOF")) + except rtc.StreamError as e: + result.set_result(e) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("disconnect-topic", handler) + + writer = await sender.local_participant.stream_text(topic="disconnect-topic") + await writer.write("partial data") + await asyncio.wait_for(handler_called.wait(), timeout=5.0) + + await receiver.disconnect() + + error = await asyncio.wait_for(result, timeout=5.0) + assert isinstance(error, rtc.StreamError) + + await writer.aclose() + finally: + await sender.disconnect() + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_send_file_round_trip(tmp_path: Any) -> None: + """send_file hands the path to the FFI, which reads and streams the file; + the receiver gets the payload plus the derived name, size, and mime type.""" + receiver, sender = await connect_rooms(unique_room_name("ds-file")) + try: + payload = os.urandom(64_000) + file_path = tmp_path / "test-payload.bin" + file_path.write_bytes(payload) + + received: asyncio.Future[bytes] = asyncio.get_event_loop().create_future() + received_info: asyncio.Future[rtc.ByteStreamInfo] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.ByteStreamReader, participant_identity: str) -> None: + assert participant_identity == "sender" + received_info.set_result(reader.info) + + async def read() -> None: + chunks = [chunk async for chunk in reader] + received.set_result(b"".join(chunks)) + + asyncio.create_task(read()) + + receiver.register_byte_stream_handler("file-topic", handler) + + info = await sender.local_participant.send_file(str(file_path), topic="file-topic") + assert info.stream_id + assert info.name == "test-payload.bin" + assert info.size == len(payload) + + assert await asyncio.wait_for(received, timeout=10.0) == payload + reader_info = await asyncio.wait_for(received_info, timeout=5.0) + assert reader_info.name == "test-payload.bin" + assert reader_info.size == len(payload) + assert reader_info.mime_type == "application/octet-stream" + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) From da90a5bbde03652f880e98e45216c47e19063913 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Fri, 31 Jul 2026 11:06:00 -0400 Subject: [PATCH 2/9] fix: clean up e2e tests implementation Use a less hard to reason about PRNG randomness source --- tests/rtc/test_e2e_data_streams.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py index 7c6f57fb..82f5183b 100644 --- a/tests/rtc/test_e2e_data_streams.py +++ b/tests/rtc/test_e2e_data_streams.py @@ -17,9 +17,11 @@ import asyncio import os +import random +import string import time import uuid -from typing import Any, List, Optional, Tuple +from typing import Any, Optional, Tuple import pytest @@ -49,22 +51,14 @@ def unique_room_name(base: str) -> str: def pseudo_random_text(length: int, seed: int = 0x5EED) -> str: - """Deterministic pseudo-random lowercase text (mulberry32 PRNG). + """Deterministic pseudo-random lowercase text. Random lowercase carries ~4.7 bits of entropy per 8-bit byte, so deflate compresses it well under its raw size — exercising the chunked-compressed wire path. """ - a = seed & 0xFFFFFFFF - chars: List[str] = [] - for _ in range(length): - a = (a + 0x6D2B79F5) & 0xFFFFFFFF - t = a - t = (t ^ (t >> 15)) * (t | 1) & 0xFFFFFFFF - t = (t + ((t ^ (t >> 7)) * (t | 61) & 0xFFFFFFFF)) & 0xFFFFFFFF - t = t ^ (t >> 14) - chars.append(chr(97 + t % 26)) - return "".join(chars) + rng = random.Random(seed) + return "".join(rng.choices(string.ascii_lowercase, k=length)) async def connect_rooms( From d21086ff3601ecabf06873a2ecf3937791da6502 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Fri, 31 Jul 2026 12:28:48 -0400 Subject: [PATCH 3/9] fix: ensure that data streams readers are never orphaned when there's an error reading via ffi --- livekit-rtc/livekit/rtc/data_stream.py | 119 +++++++++++++++++- livekit-rtc/livekit/rtc/room.py | 14 ++- tests/rtc/test_e2e_data_streams.py | 160 +++++++++++++++++++++++++ 3 files changed, 286 insertions(+), 7 deletions(-) diff --git a/livekit-rtc/livekit/rtc/data_stream.py b/livekit-rtc/livekit/rtc/data_stream.py index 50177646..a033c8f7 100644 --- a/livekit-rtc/livekit/rtc/data_stream.py +++ b/livekit-rtc/livekit/rtc/data_stream.py @@ -89,6 +89,25 @@ def _byte_stream_info_from_proto(info: proto_data_stream.ByteStreamInfo) -> Byte class TextStreamReader: + """An incoming text stream. + + Use as an async iterator to receive chunks, or :meth:`read_all` to collect + the whole payload:: + + async for chunk in reader: + process(chunk) + + A reader subscribes to the FFI event queue as soon as it is handed to your + handler, so one that is never iterated to completion — abandoned after a + ``break``, or received by a handler that never reads it — must be closed + with :meth:`close`, or its subscription lives for the rest of the process. + Iterating to the end closes the reader for you. + + If the stream terminates abnormally (aborted by the sender, oversized, or + the room disconnects mid-stream), :class:`StreamError` is raised instead of + a normal ``StopAsyncIteration``. + """ + def __init__( self, owned_info: proto_data_stream.OwnedTextStreamReader, @@ -151,25 +170,80 @@ async def read_all(self) -> str: final_string += chunk return final_string + def _unsubscribe(self) -> None: + """Drops the FFI queue subscription without ending the stream. + + Events already delivered stay in the queue and remain readable; only + the global subscription, and the filter it runs against every FFI + event, goes away. Safe to call repeatedly — unsubscribing a queue that + is no longer registered is a no-op. + + Only the subscription is released: the reader handle was consumed by + the read_incremental request in __init__, so there is nothing left to + dispose (and disposing it would drop a handle the FFI server no longer + owns). + """ + FfiClient.instance.queue.unsubscribe(self._queue) + def _close(self) -> None: if not self._closed: self._closed = True - FfiClient.instance.queue.unsubscribe(self._queue) + self._unsubscribe() if self._on_close is not None: self._on_close() + def close(self) -> None: + """Explicitly close the reader and unsubscribe. + + Call this on a reader you stop consuming before it ends. Closing a + reader that already reached end-of-stream is a no-op, and iterating a + closed reader raises ``StopAsyncIteration`` (or :class:`StreamError` + if the stream had already failed). + """ + self._close() + + async def aclose(self) -> None: + self.close() + def _signal_disconnect(self) -> None: """Injects a synthetic EOS-with-error event so pending reads wake up - and raise StreamError when the room disconnects mid-stream.""" + and raise StreamError when the room disconnects mid-stream. + + Also drops the queue subscription, which would otherwise outlive the + room: no further events can arrive once the room is gone, and the + injected one is already queued. The reader is deliberately left open + so chunks that arrived before the disconnect are still delivered + before the StreamError, as they were before this was unsubscribed + here. + """ if self._closed: return event = proto_ffi.FfiEvent() event.text_stream_reader_event.reader_handle = self._reader_handle event.text_stream_reader_event.eos.error.description = _DISCONNECT_ERROR self._queue.put_nowait(event) + self._unsubscribe() class ByteStreamReader: + """An incoming byte stream. + + Use as an async iterator to receive chunks:: + + async for chunk in reader: + process(chunk) + + A reader subscribes to the FFI event queue as soon as it is handed to your + handler, so one that is never iterated to completion — abandoned after a + ``break``, or received by a handler that never reads it — must be closed + with :meth:`close`, or its subscription lives for the rest of the process. + Iterating to the end closes the reader for you. + + If the stream terminates abnormally (aborted by the sender, oversized, or + the room disconnects mid-stream), :class:`StreamError` is raised instead of + a normal ``StopAsyncIteration``. + """ + def __init__( self, owned_info: proto_data_stream.OwnedByteStreamReader, @@ -225,22 +299,59 @@ async def __anext__(self) -> bytes: def info(self) -> ByteStreamInfo: return self._info + def _unsubscribe(self) -> None: + """Drops the FFI queue subscription without ending the stream. + + Events already delivered stay in the queue and remain readable; only + the global subscription, and the filter it runs against every FFI + event, goes away. Safe to call repeatedly — unsubscribing a queue that + is no longer registered is a no-op. + + Only the subscription is released: the reader handle was consumed by + the read_incremental request in __init__, so there is nothing left to + dispose (and disposing it would drop a handle the FFI server no longer + owns). + """ + FfiClient.instance.queue.unsubscribe(self._queue) + def _close(self) -> None: if not self._closed: self._closed = True - FfiClient.instance.queue.unsubscribe(self._queue) + self._unsubscribe() if self._on_close is not None: self._on_close() + def close(self) -> None: + """Explicitly close the reader and unsubscribe. + + Call this on a reader you stop consuming before it ends. Closing a + reader that already reached end-of-stream is a no-op, and iterating a + closed reader raises ``StopAsyncIteration`` (or :class:`StreamError` + if the stream had already failed). + """ + self._close() + + async def aclose(self) -> None: + self.close() + def _signal_disconnect(self) -> None: """Injects a synthetic EOS-with-error event so pending reads wake up - and raise StreamError when the room disconnects mid-stream.""" + and raise StreamError when the room disconnects mid-stream. + + Also drops the queue subscription, which would otherwise outlive the + room: no further events can arrive once the room is gone, and the + injected one is already queued. The reader is deliberately left open + so chunks that arrived before the disconnect are still delivered + before the StreamError, as they were before this was unsubscribed + here. + """ if self._closed: return event = proto_ffi.FfiEvent() event.byte_stream_reader_event.reader_handle = self._reader_handle event.byte_stream_reader_event.eos.error.description = _DISCONNECT_ERROR self._queue.put_nowait(event) + self._unsubscribe() class BaseStreamWriter: diff --git a/livekit-rtc/livekit/rtc/room.py b/livekit-rtc/livekit/rtc/room.py index c288e682..c67b0950 100644 --- a/livekit-rtc/livekit/rtc/room.py +++ b/livekit-rtc/livekit/rtc/room.py @@ -1196,10 +1196,18 @@ def on_close() -> None: byte_stream_handler(byte_reader, opened.participant_identity) def _error_stream_readers(self) -> None: - """Wakes up any pending stream reads with a StreamError on disconnect.""" - for text_reader in self._text_stream_readers.values(): + """Wakes up any pending stream reads with a StreamError on disconnect. + + This also drops each reader's FFI queue subscription, which would + otherwise outlive the room: a reader unsubscribes itself only once a + read consumes its end-of-stream event, so one the application never + finished reading would stay subscribed for the life of the process. + The readers stay open, so a read can still drain whatever arrived + before the disconnect and then raise StreamError. + """ + for text_reader in list(self._text_stream_readers.values()): text_reader._signal_disconnect() - for byte_reader in self._byte_stream_readers.values(): + for byte_reader in list(self._byte_stream_readers.values()): byte_reader._signal_disconnect() self._text_stream_readers.clear() self._byte_stream_readers.clear() diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py index 82f5183b..480e525f 100644 --- a/tests/rtc/test_e2e_data_streams.py +++ b/tests/rtc/test_e2e_data_streams.py @@ -26,6 +26,7 @@ import pytest from livekit import api, rtc +from livekit.rtc._ffi_client import FfiClient def skip_if_no_credentials() -> Any: @@ -61,6 +62,26 @@ def pseudo_random_text(length: int, seed: int = 0x5EED) -> str: return "".join(rng.choices(string.ascii_lowercase, k=length)) +def reader_is_subscribed(reader: Any) -> bool: + """True while the reader's filtered subscriber is still on the FFI queue. + + Checked by identity against the reader's own queue rather than by counting + subscribers, so unrelated churn (per-request callback subscriptions, the + room's own subscribers going away at disconnect) can't affect the result. + """ + queue = reader._queue + return any(q is queue for q, _, _ in FfiClient.instance.queue._subscribers) + + +async def wait_until_unsubscribed(reader: Any, timeout: float = 5.0) -> None: + deadline = asyncio.get_event_loop().time() + timeout + while asyncio.get_event_loop().time() < deadline: + if not reader_is_subscribed(reader): + return + await asyncio.sleep(0.02) + raise AssertionError("reader is still subscribed to the FFI event queue") + + async def connect_rooms( room_name: str, *, @@ -441,6 +462,145 @@ async def read() -> None: await sender.disconnect() +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_abandoned_reader_releases_subscription_on_close() -> None: + """A reader abandoned mid-iteration never consumes its terminal eos event, + so it stays subscribed to the FFI queue until close() is called.""" + receiver, sender = await connect_rooms(unique_room_name("ds-abandon")) + try: + got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future() + abandoned = asyncio.Event() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_reader.set_result(reader) + + async def read() -> None: + async for _chunk in reader: + break # walk away without draining to end-of-stream + abandoned.set() + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("abandon-topic", handler) + + await sender.local_participant.send_text(pseudo_random_text(50_000), topic="abandon-topic") + reader = await asyncio.wait_for(got_reader, timeout=10.0) + await asyncio.wait_for(abandoned.wait(), timeout=10.0) + + assert reader_is_subscribed(reader) + reader.close() + await wait_until_unsubscribed(reader) + + reader.close() # idempotent + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_never_read_reader_releases_subscription_on_close() -> None: + """A handler that never iterates its reader still leaves a subscription + behind, because readers subscribe eagerly on construction.""" + receiver, sender = await connect_rooms(unique_room_name("ds-noread")) + writer = None + try: + got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_reader.set_result(reader) # never read + + receiver.register_text_stream_handler("noread-topic", handler) + + writer = await sender.local_participant.stream_text(topic="noread-topic") + await writer.write("some data") + + reader = await asyncio.wait_for(got_reader, timeout=10.0) + assert reader_is_subscribed(reader) + + await reader.aclose() + await wait_until_unsubscribed(reader) + finally: + if writer is not None: + await writer.aclose() + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_disconnect_unsubscribes_unread_reader() -> None: + """Disconnecting must drop the subscriptions of readers the application + never consumed, without the application closing them. + + The reader is left open, so a read starting after the disconnect still + drains whatever arrived beforehand and then reports the failure — the + behaviour callers had before disconnect released the subscription. + """ + receiver, sender = await connect_rooms(unique_room_name("ds-disc-unread")) + writer = None + try: + got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_reader.set_result(reader) # never read + + receiver.register_text_stream_handler("disc-unread-topic", handler) + + writer = await sender.local_participant.stream_text(topic="disc-unread-topic") + await writer.write("buffered ") + await writer.write("data") + + reader = await asyncio.wait_for(got_reader, timeout=10.0) + assert reader_is_subscribed(reader) + # let the chunks land in the reader's queue while nothing is reading + await asyncio.sleep(0.5) + + await receiver.disconnect() + await wait_until_unsubscribed(reader) + + # chunks received before the disconnect are still delivered, and the + # stream then terminates with StreamError rather than a clean EOF + chunks = [] + with pytest.raises(rtc.StreamError): + async for chunk in reader: + chunks.append(chunk) + assert "".join(chunks) == "buffered data" + finally: + if writer is not None: + await writer.aclose() + await sender.disconnect() + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_fully_read_reader_needs_no_close() -> None: + """Reading to end-of-stream unsubscribes on its own; close() afterwards is + a no-op.""" + receiver, sender = await connect_rooms(unique_room_name("ds-drain")) + try: + got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future() + received: asyncio.Future[str] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_reader.set_result(reader) + + async def read() -> None: + received.set_result(await reader.read_all()) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("drain-topic", handler) + + await sender.local_participant.send_text("hello", topic="drain-topic") + assert await asyncio.wait_for(received, timeout=10.0) == "hello" + + reader = await asyncio.wait_for(got_reader, timeout=5.0) + await wait_until_unsubscribed(reader) + reader.close() + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + @pytest.mark.asyncio @skip_if_no_credentials() # type: ignore[untyped-decorator] async def test_send_file_round_trip(tmp_path: Any) -> None: From 59104b727faaa8cfaf18f1c33628c3d9e2dcb692 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Fri, 31 Jul 2026 13:52:56 -0400 Subject: [PATCH 4/9] fix: backport abandoned data stream writer drop from node-sdks a6f8a35 --- livekit-rtc/livekit/rtc/data_stream.py | 21 +++++++++ livekit-rtc/livekit/rtc/participant.py | 24 +++++++++++ livekit-rtc/livekit/rtc/room.py | 12 ++++++ tests/rtc/test_e2e_data_streams.py | 60 +++++++++++++++++++++++++- 4 files changed, 116 insertions(+), 1 deletion(-) diff --git a/livekit-rtc/livekit/rtc/data_stream.py b/livekit-rtc/livekit/rtc/data_stream.py index a033c8f7..cca431f3 100644 --- a/livekit-rtc/livekit/rtc/data_stream.py +++ b/livekit-rtc/livekit/rtc/data_stream.py @@ -378,6 +378,21 @@ def __init__( def _provisional_timestamp(self) -> int: return int(datetime.datetime.now().timestamp() * 1000) + def _register_open(self) -> None: + """Records the freshly opened writer so that, if it is never closed, + the room can drop its native handle on disconnect.""" + assert self._writer_handle is not None + self._local_participant._open_stream_writers.add(self._writer_handle) + + def _unregister_open(self) -> None: + """Forgets the writer because a close request has been issued for it. + + The close consumes the handle on the native side, so it must no longer + be dropped by the disconnect cleanup. + """ + if self._writer_handle is not None: + self._local_participant._open_stream_writers.discard(self._writer_handle) + async def _wait_for_callback( self, req: proto_ffi.FfiRequest, callback_field: str, response_field: str ) -> proto_ffi.FfiEvent: @@ -403,6 +418,10 @@ async def aclose( if self._writer_handle is None: raise RuntimeError("Stream is not open") self._closed = True + # unregister before sending: the request consumes the handle natively, + # so it must not be dropped again by the disconnect cleanup even if the + # close itself reports an error + self._unregister_open() await self._send_close(reason=reason, attributes=attributes) async def _send_close(self, *, reason: str, attributes: Optional[Dict[str, str]]) -> None: @@ -461,6 +480,7 @@ async def _send_header(self) -> None: cb = await self._wait_for_callback(req, "text_stream_open", "text_stream_open") owned = cb.text_stream_open.writer self._writer_handle = owned.handle.id + self._register_open() self._info = _text_stream_info_from_proto(owned.info) if self._info.size is None and self._total_size is not None: # total_size is not transmitted for text streams; keep it local @@ -546,6 +566,7 @@ async def _send_header(self) -> None: cb = await self._wait_for_callback(req, "byte_stream_open", "byte_stream_open") owned = cb.byte_stream_open.writer self._writer_handle = owned.handle.id + self._register_open() self._info = _byte_stream_info_from_proto(owned.info) async def write(self, data: bytes) -> None: diff --git a/livekit-rtc/livekit/rtc/participant.py b/livekit-rtc/livekit/rtc/participant.py index 6e6cd556..45293ec7 100644 --- a/livekit-rtc/livekit/rtc/participant.py +++ b/livekit-rtc/livekit/rtc/participant.py @@ -246,6 +246,14 @@ def __init__( self._room_queue = room_queue self._track_publications: dict[str, LocalTrackPublication] = {} self._rpc_handlers: Dict[str, RpcHandler] = {} + # Handle ids of data stream writers that have been opened but not yet + # closed. The FFI close request consumes the handle (take_handle), so + # an entry is removed as soon as a close is sent for it; whatever is + # left when the room disconnects is dropped by + # _dispose_open_stream_writers(). Only the ids are kept — wrapping them + # in FfiHandle would make GC drop a handle the close request already + # consumed. + self._open_stream_writers: set[int] = set() @property def track_publications(self) -> Mapping[str, LocalTrackPublication]: @@ -254,6 +262,22 @@ def track_publications(self) -> Mapping[str, LocalTrackPublication]: """ return self._track_publications + def _dispose_open_stream_writers(self) -> None: + """Drops the native writers of any data streams that were opened but + never closed — a writer the caller abandoned, or one left behind by a + failed write. Their handles would otherwise be held by the FFI server + for the lifetime of the process. + + The room is already gone at this point, so the handles are dropped + directly rather than closed: there is no end-of-stream left to deliver. + """ + for writer_handle in self._open_stream_writers: + try: + FfiHandle(writer_handle).dispose() + except Exception: + logger.exception("failed to dispose data stream writer handle") + self._open_stream_writers.clear() + async def publish_data( self, payload: Union[bytes, str], diff --git a/livekit-rtc/livekit/rtc/room.py b/livekit-rtc/livekit/rtc/room.py index c67b0950..8f89e8a5 100644 --- a/livekit-rtc/livekit/rtc/room.py +++ b/livekit-rtc/livekit/rtc/room.py @@ -686,6 +686,7 @@ async def disconnect( await self._drain_rpc_invocation_tasks() self._error_stream_readers() + self._dispose_open_stream_writers() req = proto_ffi.FfiRequest() req.disconnect.room_handle = self._ffi_handle.handle # type: ignore @@ -735,6 +736,7 @@ async def _listen_task(self) -> None: # Clean up any pending RPC invocation tasks await self._drain_rpc_invocation_tasks() self._error_stream_readers() + self._dispose_open_stream_writers() def _on_rpc_method_invocation(self, rpc_invocation: RpcMethodInvocationEvent) -> None: if self._local_participant is None: @@ -1195,6 +1197,16 @@ def on_close() -> None: self._byte_stream_readers[handle_id] = byte_reader byte_stream_handler(byte_reader, opened.participant_identity) + def _dispose_open_stream_writers(self) -> None: + """Drops the native writers of data streams left open at disconnect. + + A writer releases its handle only when it is closed, so one the + application abandoned — or left behind after a failed write — would + otherwise be held by the FFI server for the life of the process. + """ + if self._local_participant is not None: + self._local_participant._dispose_open_stream_writers() + def _error_stream_readers(self) -> None: """Wakes up any pending stream reads with a StreamError on disconnect. diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py index 480e525f..b1168e51 100644 --- a/tests/rtc/test_e2e_data_streams.py +++ b/tests/rtc/test_e2e_data_streams.py @@ -26,7 +26,7 @@ import pytest from livekit import api, rtc -from livekit.rtc._ffi_client import FfiClient +from livekit.rtc._ffi_client import FfiClient, FfiHandle def skip_if_no_credentials() -> Any: @@ -73,6 +73,21 @@ def reader_is_subscribed(reader: Any) -> bool: return any(q is queue for q, _, _ in FfiClient.instance.queue._subscribers) +def ffi_handle_is_dropped(handle_id: int) -> bool: + """True when the FFI server no longer holds ``handle_id``. + + Probes by trying to drop it: ``FfiHandle.dispose()`` asserts that the + native drop succeeded, so a handle that is already gone raises. Destructive + when the handle *is* still held — it gets dropped — so only call this once, + at the end of a test. + """ + try: + FfiHandle(handle_id).dispose() + except AssertionError: + return True + return False + + async def wait_until_unsubscribed(reader: Any, timeout: float = 5.0) -> None: deadline = asyncio.get_event_loop().time() + timeout while asyncio.get_event_loop().time() < deadline: @@ -601,6 +616,49 @@ async def read() -> None: await asyncio.gather(receiver.disconnect(), sender.disconnect()) +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_writer_handle_untracked_after_close() -> None: + """Closing a writer consumes its handle natively, so it must drop out of + the open-writer set and not be dropped again at disconnect.""" + receiver, sender = await connect_rooms(unique_room_name("ds-writer-close")) + try: + local = sender.local_participant + writer = await local.stream_text(topic="writer-close-topic") + handle = writer._writer_handle + assert handle is not None + assert handle in local._open_stream_writers + + await writer.write("data") + await writer.aclose() + + assert handle not in local._open_stream_writers + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_abandoned_writer_disposed_on_disconnect() -> None: + """A writer opened and never closed holds its native writer in the FFI + handle table; disconnecting must drop it.""" + receiver, sender = await connect_rooms(unique_room_name("ds-writer-abandon")) + try: + writer = await sender.local_participant.stream_text(topic="writer-abandon-topic") + handle = writer._writer_handle + assert handle is not None + + await writer.write("data") # abandoned: aclose() is never called + + await sender.disconnect() + + # asserted against the FFI handle table rather than the SDK's + # bookkeeping, so this fails on an actually-leaked native writer + assert ffi_handle_is_dropped(handle) + finally: + await receiver.disconnect() + + @pytest.mark.asyncio @skip_if_no_credentials() # type: ignore[untyped-decorator] async def test_send_file_round_trip(tmp_path: Any) -> None: From 68616e3c46d6e5247bd0427af9cc1e93d1d7a79d Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Fri, 31 Jul 2026 15:41:17 -0400 Subject: [PATCH 5/9] fix: address data stream deadlock case when reading a data stream which is closed in the background --- livekit-rtc/livekit/rtc/_ffi_client.py | 9 ++++ livekit-rtc/livekit/rtc/data_stream.py | 67 +++++++++++++++++++++----- livekit-rtc/livekit/rtc/participant.py | 25 ++++++---- 3 files changed, 78 insertions(+), 23 deletions(-) diff --git a/livekit-rtc/livekit/rtc/_ffi_client.py b/livekit-rtc/livekit/rtc/_ffi_client.py index 81ed71d3..3ae44aac 100644 --- a/livekit-rtc/livekit/rtc/_ffi_client.py +++ b/livekit-rtc/livekit/rtc/_ffi_client.py @@ -80,6 +80,15 @@ def __del__(self) -> None: def disposed(self) -> bool: return self._disposed + def mark_consumed(self) -> None: + """Marks the handle as already released on the native side. + + Some requests take ownership of their handle (``take_handle``) rather + than borrowing it. After issuing one, the handle must not be dropped + again, so this suppresses the drop that dispose/GC would otherwise do. + """ + self._disposed = True + def dispose(self) -> None: if self.handle != INVALID_HANDLE and not self._disposed: self._disposed = True diff --git a/livekit-rtc/livekit/rtc/data_stream.py b/livekit-rtc/livekit/rtc/data_stream.py index cca431f3..5c4bffe6 100644 --- a/livekit-rtc/livekit/rtc/data_stream.py +++ b/livekit-rtc/livekit/rtc/data_stream.py @@ -21,7 +21,7 @@ from typing import AsyncIterator, Optional, Dict, List from ._proto import data_stream_pb2 as proto_data_stream from ._proto import ffi_pb2 as proto_ffi -from ._ffi_client import FfiClient +from ._ffi_client import FfiClient, FfiHandle from typing import TYPE_CHECKING @@ -185,10 +185,25 @@ def _unsubscribe(self) -> None: """ FfiClient.instance.queue.unsubscribe(self._queue) + def _wake_pending_reads(self) -> None: + """Pushes a terminal event into this reader's own queue. + + Closing unsubscribes, so no further FFI event can ever arrive; without + this, a read already parked in ``await self._queue.get()`` would wait + forever. The sentinel carries no error, so the woken read raises + ``StopAsyncIteration`` — or the stored :class:`StreamError`, which + ``__anext__`` checks first. + """ + event = proto_ffi.FfiEvent() + event.text_stream_reader_event.reader_handle = self._reader_handle + event.text_stream_reader_event.eos.SetInParent() + self._queue.put_nowait(event) + def _close(self) -> None: if not self._closed: self._closed = True self._unsubscribe() + self._wake_pending_reads() if self._on_close is not None: self._on_close() @@ -314,10 +329,25 @@ def _unsubscribe(self) -> None: """ FfiClient.instance.queue.unsubscribe(self._queue) + def _wake_pending_reads(self) -> None: + """Pushes a terminal event into this reader's own queue. + + Closing unsubscribes, so no further FFI event can ever arrive; without + this, a read already parked in ``await self._queue.get()`` would wait + forever. The sentinel carries no error, so the woken read raises + ``StopAsyncIteration`` — or the stored :class:`StreamError`, which + ``__anext__`` checks first. + """ + event = proto_ffi.FfiEvent() + event.byte_stream_reader_event.reader_handle = self._reader_handle + event.byte_stream_reader_event.eos.SetInParent() + self._queue.put_nowait(event) + def _close(self) -> None: if not self._closed: self._closed = True self._unsubscribe() + self._wake_pending_reads() if self._on_close is not None: self._on_close() @@ -372,17 +402,27 @@ def __init__( # the writer handle is assigned by the FFI when the stream is opened; # the close request consumes it, so it is kept as a raw id self._writer_handle: Optional[int] = None + self._writer_ffi_handle: Optional[FfiHandle] = None self._write_lock = asyncio.Lock() self._closed = False def _provisional_timestamp(self) -> int: return int(datetime.datetime.now().timestamp() * 1000) - def _register_open(self) -> None: - """Records the freshly opened writer so that, if it is never closed, - the room can drop its native handle on disconnect.""" - assert self._writer_handle is not None - self._local_participant._open_stream_writers.add(self._writer_handle) + def _register_open(self, handle_id: int) -> None: + """Takes ownership of a freshly opened writer handle. + + Wrapping it in an FfiHandle means a writer that is simply dropped — + abandoned, or left behind by an exception or a cancelled task — still + releases the native writer when it is garbage collected. It is also + registered with the participant so the room can drop it deterministically + at disconnect, for writers still referenced at that point. The registry + holds the handle weakly, so registration does not itself keep an + abandoned writer alive. + """ + self._writer_handle = handle_id + self._writer_ffi_handle = FfiHandle(handle_id) + self._local_participant._open_stream_writers[handle_id] = self._writer_ffi_handle def _unregister_open(self) -> None: """Forgets the writer because a close request has been issued for it. @@ -390,8 +430,11 @@ def _unregister_open(self) -> None: The close consumes the handle on the native side, so it must no longer be dropped by the disconnect cleanup. """ - if self._writer_handle is not None: - self._local_participant._open_stream_writers.discard(self._writer_handle) + if self._writer_handle is None: + return + self._local_participant._open_stream_writers.pop(self._writer_handle, None) + if self._writer_ffi_handle is not None: + self._writer_ffi_handle.mark_consumed() async def _wait_for_callback( self, req: proto_ffi.FfiRequest, callback_field: str, response_field: str @@ -479,10 +522,9 @@ async def _send_header(self) -> None: req.text_stream_open.options.CopyFrom(self._options) cb = await self._wait_for_callback(req, "text_stream_open", "text_stream_open") owned = cb.text_stream_open.writer - self._writer_handle = owned.handle.id - self._register_open() + self._register_open(owned.handle.id) self._info = _text_stream_info_from_proto(owned.info) - if self._info.size is None and self._total_size is not None: + if not self._info.size and self._total_size is not None: # total_size is not transmitted for text streams; keep it local self._info.size = self._total_size @@ -565,8 +607,7 @@ async def _send_header(self) -> None: req.byte_stream_open.options.CopyFrom(self._options) cb = await self._wait_for_callback(req, "byte_stream_open", "byte_stream_open") owned = cb.byte_stream_open.writer - self._writer_handle = owned.handle.id - self._register_open() + self._register_open(owned.handle.id) self._info = _byte_stream_info_from_proto(owned.info) async def write(self, data: bytes) -> None: diff --git a/livekit-rtc/livekit/rtc/participant.py b/livekit-rtc/livekit/rtc/participant.py index 45293ec7..a9e3d799 100644 --- a/livekit-rtc/livekit/rtc/participant.py +++ b/livekit-rtc/livekit/rtc/participant.py @@ -20,6 +20,7 @@ import enum import os import mimetypes +import weakref from typing import List, Union, Callable, Dict, Awaitable, Optional, Mapping, cast, TypeVar from abc import abstractmethod, ABC @@ -246,14 +247,18 @@ def __init__( self._room_queue = room_queue self._track_publications: dict[str, LocalTrackPublication] = {} self._rpc_handlers: Dict[str, RpcHandler] = {} - # Handle ids of data stream writers that have been opened but not yet - # closed. The FFI close request consumes the handle (take_handle), so - # an entry is removed as soon as a close is sent for it; whatever is - # left when the room disconnects is dropped by - # _dispose_open_stream_writers(). Only the ids are kept — wrapping them - # in FfiHandle would make GC drop a handle the close request already - # consumed. - self._open_stream_writers: set[int] = set() + # Handles of data stream writers that have been opened but not yet + # closed, so the room can drop them at disconnect. The FFI close + # request consumes the handle (take_handle), so an entry is removed as + # soon as a close is sent for it. + # + # Held weakly: the writer itself owns the FfiHandle, and a writer that + # is simply abandoned should be collectable so its handle is dropped at + # GC rather than lingering until disconnect. A strong reference here + # would pin every writer for the life of the room. + self._open_stream_writers: weakref.WeakValueDictionary[int, FfiHandle] = ( + weakref.WeakValueDictionary() + ) @property def track_publications(self) -> Mapping[str, LocalTrackPublication]: @@ -271,9 +276,9 @@ def _dispose_open_stream_writers(self) -> None: The room is already gone at this point, so the handles are dropped directly rather than closed: there is no end-of-stream left to deliver. """ - for writer_handle in self._open_stream_writers: + for ffi_handle in list(self._open_stream_writers.values()): try: - FfiHandle(writer_handle).dispose() + ffi_handle.dispose() except Exception: logger.exception("failed to dispose data stream writer handle") self._open_stream_writers.clear() From 5c6ff6329f718c6a2576eed2afd280e6bd7821b0 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Fri, 31 Jul 2026 15:42:24 -0400 Subject: [PATCH 6/9] fix: address backwards compatibility size should be 0 like the current `main` state when a data stream does not have an associated size --- livekit-rtc/livekit/rtc/data_stream.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/livekit-rtc/livekit/rtc/data_stream.py b/livekit-rtc/livekit/rtc/data_stream.py index 5c4bffe6..35bdc6c3 100644 --- a/livekit-rtc/livekit/rtc/data_stream.py +++ b/livekit-rtc/livekit/rtc/data_stream.py @@ -70,7 +70,7 @@ def _text_stream_info_from_proto(info: proto_data_stream.TextStreamInfo) -> Text mime_type=info.mime_type, topic=info.topic, timestamp=info.timestamp, - size=info.total_length if info.HasField("total_length") else None, + size=info.total_length if info.HasField("total_length") else 0, attributes=dict(info.attributes), attachments=list(info.attached_stream_ids), ) @@ -82,7 +82,7 @@ def _byte_stream_info_from_proto(info: proto_data_stream.ByteStreamInfo) -> Byte mime_type=info.mime_type, topic=info.topic, timestamp=info.timestamp, - size=info.total_length if info.HasField("total_length") else None, + size=info.total_length if info.HasField("total_length") else 0, attributes=dict(info.attributes), name=info.name, ) From dfc8620452c1559df400cd5f29737d56bd4e90e5 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Fri, 31 Jul 2026 15:45:11 -0400 Subject: [PATCH 7/9] feat: add e2e cases for newly exposed deadlock read case from 68616e3 --- tests/rtc/test_e2e_data_streams.py | 116 +++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py index b1168e51..55b98718 100644 --- a/tests/rtc/test_e2e_data_streams.py +++ b/tests/rtc/test_e2e_data_streams.py @@ -16,6 +16,7 @@ """ import asyncio +import gc import os import random import string @@ -659,6 +660,121 @@ async def test_abandoned_writer_disposed_on_disconnect() -> None: await receiver.disconnect() +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_close_wakes_pending_read() -> None: + """Closing a reader while a read is parked awaiting the next chunk must end + that read, not leave it waiting for an event that can no longer arrive.""" + receiver, sender = await connect_rooms(unique_room_name("ds-close-wake")) + writer = None + try: + got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future() + done: asyncio.Future[list] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_reader.set_result(reader) + + async def read() -> None: + chunks = [chunk async for chunk in reader] + done.set_result(chunks) + + asyncio.create_task(read()) + + receiver.register_text_stream_handler("close-wake-topic", handler) + + writer = await sender.local_participant.stream_text(topic="close-wake-topic") + await writer.write("first") + + reader = await asyncio.wait_for(got_reader, timeout=10.0) + # let the reader consume "first" and park waiting for more, with the + # stream deliberately left open by the sender + await asyncio.sleep(0.5) + assert not done.done() + + reader.close() + + # without a wake-up sentinel this wait_for is what times out + chunks = await asyncio.wait_for(done, timeout=5.0) + assert chunks == ["first"] + finally: + if writer is not None: + await writer.aclose() + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_undeclared_stream_size_is_zero_not_none() -> None: + """An incremental text stream cannot declare its length on the wire, so the + receiver sees a size of 0 — an int, as callers have always been given, + rather than None.""" + receiver, sender = await connect_rooms(unique_room_name("ds-size-shape")) + writer = None + try: + got_info: asyncio.Future[rtc.TextStreamInfo] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_info.set_result(reader.info) + + receiver.register_text_stream_handler("size-topic", handler) + + writer = await sender.local_participant.stream_text(topic="size-topic", total_size=1234) + await writer.write("data") + + info = await asyncio.wait_for(got_info, timeout=10.0) + assert info.size == 0 + assert isinstance(info.size, int) + # the sender still sees what it declared + assert writer.info.size == 1234 + finally: + if writer is not None: + await writer.aclose() + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_send_text_size_reaches_receiver() -> None: + """send_text lets the FFI compute and transmit the length, so the receiver + does see a real size.""" + receiver, sender = await connect_rooms(unique_room_name("ds-size-sendtext")) + try: + text = "a size-carrying message" + got_info: asyncio.Future[rtc.TextStreamInfo] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_info.set_result(reader.info) + + receiver.register_text_stream_handler("sendtext-size-topic", handler) + + await sender.local_participant.send_text(text, topic="sendtext-size-topic") + + info = await asyncio.wait_for(got_info, timeout=10.0) + assert info.size == len(text.encode()) + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_abandoned_writer_disposed_on_gc() -> None: + """Dropping a writer without closing it releases the native writer at GC, + rather than holding it until the room disconnects.""" + receiver, sender = await connect_rooms(unique_room_name("ds-writer-gc")) + try: + writer = await sender.local_participant.stream_text(topic="writer-gc-topic") + handle = writer._writer_handle + assert handle is not None + await writer.write("data") + + del writer # abandoned without aclose() + gc.collect() + + assert ffi_handle_is_dropped(handle) + finally: + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + @pytest.mark.asyncio @skip_if_no_credentials() # type: ignore[untyped-decorator] async def test_send_file_round_trip(tmp_path: Any) -> None: From d28284029e9ee8465234608f98d404bfb309b551 Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Mon, 3 Aug 2026 13:00:52 -0400 Subject: [PATCH 8/9] feat: add closed reader is empty stream test --- tests/rtc/test_e2e_data_streams.py | 35 ++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py index 55b98718..ae0ab6c9 100644 --- a/tests/rtc/test_e2e_data_streams.py +++ b/tests/rtc/test_e2e_data_streams.py @@ -775,6 +775,41 @@ async def test_abandoned_writer_disposed_on_gc() -> None: await asyncio.gather(receiver.disconnect(), sender.disconnect()) +@pytest.mark.asyncio +@skip_if_no_credentials() # type: ignore[untyped-decorator] +async def test_reads_a_closed_reader_as_an_empty_stream() -> None: + """A reader closed before anything read it reads back as an empty stream. + + The handler never touches the reader and the sender writes a chunk, so a + chunk may or may not have been queued by the time close() lands. Either + way ``read_all()`` must resolve to '' promptly rather than hanging (the + failure mode a closed-but-unwoken reader had) or raising. Note this + asserts termination, not buffer-discard: the test cannot observe whether + the chunk arrived before the close, so '' is the correct result either way. + """ + receiver, sender = await connect_rooms(unique_room_name("ds-closed-readall")) + writer = None + try: + got_reader: asyncio.Future[rtc.TextStreamReader] = asyncio.get_event_loop().create_future() + + def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + got_reader.set_result(reader) # never read + + receiver.register_text_stream_handler("closed-readall-topic", handler) + + writer = await sender.local_participant.stream_text(topic="closed-readall-topic") + await writer.write("a chunk") + + reader = await asyncio.wait_for(got_reader, timeout=10.0) + reader.close() + + assert await asyncio.wait_for(reader.read_all(), timeout=5.0) == "" + finally: + if writer is not None: + await writer.aclose() + await asyncio.gather(receiver.disconnect(), sender.disconnect()) + + @pytest.mark.asyncio @skip_if_no_credentials() # type: ignore[untyped-decorator] async def test_send_file_round_trip(tmp_path: Any) -> None: From c9b01cb0ca2a60dff33ee034b73030f7019f7a3b Mon Sep 17 00:00:00 2001 From: Ryan Gaus Date: Mon, 3 Aug 2026 13:02:14 -0400 Subject: [PATCH 9/9] fix: make test_receiver_rejects_oversized_payload more robust Check the exact code paths rather than just "absence any read_all call" --- tests/rtc/test_e2e_data_streams.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/rtc/test_e2e_data_streams.py b/tests/rtc/test_e2e_data_streams.py index ae0ab6c9..2f39a42e 100644 --- a/tests/rtc/test_e2e_data_streams.py +++ b/tests/rtc/test_e2e_data_streams.py @@ -347,9 +347,12 @@ async def test_receiver_rejects_oversized_payload() -> None: ) try: text = pseudo_random_text(50_000) + handler_called = asyncio.Event() result: asyncio.Future[BaseException] = asyncio.get_event_loop().create_future() def handler(reader: rtc.TextStreamReader, _identity: str) -> None: + handler_called.set() + async def read() -> None: try: data = await reader.read_all() @@ -365,8 +368,11 @@ async def read() -> None: await sender.local_participant.send_text(text, topic="capped-topic") + await asyncio.wait_for(handler_called.wait(), timeout=10.0) error = await asyncio.wait_for(result, timeout=10.0) - assert "maximum size" in str(error) + # the full text, not just "maximum size": HeaderTooLarge reads + # "stream header exceeds maximum size" and shares that substring + assert "payload exceeds maximum size" in str(error) finally: await asyncio.gather(receiver.disconnect(), sender.disconnect())