wehub-resource-sync
377 lines
14 KiB
Python
377 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
from livekit import rtc
|
|
from livekit.browser import AudioData, BrowserPage, PaintData # type: ignore[import-untyped]
|
|
|
|
from ._keys import NATIVE_KEY_CODES as _NATIVE_KEY_CODES
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BrowserSession:
|
|
def __init__(self, *, page: BrowserPage, room: rtc.Room) -> None:
|
|
self._page = page
|
|
self._room = room
|
|
self._video_source: rtc.VideoSource | None = None
|
|
self._audio_source: rtc.AudioSource | None = None
|
|
self._video_track: rtc.LocalVideoTrack | None = None
|
|
self._audio_track: rtc.LocalAudioTrack | None = None
|
|
self._started = False
|
|
self._last_frame: rtc.VideoFrame | None = None
|
|
self._video_task: asyncio.Task[None] | None = None
|
|
self._audio_init_task: asyncio.Task[None] | None = None
|
|
self._audio_task: asyncio.Task[None] | None = None
|
|
self._audio_queue: asyncio.Queue[rtc.AudioFrame] | None = None
|
|
self._input_queue: asyncio.Queue[Any] = asyncio.Queue(maxsize=256)
|
|
self._input_task: asyncio.Task[None] | None = None
|
|
|
|
self._focus_identity: str | None = None
|
|
|
|
# Event set when a human interrupts the agent's focus
|
|
self._agent_interrupted = asyncio.Event()
|
|
|
|
@property
|
|
def focus_identity(self) -> str | None:
|
|
return self._focus_identity
|
|
|
|
@property
|
|
def agent_interrupted(self) -> asyncio.Event:
|
|
"""Event that is set when a human takes focus from the agent."""
|
|
return self._agent_interrupted
|
|
|
|
def set_agent_focus(self, active: bool) -> None:
|
|
"""Grant or revoke browser focus for the AI agent."""
|
|
if active:
|
|
self._focus_identity = "__agent__"
|
|
self._agent_interrupted.clear()
|
|
elif self._focus_identity == "__agent__":
|
|
self._focus_identity = None
|
|
|
|
async def reclaim_agent_focus(self) -> None:
|
|
"""Reclaim focus for the agent and notify all participants."""
|
|
self.set_agent_focus(True)
|
|
await self._page.send_focus_event(True)
|
|
await self._broadcast_focus()
|
|
|
|
async def start(self) -> None:
|
|
if self._started:
|
|
return
|
|
self._started = True
|
|
|
|
opts = self._page._opts
|
|
|
|
self._video_source = rtc.VideoSource(opts.width, opts.height, is_screencast=True)
|
|
|
|
self._video_track = rtc.LocalVideoTrack.create_video_track(
|
|
"browser-video", self._video_source
|
|
)
|
|
|
|
video_opts = rtc.TrackPublishOptions(
|
|
source=rtc.TrackSource.SOURCE_SCREENSHARE,
|
|
video_encoding=rtc.VideoEncoding(max_bitrate=8_000_000, max_framerate=opts.framerate),
|
|
simulcast=False,
|
|
)
|
|
self._page.on("paint", self._on_paint)
|
|
self._page.on("audio", self._on_audio)
|
|
self._page.on("cursor_changed", self._on_cursor)
|
|
self._page.on("url_changed", self._on_url_changed)
|
|
|
|
await self._room.local_participant.publish_track(self._video_track, video_opts)
|
|
|
|
self._video_task = asyncio.create_task(self._video_loop(opts.framerate))
|
|
|
|
# Single persistent task for sending input events to subprocess
|
|
self._input_task = asyncio.create_task(self._input_sender_loop())
|
|
|
|
# Register RPC methods for focus management
|
|
@self._room.local_participant.register_rpc_method("browser/request-focus") # type: ignore[arg-type]
|
|
async def _handle_request_focus(
|
|
data: rtc.rpc.RpcInvocationData,
|
|
) -> str:
|
|
if self._focus_identity is None:
|
|
self._focus_identity = data.caller_identity
|
|
await self._page.send_focus_event(True)
|
|
await self._broadcast_focus()
|
|
return json.dumps({"granted": True})
|
|
|
|
# If agent has focus, allow human to interrupt
|
|
if self._focus_identity == "__agent__":
|
|
self._focus_identity = data.caller_identity
|
|
self._agent_interrupted.set()
|
|
await self._page.send_focus_event(True)
|
|
await self._broadcast_focus()
|
|
return json.dumps({"granted": True})
|
|
|
|
return json.dumps({"granted": False, "holder": self._focus_identity})
|
|
|
|
@self._room.local_participant.register_rpc_method("browser/release-focus") # type: ignore[arg-type]
|
|
async def _handle_release_focus(
|
|
data: rtc.rpc.RpcInvocationData,
|
|
) -> str:
|
|
if self._focus_identity == data.caller_identity:
|
|
self._focus_identity = None
|
|
await self._page.send_focus_event(False)
|
|
await self._broadcast_focus()
|
|
return json.dumps({"released": True})
|
|
return json.dumps({"released": False})
|
|
|
|
@self._room.local_participant.register_rpc_method("browser/navigate") # type: ignore[arg-type]
|
|
async def _handle_navigate(
|
|
data: rtc.rpc.RpcInvocationData,
|
|
) -> str:
|
|
payload = json.loads(data.payload)
|
|
url = payload.get("url", "")
|
|
if url:
|
|
self._queue_input(self._page.navigate(url))
|
|
return json.dumps({"status": "ok"})
|
|
|
|
@self._room.local_participant.register_rpc_method("browser/go-back") # type: ignore[arg-type]
|
|
async def _handle_go_back(
|
|
data: rtc.rpc.RpcInvocationData,
|
|
) -> str:
|
|
self._queue_input(self._page.go_back())
|
|
return json.dumps({"status": "ok"})
|
|
|
|
@self._room.local_participant.register_rpc_method("browser/go-forward") # type: ignore[arg-type]
|
|
async def _handle_go_forward(
|
|
data: rtc.rpc.RpcInvocationData,
|
|
) -> str:
|
|
self._queue_input(self._page.go_forward())
|
|
return json.dumps({"status": "ok"})
|
|
|
|
# Listen for input data from participants
|
|
@self._room.on("data_received")
|
|
def _on_data_received(packet: rtc.DataPacket) -> None:
|
|
if packet.topic != "browser-input":
|
|
return
|
|
if packet.participant is None:
|
|
return
|
|
if packet.participant.identity != self._focus_identity:
|
|
return
|
|
try:
|
|
events = json.loads(packet.data)
|
|
except (json.JSONDecodeError, UnicodeDecodeError):
|
|
return
|
|
for evt in events:
|
|
self._dispatch_input(evt)
|
|
|
|
self._on_data_received = _on_data_received
|
|
|
|
# Release focus when the holder disconnects
|
|
@self._room.on("participant_disconnected")
|
|
def _on_participant_disconnected(participant: rtc.RemoteParticipant) -> None:
|
|
if participant.identity == self._focus_identity:
|
|
self._focus_identity = None
|
|
self._queue_input(self._page.send_focus_event(False))
|
|
self._queue_input(self._broadcast_focus())
|
|
|
|
self._on_participant_disconnected = _on_participant_disconnected
|
|
|
|
def _queue_input(self, coro: Any) -> None:
|
|
try:
|
|
self._input_queue.put_nowait(coro)
|
|
except asyncio.QueueFull:
|
|
pass
|
|
|
|
async def _input_sender_loop(self) -> None:
|
|
while True:
|
|
coro = await self._input_queue.get()
|
|
try:
|
|
await coro
|
|
except Exception:
|
|
pass
|
|
|
|
def _dispatch_input(self, evt: dict[str, Any]) -> None:
|
|
t = evt.get("type")
|
|
if t == "mousemove":
|
|
self._queue_input(self._page.send_mouse_move(evt["x"], evt["y"]))
|
|
elif t == "mousedown":
|
|
self._queue_input(
|
|
self._page.send_mouse_click(evt["x"], evt["y"], evt.get("button", 0), False, 1)
|
|
)
|
|
elif t == "mouseup":
|
|
self._queue_input(
|
|
self._page.send_mouse_click(evt["x"], evt["y"], evt.get("button", 0), True, 1)
|
|
)
|
|
elif t == "wheel":
|
|
self._queue_input(
|
|
self._page.send_mouse_wheel(
|
|
evt["x"], evt["y"], evt.get("deltaX", 0), evt.get("deltaY", 0)
|
|
)
|
|
)
|
|
elif t == "keydown":
|
|
wkc = evt["keyCode"]
|
|
nkc = _NATIVE_KEY_CODES.get(wkc, 0)
|
|
self._queue_input(self._page.send_key_event(0, evt.get("modifiers", 0), wkc, nkc, 0))
|
|
elif t == "keyup":
|
|
wkc = evt["keyCode"]
|
|
nkc = _NATIVE_KEY_CODES.get(wkc, 0)
|
|
self._queue_input(self._page.send_key_event(2, evt.get("modifiers", 0), wkc, 0, 0))
|
|
elif t == "char":
|
|
wkc = evt["keyCode"]
|
|
nkc = _NATIVE_KEY_CODES.get(wkc, 0)
|
|
self._queue_input(
|
|
self._page.send_key_event(
|
|
3, evt.get("modifiers", 0), wkc, nkc, evt.get("charCode", 0)
|
|
)
|
|
)
|
|
|
|
async def _broadcast_focus(self) -> None:
|
|
payload = json.dumps({"identity": self._focus_identity}).encode()
|
|
await self._room.local_participant.publish_data(
|
|
payload, reliable=True, topic="browser-focus"
|
|
)
|
|
|
|
async def _init_audio(self, frame: rtc.AudioFrame) -> None:
|
|
try:
|
|
self._audio_source = rtc.AudioSource(
|
|
frame.sample_rate, frame.num_channels, queue_size_ms=100
|
|
)
|
|
self._audio_track = rtc.LocalAudioTrack.create_audio_track(
|
|
"browser-audio", self._audio_source
|
|
)
|
|
audio_opts = rtc.TrackPublishOptions(
|
|
source=rtc.TrackSource.SOURCE_SCREENSHARE_AUDIO,
|
|
)
|
|
await self._room.local_participant.publish_track(self._audio_track, audio_opts)
|
|
self._audio_queue = asyncio.Queue(maxsize=50)
|
|
self._audio_queue.put_nowait(frame)
|
|
self._audio_task = asyncio.create_task(self._audio_loop())
|
|
except Exception:
|
|
logger.exception("failed to initialize audio, will retry on next packet")
|
|
self._audio_source = None
|
|
self._audio_track = None
|
|
self._audio_init_task = None
|
|
|
|
async def _audio_loop(self) -> None:
|
|
assert self._audio_queue is not None
|
|
assert self._audio_source is not None
|
|
while True:
|
|
frame = await self._audio_queue.get()
|
|
try:
|
|
await self._audio_source.capture_frame(frame)
|
|
except Exception:
|
|
logger.exception("audio capture error")
|
|
|
|
def _on_audio(self, data: AudioData) -> None:
|
|
if self._audio_source is None:
|
|
if self._audio_init_task is not None:
|
|
return
|
|
self._audio_init_task = asyncio.get_event_loop().create_task(
|
|
self._init_audio(data.frame)
|
|
)
|
|
return
|
|
|
|
if self._audio_queue is None:
|
|
return
|
|
|
|
try:
|
|
self._audio_queue.put_nowait(data.frame)
|
|
except asyncio.QueueFull:
|
|
pass
|
|
|
|
_CEF_CURSOR_MAP: dict[int, str] = {
|
|
0: "default", # CT_POINTER
|
|
1: "crosshair", # CT_CROSS
|
|
2: "pointer", # CT_HAND
|
|
3: "text", # CT_IBEAM
|
|
4: "wait", # CT_WAIT
|
|
5: "help", # CT_HELP
|
|
6: "ew-resize", # CT_EASTRESIZE
|
|
7: "ns-resize", # CT_NORTHRESIZE
|
|
8: "nesw-resize", # CT_NORTHEASTRESIZE
|
|
9: "nwse-resize", # CT_NORTHWESTRESIZE
|
|
10: "ns-resize", # CT_SOUTHRESIZE
|
|
11: "nwse-resize", # CT_SOUTHEASTRESIZE
|
|
12: "nesw-resize", # CT_SOUTHWESTRESIZE
|
|
13: "ew-resize", # CT_WESTRESIZE
|
|
14: "ns-resize", # CT_NORTHSOUTHRESIZE
|
|
15: "ew-resize", # CT_EASTWESTRESIZE
|
|
16: "nesw-resize", # CT_NORTHEASTSOUTHWESTRESIZE
|
|
17: "nwse-resize", # CT_NORTHWESTSOUTHEASTRESIZE
|
|
18: "ew-resize", # CT_COLUMNRESIZE
|
|
19: "ns-resize", # CT_ROWRESIZE
|
|
20: "move", # CT_MIDDLEPANNING
|
|
28: "not-allowed", # CT_NOTALLOWED
|
|
29: "grab", # CT_GRAB
|
|
30: "grabbing", # CT_GRABBING
|
|
32: "move", # CT_MOVE
|
|
}
|
|
|
|
def _on_url_changed(self, url: str) -> None:
|
|
payload = json.dumps({"url": url}).encode()
|
|
self._queue_input(
|
|
self._room.local_participant.publish_data(payload, reliable=True, topic="browser-url")
|
|
)
|
|
|
|
def _on_cursor(self, cursor_type: int) -> None:
|
|
css_cursor = self._CEF_CURSOR_MAP.get(cursor_type, "default")
|
|
payload = json.dumps({"cursor": css_cursor}).encode()
|
|
self._queue_input(
|
|
self._room.local_participant.publish_data(
|
|
payload, reliable=True, topic="cursor_changed"
|
|
)
|
|
)
|
|
|
|
def _on_paint(self, data: PaintData) -> None:
|
|
self._last_frame = data.frame
|
|
|
|
async def _video_loop(self, fps: float) -> None:
|
|
interval = 1.0 / fps
|
|
loop = asyncio.get_event_loop()
|
|
next_time = loop.time()
|
|
while True:
|
|
if self._last_frame is not None and self._video_source is not None:
|
|
ts_us = time.monotonic_ns() // 1000
|
|
self._video_source.capture_frame(self._last_frame, timestamp_us=ts_us)
|
|
next_time += interval
|
|
delay = next_time - loop.time()
|
|
if delay > 0:
|
|
await asyncio.sleep(delay)
|
|
else:
|
|
next_time = loop.time()
|
|
|
|
async def aclose(self) -> None:
|
|
if not self._started:
|
|
return
|
|
|
|
self._page.off("paint", self._on_paint)
|
|
self._page.off("audio", self._on_audio)
|
|
self._page.off("cursor_changed", self._on_cursor)
|
|
self._page.off("url_changed", self._on_url_changed)
|
|
|
|
if hasattr(self, "_on_data_received"):
|
|
self._room.off("data_received", self._on_data_received)
|
|
if hasattr(self, "_on_participant_disconnected"):
|
|
self._room.off("participant_disconnected", self._on_participant_disconnected)
|
|
|
|
if self._video_task:
|
|
self._video_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await self._video_task
|
|
|
|
if self._audio_task:
|
|
self._audio_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await self._audio_task
|
|
|
|
if self._input_task:
|
|
self._input_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await self._input_task
|
|
|
|
try:
|
|
if self._video_track:
|
|
await self._room.local_participant.unpublish_track(self._video_track.sid)
|
|
if self._audio_track:
|
|
await self._room.local_participant.unpublish_track(self._audio_track.sid)
|
|
except Exception:
|
|
pass # tracks may already be gone if room disconnected first
|