From 11a77d732f0ad0bf6314a50ce94cb7701b15af6a Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Wed, 19 Aug 2026 23:27:51 +0000 Subject: [PATCH 1/6] Add a standalone WebRTC client window Provide v2 browser input buffering and video delivery without depending on the existing serving stack. --- .../runtime_v2/serving/webrtc_server.py | 523 ++++++++++++++++++ .../runtime_v2/webrtc_client_window.py | 57 ++ .../test_v2/test_webrtc_client_window.py | 147 +++++ 3 files changed, 727 insertions(+) create mode 100644 flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py create mode 100644 flashdreams/flashdreams/runtime_v2/webrtc_client_window.py create mode 100644 flashdreams/test_v2/test_webrtc_client_window.py diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py new file mode 100644 index 000000000..9ff387a62 --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -0,0 +1,523 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Standalone WebRTC server used by the v2 client window.""" + +from __future__ import annotations + +import asyncio +import json +import socket +import threading +import time +from collections import deque +from fractions import Fraction +from typing import Any + +import numpy as np +import torch +from aiohttp import web +from aiortc import MediaStreamTrack, RTCPeerConnection, RTCSessionDescription +from aiortc.mediastreams import MediaStreamError +from av import VideoFrame + +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import ( + CloseUserInputEventData, + KeyboardUserInputEventData, + ResetUserInputEventData, + UserInputEvent, +) +from flashdreams.runtime_v2.user_input_events import UserInputEvents +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout + +_BROWSER_PAGE = """ + + + + + + + +""" + + +class _VideoTrack(MediaStreamTrack): + """Video track whose frames are supplied by the server.""" + + kind = "video" + + def __init__(self, frames_per_second: int) -> None: + super().__init__() + self._frames_per_second = frames_per_second + self._time_base = Fraction(1, frames_per_second) + self._frames: asyncio.Queue[np.ndarray[Any, np.dtype[np.uint8]] | None] = ( + asyncio.Queue() + ) + self._next_frame_time: float | None = None + self._pts = 0 + self._closed = False + + async def enqueue( + self, frames: tuple[np.ndarray[Any, np.dtype[np.uint8]], ...] + ) -> None: + """Append generated RGB frames for the WebRTC sender.""" + if self._closed: + return + for frame in frames: + await self._frames.put(frame) + + async def recv(self) -> VideoFrame: + """Return the next generated frame when aiortc requests one.""" + if self._closed: + raise MediaStreamError + frame = await self._frames.get() + if frame is None: + raise MediaStreamError + + loop = asyncio.get_running_loop() + now = loop.time() + if self._next_frame_time is None: + self._next_frame_time = now + else: + self._next_frame_time += 1.0 / self._frames_per_second + await asyncio.sleep(max(0.0, self._next_frame_time - now)) + + video_frame = VideoFrame.from_ndarray(frame, format="rgb24") + video_frame.pts = self._pts + video_frame.time_base = self._time_base + self._pts += 1 + return video_frame + + async def close(self) -> None: + """Stop the track and release a pending receiver.""" + if self._closed: + return + self._closed = True + self._frames.put_nowait(None) + self.stop() + + +class WebRTCServer: + """Own the HTTP, signaling, input buffering, and media transport.""" + + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 0, + startup_timeout_seconds: float = 10.0, + ) -> None: + """ + Args: + host: Interface on which the HTTP server listens. + port: Listening port. Zero asks the operating system to choose one. + startup_timeout_seconds: Maximum time to wait for server startup. + + Raises: + RuntimeError: The server cannot start. + TimeoutError: The server does not start before the timeout. + """ + if not host: + raise ValueError("host must not be empty.") + if port < 0 or port > 65535: + raise ValueError("port must be between 0 and 65535.") + if startup_timeout_seconds <= 0: + raise ValueError("startup_timeout_seconds must be > 0.") + + self._host = host + self._port = port + self._startup_timeout_seconds = startup_timeout_seconds + self._events: deque[UserInputEvent] = deque() + self._events_lock = threading.Lock() + self._started = threading.Event() + self._startup_error: BaseException | None = None + self._loop: asyncio.AbstractEventLoop | None = None + self._runner: web.AppRunner | None = None + self._peer_connection: RTCPeerConnection | None = None + self._video_track: _VideoTrack | None = None + self._session_desc: SessionDesc | None = None + self._session_start_ns: int | None = None + self._closed = False + self._client_connected = False + self._thread = threading.Thread( + target=self._run_server, + name="flashdreams-webrtc", + daemon=True, + ) + self._thread.start() + if not self._started.wait(startup_timeout_seconds): + raise TimeoutError("WebRTC server did not start before the timeout.") + if self._startup_error is not None: + raise RuntimeError( + "WebRTC server failed to start." + ) from self._startup_error + + @property + def host(self) -> str: + """Return the interface on which the server is listening.""" + return self._host + + @property + def port(self) -> int: + """Return the bound server port.""" + return self._port + + @property + def url(self) -> str: + """Return the browser URL for this server.""" + return f"http://{self._host}:{self._port}/" + + def open(self, session_desc: SessionDesc) -> None: + """Configure the server for one session's generated video. + + Args: + session_desc: Resolved dimensions, frame rate, and tensor layout. + + Raises: + RuntimeError: The server is closed or already open. + """ + if self._closed: + raise RuntimeError("Cannot open a closed WebRTC server.") + if self._session_desc is not None: + raise RuntimeError("WebRTC server is already open.") + self._session_desc = session_desc + self._session_start_ns = time.monotonic_ns() + + def get_user_input_events(self) -> UserInputEvents: + """Drain and return browser events in timestamp order.""" + with self._events_lock: + events = list(self._events) + self._events.clear() + return UserInputEvents(events) + + def write(self, result: StepResult) -> None: + """Deliver one generated result to the browser's video track. + + Args: + result: Generated frames matching the description passed to + :meth:`open`. + + Raises: + RuntimeError: The server is not open or has been closed. + ValueError: The result shape or layout does not match the session. + """ + if self._closed: + raise RuntimeError("Cannot write to a closed WebRTC server.") + session_desc = self._session_desc + if session_desc is None: + raise RuntimeError("Open the WebRTC server before writing.") + frames = _result_to_rgb_frames(result, session_desc) + loop = self._loop + if loop is None: + raise RuntimeError("WebRTC server is not running.") + future = asyncio.run_coroutine_threadsafe(self._enqueue_frames(frames), loop) + future.result() + + def close(self) -> None: + """Close the peer connection and stop the WebRTC server.""" + if self._closed: + return + self._closed = True + loop = self._loop + if loop is None: + return + future = asyncio.run_coroutine_threadsafe(self._shutdown(), loop) + future.result(timeout=self._startup_timeout_seconds) + loop.call_soon_threadsafe(loop.stop) + self._thread.join(timeout=self._startup_timeout_seconds) + if self._thread.is_alive(): + raise TimeoutError("WebRTC server did not stop before the timeout.") + + def _run_server(self) -> None: + """Own the WebRTC asyncio loop for the lifetime of the server.""" + loop = asyncio.new_event_loop() + self._loop = loop + asyncio.set_event_loop(loop) + try: + loop.run_until_complete(self._start_server()) + except BaseException as error: + self._startup_error = error + self._started.set() + loop.close() + return + self._started.set() + try: + loop.run_forever() + finally: + loop.close() + + async def _start_server(self) -> None: + """Create and bind the standalone aiohttp application.""" + app = web.Application() + app.router.add_get("/", self._serve_browser) + app.router.add_get("/healthz", self._health) + app.router.add_post("/api/webrtc/offer", self._offer) + runner = web.AppRunner(app) + await runner.setup() + address_family = socket.AF_INET6 if ":" in self._host else socket.AF_INET + server_socket = socket.socket(address_family, socket.SOCK_STREAM) + try: + server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server_socket.bind((self._host, self._port)) + server_socket.setblocking(False) + server_socket.listen(128) + self._port = int(server_socket.getsockname()[1]) + site = web.SockSite(runner, server_socket) + await site.start() + except Exception: + server_socket.close() + await runner.cleanup() + raise + self._runner = runner + + async def _serve_browser(self, _: web.Request) -> web.Response: + """Return the minimal browser client.""" + return web.Response(text=_BROWSER_PAGE, content_type="text/html") + + async def _health(self, _: web.Request) -> web.Response: + """Report whether the server has an open session and client.""" + return web.json_response( + { + "open": self._session_desc is not None, + "client_connected": self._client_connected, + } + ) + + async def _offer(self, request: web.Request) -> web.Response: + """Negotiate one browser peer connection.""" + if self._closed: + raise web.HTTPServiceUnavailable(reason="WebRTC server is closed.") + session_desc = self._session_desc + if session_desc is None: + raise web.HTTPConflict(reason="WebRTC server is not open.") + if self._peer_connection is not None: + raise web.HTTPConflict(reason="A WebRTC client is already connected.") + + try: + payload = await request.json() + except (json.JSONDecodeError, web.HTTPException) as error: + raise web.HTTPBadRequest(reason="Expected a JSON WebRTC offer.") from error + if not isinstance(payload, dict): + raise web.HTTPBadRequest(reason="WebRTC offer must be an object.") + sdp = payload.get("sdp") + offer_type = payload.get("type") + if not isinstance(sdp, str) or not isinstance(offer_type, str): + raise web.HTTPBadRequest( + reason="WebRTC offer requires string sdp and type." + ) + + peer_connection = RTCPeerConnection() + video_track = _VideoTrack(session_desc.frames_per_second_for_ui) + peer_connection.addTrack(video_track) + self._peer_connection = peer_connection + self._video_track = video_track + + @peer_connection.on("datachannel") + def on_datachannel(channel: Any) -> None: + self._client_connected = True + + @channel.on("message") + def on_message(message: Any) -> None: + try: + self._buffer_browser_message(message) + except ValueError as error: + channel.send(json.dumps({"type": "error", "message": str(error)})) + + @channel.on("close") + def on_close() -> None: + self._record_client_disconnect() + + @peer_connection.on("connectionstatechange") + async def on_connectionstatechange() -> None: + if peer_connection.connectionState in {"failed", "disconnected", "closed"}: + self._record_client_disconnect() + + try: + await peer_connection.setRemoteDescription( + RTCSessionDescription(sdp=sdp, type=offer_type) + ) + await peer_connection.setLocalDescription( + await peer_connection.createAnswer() + ) + except Exception: + self._peer_connection = None + self._video_track = None + await video_track.close() + await peer_connection.close() + raise + + local_description = peer_connection.localDescription + if local_description is None: + raise web.HTTPInternalServerError( + reason="WebRTC peer did not create an answer." + ) + return web.json_response( + {"sdp": local_description.sdp, "type": local_description.type} + ) + + def _buffer_browser_message(self, raw_message: object) -> None: + """Validate and append one data-channel message.""" + if not isinstance(raw_message, str): + raise ValueError("Browser event must be a JSON string.") + try: + payload = json.loads(raw_message) + except json.JSONDecodeError as error: + raise ValueError("Browser event must contain valid JSON.") from error + if not isinstance(payload, dict): + raise ValueError("Browser event must be a JSON object.") + + event_type = payload.get("type") + if event_type == "keyboard": + key = payload.get("key") + pressed = payload.get("pressed") + if not isinstance(key, str) or not key: + raise ValueError("Keyboard event requires a non-empty key.") + if not isinstance(pressed, bool): + raise ValueError("Keyboard event requires a boolean pressed value.") + event_data = KeyboardUserInputEventData(key=key, pressed=pressed) + elif event_type == "reset": + event_data = ResetUserInputEventData() + elif event_type == "close": + event_data = CloseUserInputEventData() + else: + raise ValueError( + "Browser event type must be 'keyboard', 'reset', or 'close'." + ) + self._append_event(event_data) + + def _append_event( + self, + event_data: ( + KeyboardUserInputEventData + | ResetUserInputEventData + | CloseUserInputEventData + ), + ) -> None: + """Timestamp and buffer one validated browser event.""" + session_start_ns = self._session_start_ns + if session_start_ns is None: + return + timestamp_us = np.uint64((time.monotonic_ns() - session_start_ns) // 1_000) + with self._events_lock: + self._events.append( + UserInputEvent(timestamp=timestamp_us, event_data=event_data) + ) + + def _record_client_disconnect(self) -> None: + """Buffer one close event when the active browser disconnects.""" + if not self._client_connected: + return + self._client_connected = False + if not self._closed: + self._append_event(CloseUserInputEventData()) + + async def _enqueue_frames( + self, frames: tuple[np.ndarray[Any, np.dtype[np.uint8]], ...] + ) -> None: + """Append frames to the active media track, if connected.""" + track = self._video_track + if track is not None: + await track.enqueue(frames) + + async def _shutdown(self) -> None: + """Release async server resources on their owning loop.""" + peer_connection = self._peer_connection + self._peer_connection = None + track = self._video_track + self._video_track = None + if track is not None: + await track.close() + if peer_connection is not None: + await peer_connection.close() + runner = self._runner + self._runner = None + if runner is not None: + await runner.cleanup() + + +def _result_to_rgb_frames( + result: StepResult, session_desc: SessionDesc +) -> tuple[np.ndarray[Any, np.dtype[np.uint8]], ...]: + """Convert one result to time-major RGB uint8 frames.""" + output = result.output.detach() + if result.output_layout == VideoTensorLayout.tchw: + frames = output + elif result.output_layout == VideoTensorLayout.btchw: + if output.ndim != 5 or output.shape[0] != 1: + raise ValueError("btchw WebRTC output requires a batch size of one.") + frames = output[0] + elif result.output_layout == VideoTensorLayout.bcthw: + if output.ndim != 5 or output.shape[0] != 1: + raise ValueError("bcthw WebRTC output requires a batch size of one.") + frames = output[0].permute(1, 0, 2, 3) + elif result.output_layout == VideoTensorLayout.bvtchw: + if output.ndim != 6 or output.shape[:2] != (1, 1): + raise ValueError( + "bvtchw WebRTC output requires one batch and one video view." + ) + frames = output[0, 0] + else: + raise ValueError(f"Unsupported WebRTC output layout: {result.output_layout}.") + + if frames.ndim != 4: + raise ValueError("WebRTC output must resolve to a tchw tensor.") + if frames.shape[0] != result.frame_count: + raise ValueError("StepResult.frame_count does not match its output tensor.") + if frames.shape[1] not in (1, 3): + raise ValueError("WebRTC output must have one or three color channels.") + if frames.shape[2:] != (session_desc.video_height, session_desc.video_width): + raise ValueError("WebRTC output dimensions do not match SessionDesc.") + if result.output_layout != session_desc.output_layout: + raise ValueError("StepResult.output_layout does not match SessionDesc.") + + if frames.shape[1] == 1: + frames = frames.repeat(1, 3, 1, 1) + if frames.is_floating_point(): + frames = ((frames.to(torch.float32).clamp(-1.0, 1.0) + 1.0) * 127.5).round() + frames = frames.clamp(0, 255).to(torch.uint8) + frames = frames.permute(0, 2, 3, 1).contiguous().cpu() + return tuple(np.asarray(frame.numpy()) for frame in frames) diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py new file mode 100644 index 000000000..668626e53 --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""WebRTC client window for the v2 runtime.""" + +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.serving.webrtc_server import WebRTCServer +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_events import UserInputEvents + + +class WebRTCClientWindow(IClientWindow): + """Adapt a standalone WebRTC server to the v2 client-window interface.""" + + def __init__( + self, + *, + host: str = "127.0.0.1", + port: int = 0, + startup_timeout_seconds: float = 10.0, + ) -> None: + """ + Args: + host: Interface on which the HTTP server listens. + port: Listening port. Zero asks the operating system to choose one. + startup_timeout_seconds: Maximum time to wait for server startup. + """ + self.server = WebRTCServer( + host=host, + port=port, + startup_timeout_seconds=startup_timeout_seconds, + ) + + def open(self, session_desc: SessionDesc) -> None: + """Configure WebRTC output for the session. + + Args: + session_desc: Resolved dimensions, frame rate, and tensor layout. + """ + self.server.open(session_desc) + + def get_user_input_events(self) -> UserInputEvents: + """Drain and return buffered browser events in timestamp order.""" + return self.server.get_user_input_events() + + def write(self, result: StepResult) -> None: + """Deliver one generated result to the browser. + + Args: + result: Generated frames matching the opened session. + """ + self.server.write(result) + + def close(self) -> None: + """Close the WebRTC connection and server.""" + self.server.close() diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py new file mode 100644 index 000000000..149b17698 --- /dev/null +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -0,0 +1,147 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the v2 WebRTC client window.""" + +import asyncio +import json + +import pytest +import torch + +pytestmark = pytest.mark.ci_cpu + +pytest.importorskip("aiohttp") +pytest.importorskip("aiortc") + +from aiohttp import ClientSession +from aiortc import ( + MediaStreamTrack, + RTCDataChannel, + RTCPeerConnection, + RTCSessionDescription, +) +from av import VideoFrame + +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData +from flashdreams.runtime_v2.video_tensor import VideoTensorLayout +from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow + + +def _session_desc() -> SessionDesc: + return SessionDesc( + output_layout=VideoTensorLayout.tchw, + frames_per_second_for_ui=30, + frames_per_second_for_step=30, + video_width=16, + video_height=16, + ) + + +async def _connect_browser( + window: WebRTCClientWindow, +) -> tuple[RTCPeerConnection, RTCDataChannel, asyncio.Future[MediaStreamTrack]]: + peer = RTCPeerConnection() + channel = peer.createDataChannel("controls") + peer.addTransceiver("video", direction="recvonly") + channel_opened = asyncio.Event() + video_track: asyncio.Future[MediaStreamTrack] = ( + asyncio.get_running_loop().create_future() + ) + + @channel.on("open") + def on_open() -> None: + channel_opened.set() + + @peer.on("track") + def on_track(track: MediaStreamTrack) -> None: + if not video_track.done(): + video_track.set_result(track) + + await peer.setLocalDescription(await peer.createOffer()) + async with ClientSession() as client: + async with client.post( + f"{window.server.url}api/webrtc/offer", + json={ + "sdp": peer.localDescription.sdp, + "type": peer.localDescription.type, + }, + ) as response: + assert response.status == 200 + answer = await response.json() + await peer.setRemoteDescription( + RTCSessionDescription(sdp=answer["sdp"], type=answer["type"]) + ) + await asyncio.wait_for(channel_opened.wait(), timeout=5) + return peer, channel, video_track + + +@pytest.mark.asyncio +async def test_server_buffers_browser_events_until_drained() -> None: + window = WebRTCClientWindow() + peer: RTCPeerConnection | None = None + try: + async with ClientSession() as client: + async with client.get(f"{window.server.url}healthz") as response: + assert response.status == 200 + assert await response.json() == { + "open": False, + "client_connected": False, + } + + window.open(_session_desc()) + peer, channel, _ = await _connect_browser(window) + channel.send(json.dumps({"type": "keyboard", "key": "w", "pressed": True})) + channel.send(json.dumps({"type": "keyboard", "key": "w", "pressed": False})) + + events = [] + for _ in range(100): + events.extend(window.get_user_input_events().get_events()) + if len(events) == 2: + break + await asyncio.sleep(0.01) + + assert len(events) == 2 + assert [ + (event.get_event_data().key, event.get_event_data().pressed) + for event in events + if isinstance(event.get_event_data(), KeyboardUserInputEventData) + ] == [("w", True), ("w", False)] + assert events[0].get_timestamp() <= events[1].get_timestamp() + assert window.get_user_input_events().get_events() == [] + finally: + if peer is not None: + await peer.close() + window.close() + + +@pytest.mark.asyncio +async def test_write_delivers_a_video_frame_to_the_browser() -> None: + window = WebRTCClientWindow() + peer: RTCPeerConnection | None = None + try: + window.open(_session_desc()) + peer, _, video_track = await _connect_browser(window) + track = await asyncio.wait_for(video_track, timeout=5) + + window.write( + StepResult( + step_index=0, + output=torch.full((2, 3, 16, 16), 17, dtype=torch.uint8), + frame_count=2, + output_layout=VideoTensorLayout.tchw, + metrics={}, + ) + ) + + frame = await asyncio.wait_for(track.recv(), timeout=5) + assert isinstance(frame, VideoFrame) + pixels = frame.to_ndarray(format="rgb24") + assert pixels.shape == (16, 16, 3) + assert abs(float(pixels.mean()) - 17.0) <= 2.0 + finally: + if peer is not None: + await peer.close() + window.close() From 49de20e142575b7f6ec850dcf1125c2bff201226 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 20 Aug 2026 00:55:58 +0000 Subject: [PATCH 2/6] Move WebRTC input buffering into the client window Keep transport callbacks internal while preserving the polling interface used by the session runner. --- .../flashdreams/api_v2/client_window.py | 2 +- .../runtime_v2/serving/webrtc_server.py | 39 ++++++++++++------- .../runtime_v2/webrtc_client_window.py | 16 +++++++- .../test_v2/test_webrtc_client_window.py | 2 +- 4 files changed, 42 insertions(+), 17 deletions(-) diff --git a/flashdreams/flashdreams/api_v2/client_window.py b/flashdreams/flashdreams/api_v2/client_window.py index 62b829245..8befe1c02 100644 --- a/flashdreams/flashdreams/api_v2/client_window.py +++ b/flashdreams/flashdreams/api_v2/client_window.py @@ -20,7 +20,7 @@ class IClientWindow(InputSource, OutputSink, ABC): is given that description in :meth:`OutputSink.open`. One thread makes every call on a window, so an implementation needs no - locking of its own. + locking except when its backend delivers input from another thread. Created by the runtime, never by an application. """ diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py index 9ff387a62..30717c166 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -10,7 +10,7 @@ import socket import threading import time -from collections import deque +from collections.abc import Callable from fractions import Fraction from typing import Any @@ -29,7 +29,6 @@ ResetUserInputEventData, UserInputEvent, ) -from flashdreams.runtime_v2.user_input_events import UserInputEvents from flashdreams.runtime_v2.video_tensor import VideoTensorLayout _BROWSER_PAGE = """ @@ -171,8 +170,7 @@ def __init__( self._host = host self._port = port self._startup_timeout_seconds = startup_timeout_seconds - self._events: deque[UserInputEvent] = deque() - self._events_lock = threading.Lock() + self._input_callback: Callable[[UserInputEvent], None] | None = None self._started = threading.Event() self._startup_error: BaseException | None = None self._loop: asyncio.AbstractEventLoop | None = None @@ -224,15 +222,25 @@ def open(self, session_desc: SessionDesc) -> None: raise RuntimeError("Cannot open a closed WebRTC server.") if self._session_desc is not None: raise RuntimeError("WebRTC server is already open.") + if self._input_callback is None: + raise RuntimeError("Register an input callback before opening WebRTC.") self._session_desc = session_desc self._session_start_ns = time.monotonic_ns() - def get_user_input_events(self) -> UserInputEvents: - """Drain and return browser events in timestamp order.""" - with self._events_lock: - events = list(self._events) - self._events.clear() - return UserInputEvents(events) + def register_input_callback( + self, callback: Callable[[UserInputEvent], None] + ) -> None: + """Register the function called for each received browser event. + + Args: + callback: Function that accepts one validated, timestamped event. + + Raises: + RuntimeError: A callback has already been registered. + """ + if self._input_callback is not None: + raise RuntimeError("An input callback is already registered.") + self._input_callback = callback def write(self, result: StepResult) -> None: """Deliver one generated result to the browser's video track. @@ -442,10 +450,13 @@ def _append_event( if session_start_ns is None: return timestamp_us = np.uint64((time.monotonic_ns() - session_start_ns) // 1_000) - with self._events_lock: - self._events.append( - UserInputEvent(timestamp=timestamp_us, event_data=event_data) - ) + event = UserInputEvent(timestamp=timestamp_us, event_data=event_data) + callback = self._input_callback + if callback is None: + raise RuntimeError("WebRTC input callback is not registered.") + # Pass that UserInputEvent to the callback. + # The callback stores it in WebRTCClientWindow’s thread-safe queue. + callback(event) def _record_client_disconnect(self) -> None: """Buffer one close event when the active browser disconnects.""" diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py index 668626e53..ea2185aa3 100644 --- a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -3,10 +3,13 @@ """WebRTC client window for the v2 runtime.""" +import queue + from flashdreams.api_v2.client_window import IClientWindow from flashdreams.runtime_v2.serving.webrtc_server import WebRTCServer from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult +from flashdreams.runtime_v2.user_input_event import UserInputEvent from flashdreams.runtime_v2.user_input_events import UserInputEvents @@ -26,12 +29,18 @@ def __init__( port: Listening port. Zero asks the operating system to choose one. startup_timeout_seconds: Maximum time to wait for server startup. """ + self._input_events: queue.SimpleQueue[UserInputEvent] = queue.SimpleQueue() self.server = WebRTCServer( host=host, port=port, startup_timeout_seconds=startup_timeout_seconds, ) + def handle_input(event: UserInputEvent) -> None: + self._input_events.put(event) + + self.server.register_input_callback(handle_input) + def open(self, session_desc: SessionDesc) -> None: """Configure WebRTC output for the session. @@ -42,7 +51,12 @@ def open(self, session_desc: SessionDesc) -> None: def get_user_input_events(self) -> UserInputEvents: """Drain and return buffered browser events in timestamp order.""" - return self.server.get_user_input_events() + events = [] + while True: + try: + events.append(self._input_events.get_nowait()) + except queue.Empty: + return UserInputEvents(events) def write(self, result: StepResult) -> None: """Deliver one generated result to the browser. diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index 149b17698..867a9868f 100644 --- a/flashdreams/test_v2/test_webrtc_client_window.py +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -79,7 +79,7 @@ def on_track(track: MediaStreamTrack) -> None: @pytest.mark.asyncio -async def test_server_buffers_browser_events_until_drained() -> None: +async def test_window_buffers_browser_events_until_drained() -> None: window = WebRTCClientWindow() peer: RTCPeerConnection | None = None try: From b2e5262bd5edef987ca688ceccb4d4c434791729 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 20 Aug 2026 04:57:50 +0000 Subject: [PATCH 3/6] Add a browser runner for the v2 red screen app Wire application lifecycle and WebRTC presentation end to end, including browser controls and responsive red-screen output. --- .../runtime_v2/application_runner.py | 44 ++++++ .../runtime_v2/serving/webrtc_server.py | 12 +- .../test_v2/test_application_runner.py | 136 ++++++++++++++++++ .../test_v2/test_webrtc_client_window.py | 17 ++- integrations_v2/red_screen/README.md | 32 ++++- integrations_v2/red_screen/pyproject.toml | 5 +- integrations_v2/red_screen/red_screen/app.py | 116 +++++++++++++-- .../red_screen/tests/test_red_screen.py | 32 ++++- .../red_screen/tests/test_webrtc.py | 32 +++++ uv.lock | 4 +- 10 files changed, 400 insertions(+), 30 deletions(-) create mode 100644 flashdreams/flashdreams/runtime_v2/application_runner.py create mode 100644 flashdreams/test_v2/test_application_runner.py create mode 100644 integrations_v2/red_screen/red_screen/tests/test_webrtc.py diff --git a/flashdreams/flashdreams/runtime_v2/application_runner.py b/flashdreams/flashdreams/runtime_v2/application_runner.py new file mode 100644 index 000000000..a985ee377 --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/application_runner.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Application lifecycle runner for the v2 runtime.""" + +from collections.abc import Sequence + +from flashdreams.api_v2.application import IApplication +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.session_desc import SessionDesc +from flashdreams.runtime_v2.session_runner import run_session + + +class ApplicationRunner: + """Create and run one application session against one client window.""" + + def __init__(self, application: IApplication, client_window: IClientWindow) -> None: + """ + Args: + application: Long-lived application that creates the session. + client_window: Window that supplies input and presents generated output. + """ + self._application = application + self._client_window = client_window + + def run( + self, + session_desc: SessionDesc, + commandline_args: Sequence[str] = (), + ) -> None: + """Initialize the application, create one session, and run it. + + The application is closed before this method returns or raises. + + Args: + session_desc: Output shape and timing requested for the session. + commandline_args: Arguments owned and parsed by the application. + """ + try: + self._application.init(commandline_args) + session = self._application.create_session(session_desc) + run_session(session, self._client_window) + finally: + self._application.close() diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py index 30717c166..6a7a0aa25 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -35,6 +35,7 @@ + + + diff --git a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py index 6a7a0aa25..68041879d 100644 --- a/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py +++ b/flashdreams/flashdreams/runtime_v2/serving/webrtc_server.py @@ -12,6 +12,7 @@ import time from collections.abc import Callable from fractions import Fraction +from importlib.resources import files from typing import Any import numpy as np @@ -31,67 +32,9 @@ ) from flashdreams.runtime_v2.video_tensor import VideoTensorLayout -_BROWSER_PAGE = """ - - - - - - - - -""" +_WEB_RESOURCES = files("flashdreams.runtime_v2.serving").joinpath("web") +_BROWSER_PAGE = _WEB_RESOURCES.joinpath("index.html").read_text(encoding="utf-8") +_BROWSER_SCRIPT = _WEB_RESOURCES.joinpath("app.js").read_text(encoding="utf-8") class _VideoTrack(MediaStreamTrack): @@ -312,6 +255,7 @@ async def _start_server(self) -> None: """Create and bind the standalone aiohttp application.""" app = web.Application() app.router.add_get("/", self._serve_browser) + app.router.add_get("/app.js", self._serve_browser_script) app.router.add_get("/healthz", self._health) app.router.add_post("/api/webrtc/offer", self._offer) runner = web.AppRunner(app) @@ -336,6 +280,10 @@ async def _serve_browser(self, _: web.Request) -> web.Response: """Return the minimal browser client.""" return web.Response(text=_BROWSER_PAGE, content_type="text/html") + async def _serve_browser_script(self, _: web.Request) -> web.Response: + """Return the browser client's JavaScript.""" + return web.Response(text=_BROWSER_SCRIPT, content_type="text/javascript") + async def _health(self, _: web.Request) -> web.Response: """Report whether the server has an open session and client.""" return web.json_response( diff --git a/flashdreams/pyproject.toml b/flashdreams/pyproject.toml index 5a0636e33..0206e6005 100644 --- a/flashdreams/pyproject.toml +++ b/flashdreams/pyproject.toml @@ -143,6 +143,7 @@ exclude = ["tests", "flashdreams._pytest_plugins*"] [tool.setuptools.package-data] "flashdreams.serving.webrtc" = ["web/*.html", "web/*.css", "web/*.js", "web/assets/*.svg"] +"flashdreams.runtime_v2.serving" = ["web/*.html", "web/*.js"] [dependency-groups] # Default CUDA 13 profile. No source binding -- Linux falls through to diff --git a/flashdreams/test_v2/test_webrtc_client_window.py b/flashdreams/test_v2/test_webrtc_client_window.py index 1e1fda4a9..b7f119651 100644 --- a/flashdreams/test_v2/test_webrtc_client_window.py +++ b/flashdreams/test_v2/test_webrtc_client_window.py @@ -94,7 +94,11 @@ async def test_window_buffers_browser_events_until_drained() -> None: browser_page = await response.text() assert response.status == 200 assert 'id="activate"' in browser_page - assert 'key: "r", pressed: activationPressed' in browser_page + assert '' in browser_page + async with client.get(f"{window.server.url}app.js") as response: + browser_script = await response.text() + assert response.status == 200 + assert 'key: "r", pressed: activationPressed' in browser_script window.open(_session_desc()) peer, channel, _ = await _connect_browser(window) From 907dcb6a87dc78acd7e1055a63352c3266828874 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 20 Aug 2026 05:14:00 +0000 Subject: [PATCH 5/6] Move client-window creation into the v2 runtime Centralize presentation mode selection so integrations reuse one runtime-owned factory. --- .../runtime_v2/client_window_factory.py | 27 +++++++++++++++++++ .../test_v2/test_client_window_factory.py | 5 ++-- integrations_v2/red_screen/red_screen/app.py | 20 +------------- 3 files changed, 30 insertions(+), 22 deletions(-) create mode 100644 flashdreams/flashdreams/runtime_v2/client_window_factory.py rename integrations_v2/red_screen/red_screen/tests/test_webrtc.py => flashdreams/test_v2/test_client_window_factory.py (86%) diff --git a/flashdreams/flashdreams/runtime_v2/client_window_factory.py b/flashdreams/flashdreams/runtime_v2/client_window_factory.py new file mode 100644 index 000000000..c5f2327bb --- /dev/null +++ b/flashdreams/flashdreams/runtime_v2/client_window_factory.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Create v2 client windows from runtime arguments.""" + +import argparse + +from flashdreams.api_v2.client_window import IClientWindow +from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow + + +def create_client_window(parsed_args: argparse.Namespace) -> IClientWindow: + """Create the client window selected by the presentation mode. + + Args: + parsed_args: Runtime arguments. Mode-specific fields are read only by + the selected mode. + + Returns: + Client window for the selected mode. + + Raises: + ValueError: ``mode`` is unsupported. + """ + if parsed_args.mode == "webrtc": + return WebRTCClientWindow(host=parsed_args.host, port=parsed_args.port) + raise ValueError(f"Unsupported client-window mode: {parsed_args.mode!r}.") diff --git a/integrations_v2/red_screen/red_screen/tests/test_webrtc.py b/flashdreams/test_v2/test_client_window_factory.py similarity index 86% rename from integrations_v2/red_screen/red_screen/tests/test_webrtc.py rename to flashdreams/test_v2/test_client_window_factory.py index eb3fb434c..ee05abe6a 100644 --- a/integrations_v2/red_screen/red_screen/tests/test_webrtc.py +++ b/flashdreams/test_v2/test_client_window_factory.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""CPU tests for the red-screen client-window factory.""" +"""CPU tests for the v2 client-window factory.""" import argparse @@ -12,8 +12,7 @@ pytest.importorskip("aiohttp") pytest.importorskip("aiortc") -from red_screen.app import create_client_window - +from flashdreams.runtime_v2.client_window_factory import create_client_window from flashdreams.runtime_v2.webrtc_client_window import WebRTCClientWindow diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py index 0a133114d..95975d559 100644 --- a/integrations_v2/red_screen/red_screen/app.py +++ b/integrations_v2/red_screen/red_screen/app.py @@ -11,9 +11,9 @@ from torch import Tensor from flashdreams.api_v2.application import IApplication -from flashdreams.api_v2.client_window import IClientWindow from flashdreams.api_v2.session import ISession from flashdreams.runtime_v2.application_runner import ApplicationRunner +from flashdreams.runtime_v2.client_window_factory import create_client_window from flashdreams.runtime_v2.session_desc import SessionDesc from flashdreams.runtime_v2.step_result import StepResult from flashdreams.runtime_v2.user_input_event import KeyboardUserInputEventData @@ -169,24 +169,6 @@ def create_app() -> IApplication: return RedScreenApplication() -def create_client_window(parsed_args: argparse.Namespace) -> IClientWindow: - """Create the client window selected by the presentation mode. - - Args: - parsed_args: Runtime arguments. Mode-specific fields are read only by - the selected mode. - - Returns: - Client window for the selected mode. - - Raises: - ValueError: ``mode`` is unsupported. - """ - if parsed_args.mode == "webrtc": - return WebRTCClientWindow(host=parsed_args.host, port=parsed_args.port) - raise ValueError(f"Unsupported client-window mode: {parsed_args.mode!r}.") - - def _parse_args(commandline_args: Sequence[str] | None) -> argparse.Namespace: """Parse runtime arguments and preserve application arguments.""" parser = argparse.ArgumentParser( From cc7ff22d45c7ad5f1b85e1a67d1d5fe15df57b28 Mon Sep 17 00:00:00 2001 From: Gangzheng Tong Date: Thu, 20 Aug 2026 05:17:10 +0000 Subject: [PATCH 6/6] Document WebRTC client-window protocol methods Make each implementation method's InputSource or OutputSink responsibility explicit. --- .../runtime_v2/webrtc_client_window.py | 21 +++++++++++++------ integrations_v2/red_screen/red_screen/app.py | 7 ++++++- 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py index ea2185aa3..28936d1fc 100644 --- a/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py +++ b/flashdreams/flashdreams/runtime_v2/webrtc_client_window.py @@ -14,7 +14,7 @@ class WebRTCClientWindow(IClientWindow): - """Adapt a standalone WebRTC server to the v2 client-window interface.""" + """Implement ``IClientWindow`` with WebRTC input and presentation.""" def __init__( self, @@ -23,7 +23,11 @@ def __init__( port: int = 0, startup_timeout_seconds: float = 10.0, ) -> None: - """ + """Create the WebRTC backend. + + Construction is specific to this implementation; it is not part of the + ``IClientWindow`` protocol. + Args: host: Interface on which the HTTP server listens. port: Listening port. Zero asks the operating system to choose one. @@ -37,12 +41,13 @@ def __init__( ) def handle_input(event: UserInputEvent) -> None: + """Buffer one backend event for the ``InputSource`` protocol.""" self._input_events.put(event) self.server.register_input_callback(handle_input) def open(self, session_desc: SessionDesc) -> None: - """Configure WebRTC output for the session. + """Implement ``OutputSink.open`` by configuring WebRTC output. Args: session_desc: Resolved dimensions, frame rate, and tensor layout. @@ -50,7 +55,11 @@ def open(self, session_desc: SessionDesc) -> None: self.server.open(session_desc) def get_user_input_events(self) -> UserInputEvents: - """Drain and return buffered browser events in timestamp order.""" + """Implement ``InputSource.get_user_input_events`` for browser input. + + Returns: + Buffered browser events in timestamp order, each returned once. + """ events = [] while True: try: @@ -59,7 +68,7 @@ def get_user_input_events(self) -> UserInputEvents: return UserInputEvents(events) def write(self, result: StepResult) -> None: - """Deliver one generated result to the browser. + """Implement ``OutputSink.write`` by delivering a result to the browser. Args: result: Generated frames matching the opened session. @@ -67,5 +76,5 @@ def write(self, result: StepResult) -> None: self.server.write(result) def close(self) -> None: - """Close the WebRTC connection and server.""" + """Implement ``OutputSink.close`` by releasing WebRTC resources.""" self.server.close() diff --git a/integrations_v2/red_screen/red_screen/app.py b/integrations_v2/red_screen/red_screen/app.py index 95975d559..c65ac9a58 100644 --- a/integrations_v2/red_screen/red_screen/app.py +++ b/integrations_v2/red_screen/red_screen/app.py @@ -85,7 +85,7 @@ def step(self, step_index: int, events: UserInputEvents) -> StepResult: import time # Simulate the real model inference time - time.sleep(1) + time.sleep(0.1) return StepResult( step_index=step_index, output=self._frame(), @@ -201,6 +201,11 @@ def main(commandline_args: Sequence[str] | None = None) -> int: if isinstance(window, WebRTCClientWindow): print(f"Open {window.server.url} in a browser.", flush=True) try: + # ApplicationRunner is a FlashDreams runtime component that takes an IApplication instance, a IClientWindow instance, + # and drives the main loop. + + # TODO: in production, commandline argument parsing and IClientWindow creation should be done by flashdreams-run, a CLI tool + # basically, we need to generailze this main function to be shared by all applications ApplicationRunner(app, window).run( SessionDesc( output_layout=VideoTensorLayout.bcthw,