"""Ordered Firefox Remote Debugging Protocol connection."""
from __future__ import annotations
import asyncio
from collections import deque
from dataclasses import dataclass
import json
import logging
from typing import Any
from .errors import (
ConnectionLostError,
PageScriptError,
RdpProtocolError,
RdpTimeoutError,
)
from .websocket_client import WebSocketClient
LOGGER = logging.getLogger("firefox_mcp.rdp")
def _utf16Length(value: str) -> int:
return len(value.encode("utf-16-le")) // 2
@dataclass(frozen=True)
class TabRecord:
"""Describe an ordinary top-level Firefox tab actor."""
actor_id: str
title: str
url: str
console_actor_id: str | None = None
@dataclass
class PendingEvaluation:
"""Track the asynchronous completion phase of an evaluation."""
actor_id: str
result_id: str
future: asyncio.Future[dict[str, object]]
class RdpConnection:
"""Serialize RDP actor requests and route asynchronous evaluations."""
def __init__(
self,
host: str,
port: int,
timeout_ms: int,
*,
max_result_chars: int = 64 * 1024 * 1024,
) -> None:
self.host = host
self.port = port
self.timeout_ms = timeout_ms
self.max_result_chars = max_result_chars
self._websocket = WebSocketClient(host, port)
self._request_lock = asyncio.Lock()
self._connection_lock = asyncio.Lock()
self._direct_future: asyncio.Future[dict[str, object]] | None = None
self._direct_actor: str | None = None
self._evaluations: dict[str, PendingEvaluation] = {}
self._early_evaluations: dict[str, dict[str, object]] = {}
self._receiver_task: asyncio.Task[None] | None = None
self._greeting: dict[str, object] | None = None
self._generation = 0
self._target_generation = 0
self.unmatched_packets: deque[dict[str, object]] = deque(
maxlen=100
)
@property
def connected(self) -> bool:
"""Return whether the WebSocket and receiver task are healthy."""
return (
self._websocket.connected
and self._receiver_task is not None
and not self._receiver_task.done()
)
@property
def generation(self) -> int:
"""Return the connection recreation counter."""
return self._generation
@property
def greeting(self) -> dict[str, object] | None:
"""Return the last validated root greeting."""
return self._greeting
@property
def target_generation(self) -> int:
"""Return the target-destruction notification counter."""
return self._target_generation
async def connect(self) -> None:
"""Connect, receive the root greeting, and start packet routing."""
async with self._connection_lock:
if self.connected:
return
await self._closeUnlocked()
timeout = self.timeout_ms / 1000
await self._websocket.connect(timeout)
try:
raw = await asyncio.wait_for(
self._websocket.receiveText(),
timeout,
)
greeting = self._decodePacket(raw)
except asyncio.TimeoutError as error:
await self._websocket.close()
raise RdpTimeoutError(
"Timed out waiting for the Firefox root greeting."
) from error
if greeting.get("from") != "root":
await self._websocket.close()
raise RdpProtocolError(
"Firefox RDP root greeting has an unexpected actor."
)
application_type = greeting.get("applicationType")
if application_type not in {None, "browser"}:
await self._websocket.close()
raise RdpProtocolError(
"Firefox RDP endpoint is not a browser endpoint."
)
self._greeting = greeting
self._generation += 1
self._target_generation += 1
self._receiver_task = asyncio.create_task(
self._receiveLoop(),
name="firefox-rdp-receiver",
)
async def close(self) -> None:
"""Close the RDP connection and fail all pending operations."""
async with self._connection_lock:
await self._closeUnlocked()
async def _closeUnlocked(self) -> None:
task = self._receiver_task
self._receiver_task = None
if task is not None and task is not asyncio.current_task():
task.cancel()
try:
await task
except (asyncio.CancelledError, ConnectionLostError):
pass
await self._websocket.close()
self._greeting = None
self._failPending(
ConnectionLostError("Firefox RDP connection closed.")
)
@staticmethod
def _decodePacket(raw: str) -> dict[str, object]:
try:
packet = json.loads(raw)
except json.JSONDecodeError as error:
raise RdpProtocolError(
"Firefox RDP sent invalid JSON."
) from error
if not isinstance(packet, dict):
raise RdpProtocolError(
"Firefox RDP packet must be a JSON object."
)
return packet
async def _receiveLoop(self) -> None:
try:
while True:
raw = await self._websocket.receiveText()
packet = self._decodePacket(raw)
self._routePacket(packet)
except asyncio.CancelledError:
raise
except Exception as error:
if isinstance(error, ConnectionLostError):
failure = error
else:
failure = ConnectionLostError(
"Firefox RDP receiver stopped unexpectedly."
)
failure.__cause__ = error
self._failPending(failure)
def _routePacket(self, packet: dict[str, object]) -> None:
result_id = packet.get("resultID")
if isinstance(result_id, str) and result_id in self._evaluations:
pending = self._evaluations.pop(result_id)
if not pending.future.done():
pending.future.set_result(packet)
return
if (
isinstance(result_id, str)
and packet.get("type") == "evaluationResult"
):
if len(self._early_evaluations) >= 100:
oldest = next(iter(self._early_evaluations))
self._early_evaluations.pop(oldest)
self._early_evaluations[result_id] = packet
return
actor = packet.get("from")
if packet.get("type") in {
"targetDestroyed",
"targetDestroyedForm",
}:
self._target_generation += 1
if (
self._direct_future is not None
and not self._direct_future.done()
and actor == self._direct_actor
and "type" not in packet
):
self._direct_future.set_result(packet)
return
self.unmatched_packets.append(self._packetSummary(packet))
@staticmethod
def _packetSummary(packet: dict[str, object]) -> dict[str, object]:
keys = ("to", "from", "type", "resultID", "error")
return {key: packet[key] for key in keys if key in packet}
def _failPending(self, error: Exception) -> None:
if (
self._direct_future is not None
and not self._direct_future.done()
):
self._direct_future.set_exception(error)
for pending in self._evaluations.values():
if not pending.future.done():
pending.future.set_exception(error)
self._evaluations.clear()
self._early_evaluations.clear()
async def request(
self,
actor_id: str,
request_type: str,
fields: dict[str, object] | None = None,
*,
timeout_ms: int | None = None,
) -> dict[str, object]:
"""Send one actor request and await its ordered direct reply."""
if not self.connected:
raise ConnectionLostError("Firefox RDP is not connected.")
timeout = (timeout_ms or self.timeout_ms) / 1000
async with self._request_lock:
loop = asyncio.get_running_loop()
future: asyncio.Future[dict[str, object]] = loop.create_future()
self._direct_future = future
self._direct_actor = actor_id
packet = {"to": actor_id, "type": request_type}
if fields:
packet.update(fields)
try:
await self._websocket.sendText(
json.dumps(packet, separators=(",", ":"))
)
response = await asyncio.wait_for(future, timeout)
except asyncio.TimeoutError as error:
future.cancel()
await self.close()
raise RdpTimeoutError(
f"Firefox RDP request timed out: {request_type}"
) from error
finally:
self._direct_future = None
self._direct_actor = None
self._raiseActorError(response)
return response
@staticmethod
def _raiseActorError(packet: dict[str, object]) -> None:
error_name = packet.get("error")
if error_name is not None:
message = packet.get("message")
raise RdpProtocolError(
f"Firefox RDP actor error {error_name}: {message or ''}"
)
async def listTabs(self) -> list[TabRecord]:
"""List current top-level tab descriptors from the root actor."""
packet = await self.request("root", "listTabs")
tabs = packet.get("tabs")
if not isinstance(tabs, list):
raise RdpProtocolError("listTabs reply has no tabs array.")
result: list[TabRecord] = []
for item in tabs:
if not isinstance(item, dict):
continue
actor = item.get("actor")
if not isinstance(actor, str):
actor = item.get("actorID")
if not isinstance(actor, str):
continue
result.append(TabRecord(
actor_id=actor,
title=str(item.get("title", "")),
url=str(item.get("url", "")),
))
return result
async def getConsoleActor(self, tab_actor_id: str) -> str:
"""Resolve a tab descriptor to its target console actor."""
packet = await self.request(tab_actor_id, "getTarget")
candidates: list[object] = [packet]
for name in ("frame", "target", "form"):
value = packet.get(name)
if isinstance(value, dict):
candidates.append(value)
for candidate in candidates:
if not isinstance(candidate, dict):
continue
actor = candidate.get("consoleActor")
if isinstance(actor, str):
return actor
raise RdpProtocolError(
"Firefox target reply does not contain a console actor."
)
async def evaluate(
self,
console_actor_id: str,
expression: str,
*,
timeout_ms: int | None = None,
) -> object:
"""Evaluate JavaScript and decode the resulting RDP value grip."""
timeout = timeout_ms or self.timeout_ms
async with self._request_lock:
if not self.connected:
raise ConnectionLostError("Firefox RDP is not connected.")
loop = asyncio.get_running_loop()
direct: asyncio.Future[dict[str, object]] = loop.create_future()
self._direct_future = direct
self._direct_actor = console_actor_id
request = {
"to": console_actor_id,
"type": "evaluateJSAsync",
"text": expression,
"mapped": {"await": True},
}
try:
await self._websocket.sendText(
json.dumps(request, separators=(",", ":"))
)
acknowledgement = await asyncio.wait_for(
direct,
timeout / 1000,
)
except asyncio.TimeoutError as error:
direct.cancel()
await self.close()
raise RdpTimeoutError(
"Firefox did not acknowledge page evaluation."
) from error
finally:
self._direct_future = None
self._direct_actor = None
self._raiseActorError(acknowledgement)
result_id = acknowledgement.get("resultID")
if not isinstance(result_id, str):
raise RdpProtocolError(
"evaluateJSAsync reply has no resultID."
)
result_future: asyncio.Future[dict[str, object]] = (
loop.create_future()
)
self._evaluations[result_id] = PendingEvaluation(
console_actor_id,
result_id,
result_future,
)
early = self._early_evaluations.pop(result_id, None)
if early is not None:
self._evaluations.pop(result_id, None)
result_future.set_result(early)
try:
completion = await asyncio.wait_for(
result_future,
timeout / 1000,
)
except asyncio.TimeoutError as error:
self._evaluations.pop(result_id, None)
await self.close()
raise RdpTimeoutError(
"Firefox did not complete page evaluation."
) from error
self._raiseActorError(completion)
if completion.get("exception") is not None:
LOGGER.debug(
"Page evaluation exception: %r",
completion.get("exception"),
)
raise PageScriptError("Page JavaScript evaluation failed.")
result = completion.get("result")
if isinstance(result, dict) and "result" in result:
if result.get("exception") is not None:
LOGGER.debug(
"Page evaluation exception: %r",
result.get("exception"),
)
raise PageScriptError("Page JavaScript evaluation failed.")
result = result.get("result")
return await self._readGrip(result)
async def _readGrip(self, grip: object) -> object:
if not isinstance(grip, dict):
return grip
grip_type = grip.get("type")
if grip_type == "longString":
return await self._readLongString(grip)
if grip_type == "undefined":
return None
if grip_type in {"NaN", "Infinity", "-Infinity", "-0"}:
return str(grip_type)
if "value" in grip:
return grip["value"]
return grip
async def _readLongString(self, grip: dict[str, object]) -> str:
actor = grip.get("actor")
initial = grip.get("initial")
length = grip.get("length")
if (
not isinstance(actor, str)
or not isinstance(initial, str)
or not isinstance(length, int)
):
raise RdpProtocolError("Malformed Firefox long-string grip.")
if length > self.max_result_chars:
raise RdpProtocolError("Firefox result exceeds size limit.")
pieces = [initial]
offset = _utf16Length(initial)
try:
while offset < length:
end = min(offset + 1_048_576, length)
response = await self.request(
actor,
"substring",
{"start": offset, "end": end},
)
substring = response.get("substring")
if not isinstance(substring, str) or not substring:
raise RdpProtocolError(
"Firefox long-string actor returned invalid data."
)
pieces.append(substring)
offset += _utf16Length(substring)
return "".join(pieces)
finally:
try:
await self.request(actor, "release")
except Exception:
pass
async def evaluateJson(
self,
console_actor_id: str,
expression: str,
*,
timeout_ms: int | None = None,
) -> dict[str, object]:
"""Evaluate a page script whose successful result is JSON text."""
raw = await self.evaluate(
console_actor_id,
expression,
timeout_ms=timeout_ms,
)
if not isinstance(raw, str):
raise PageScriptError("Page script did not return JSON text.")
try:
value = json.loads(raw)
except json.JSONDecodeError as error:
raise PageScriptError(
"Page script returned invalid JSON text."
) from error
if not isinstance(value, dict):
raise PageScriptError("Page script result must be an object.")
return value