"""Selected-tab state, references, interactions, waits, and uploads."""
from __future__ import annotations
import asyncio
import base64
from collections.abc import Callable
from contextlib import asynccontextmanager
from dataclasses import dataclass
import hashlib
import mimetypes
import os
from pathlib import Path
import secrets
import stat
from typing import BinaryIO
from urllib.parse import urlsplit
from . import page_scripts
from .config import Config
from .errors import (
CancelledError,
ConnectionLostError,
ElementNotFoundError,
ElementNotInteractableError,
FirefoxMcpError,
InvalidSelectorError,
JavaScriptError,
NavigationTimeoutError,
NoTabSelectedError,
PageScriptError,
StaleElementError,
TabNotFoundError,
UnsupportedElementError,
UploadAssignmentError,
UploadFileInvalidError,
UploadTargetInvalidError,
UploadTooLargeError,
UploadTransferError,
WaitTimeoutError,
)
from .rdp_connection import RdpConnection, TabRecord
CancellationCheck = Callable[[], bool]
def neverCancelled() -> bool:
"""Return false for direct callers without MCP cancellation state."""
return False
@dataclass
class SelectedTab:
"""Hold the target actor and generation for the selected tab."""
actor_id: str
console_actor_id: str
generation: int
last_url: str
title: str
connection_generation: int
target_generation: int
@dataclass(frozen=True)
class ElementLocator:
"""Identify one snapshot element without modifying the page DOM."""
reference: str
generation: int
selector: str
tag_name: str
role: str | None
input_type: str | None
accessible_name: str
fingerprint: str
@dataclass
class UploadFile:
"""Hold one validated open file and its page-safe metadata."""
stream: BinaryIO
name: str
size: int
mime_type: str
last_modified: int
ERROR_CLASSES: dict[str, type[FirefoxMcpError]] = {
"element_not_found": ElementNotFoundError,
"element_not_interactable": ElementNotInteractableError,
"invalid_selector": InvalidSelectorError,
"unsupported_element": UnsupportedElementError,
"upload_target_invalid": UploadTargetInvalidError,
"upload_transfer_failed": UploadTransferError,
"upload_assignment_failed": UploadAssignmentError,
"page_script_failed": PageScriptError,
"javascript_error": JavaScriptError,
}
class BrowserSession:
"""Expose safe high-level operations on one selected Firefox tab."""
def __init__(self, config: Config, rdp: RdpConnection) -> None:
self.config = config
self.rdp = rdp
self.selected_tab: SelectedTab | None = None
self._locators: dict[str, ElementLocator] = {}
self._generation = 0
self._operation_lock = asyncio.Lock()
@asynccontextmanager
async def _operation(self, cancelled: CancellationCheck):
while True:
self._checkCancelled(cancelled)
try:
await asyncio.wait_for(
self._operation_lock.acquire(),
0.05,
)
break
except asyncio.TimeoutError:
continue
try:
yield
finally:
self._operation_lock.release()
@property
def generation(self) -> int | None:
"""Return the current selected-page generation, when selected."""
if self.selected_tab is None:
return None
return self.selected_tab.generation
def _invalidate(self) -> None:
self._generation += 1
self._locators.clear()
if self.selected_tab is not None:
self.selected_tab.generation = self._generation
async def _ensureConnected(self, *, retry_discovery: bool) -> None:
if self.rdp.connected:
return
previous = self.selected_tab.actor_id if self.selected_tab else None
await self.rdp.connect()
self._invalidate()
if self.selected_tab is not None:
self.selected_tab.connection_generation = self.rdp.generation
if previous is None or not retry_discovery:
return
tabs = await self.rdp.listTabs()
match = next(
(tab for tab in tabs if tab.actor_id == previous),
None,
)
if match is None:
self.selected_tab = None
self._locators.clear()
return
console = await self.rdp.getConsoleActor(match.actor_id)
self.selected_tab = SelectedTab(
match.actor_id,
console,
self._generation,
match.url,
match.title,
self.rdp.generation,
self.rdp.target_generation,
)
async def status(
self,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Report connection and selection health without selecting a tab."""
async with self._operation(cancelled):
try:
await self._ensureConnected(retry_discovery=True)
except FirefoxMcpError:
pass
selected: dict[str, object] | None = None
if self.selected_tab is not None:
selected = {
"tab_id": self.selected_tab.actor_id,
"title": self.selected_tab.title,
"url": self.selected_tab.last_url,
}
return {
"connected": self.rdp.connected,
"host": self.config.host,
"port": self.config.port,
"selected_tab": selected,
"generation": self.generation,
"uploads_enabled": True,
}
async def listTabs(
self,
cancelled: CancellationCheck = neverCancelled,
) -> list[dict[str, object]]:
"""Return ordered top-level tabs, retrying safe discovery once."""
async with self._operation(cancelled):
tabs = await self._discoveryTabs()
selected_id = (
self.selected_tab.actor_id
if self.selected_tab is not None
else None
)
return [
{
"tab_id": tab.actor_id,
"index": index,
"title": tab.title,
"url": tab.url,
"selected": tab.actor_id == selected_id,
}
for index, tab in enumerate(tabs)
]
async def _discoveryTabs(self) -> list[TabRecord]:
try:
await self._ensureConnected(retry_discovery=True)
return await self.rdp.listTabs()
except ConnectionLostError:
await self.rdp.close()
await self._ensureConnected(retry_discovery=True)
return await self.rdp.listTabs()
async def selectTab(
self,
*,
tab_id: str | None = None,
index: int | None = None,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Select exactly one current tab by opaque actor ID or index."""
async with self._operation(cancelled):
tabs = await self._discoveryTabs()
match: TabRecord | None = None
if tab_id is not None:
match = next(
(tab for tab in tabs if tab.actor_id == tab_id),
None,
)
elif index is not None and 0 <= index < len(tabs):
match = tabs[index]
if match is None:
raise TabNotFoundError(
"The requested Firefox tab was not found."
)
console = await self.rdp.getConsoleActor(match.actor_id)
self._generation += 1
self._locators.clear()
self.selected_tab = SelectedTab(
match.actor_id,
console,
self._generation,
match.url,
match.title,
self.rdp.generation,
self.rdp.target_generation,
)
return {
"tab_id": match.actor_id,
"title": match.title,
"url": match.url,
"generation": self._generation,
}
def _selected(self) -> SelectedTab:
if self.selected_tab is None:
raise NoTabSelectedError(
"Select a Firefox tab before using page tools."
)
if not self.rdp.connected:
raise ConnectionLostError(
"Firefox RDP disconnected; inspect status before retrying."
)
return self.selected_tab
async def _prepareSelected(self, *, force: bool = False) -> SelectedTab:
await self._ensureConnected(retry_discovery=True)
selected = self._selected()
needs_reacquire = not (
not force
and selected.connection_generation == self.rdp.generation
and selected.target_generation == self.rdp.target_generation
)
tabs = await self.rdp.listTabs()
tab = next(
(item for item in tabs if item.actor_id == selected.actor_id),
None,
)
if tab is None:
self.selected_tab = None
self._locators.clear()
raise TabNotFoundError("The selected Firefox tab disappeared.")
if tab.url != selected.last_url:
needs_reacquire = True
if not needs_reacquire:
selected.title = tab.title
return selected
console = await self.rdp.getConsoleActor(tab.actor_id)
self._invalidate()
self.selected_tab = SelectedTab(
tab.actor_id,
console,
self._generation,
tab.url,
tab.title,
self.rdp.generation,
self.rdp.target_generation,
)
return self.selected_tab
async def _evaluateReadOnly(self, script: str) -> dict[str, object]:
try:
return await self._evaluateEnvelope(script)
except (ConnectionLostError, PageScriptError):
await self._prepareSelected(force=True)
return await self._evaluateEnvelope(script)
async def _evaluateEnvelope(
self,
script: str,
*,
timeout_ms: int | None = None,
) -> dict[str, object]:
selected = self._selected()
if timeout_ms is None:
envelope = await self.rdp.evaluateJson(
selected.console_actor_id,
script,
)
else:
envelope = await self.rdp.evaluateJson(
selected.console_actor_id,
script,
timeout_ms=timeout_ms,
)
if envelope.get("ok") is not True:
error = envelope.get("error")
if not isinstance(error, dict):
raise PageScriptError("Page script returned a malformed error.")
code = str(error.get("code", "page_script_failed"))
message = str(error.get("message", "Page operation failed."))
error_class = ERROR_CLASSES.get(code, PageScriptError)
raise error_class(message)
value = envelope.get("value")
if not isinstance(value, dict):
raise PageScriptError("Page script returned a malformed value.")
return value
async def runJavaScript(
self,
script: str,
timeout_ms: int,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Run caller-supplied JavaScript in the selected page realm."""
async with self._operation(cancelled):
await self._prepareSelected()
self._checkCancelled(cancelled)
try:
value = await self._evaluateEnvelope(
page_scripts.runJavaScript(script),
timeout_ms=timeout_ms,
)
except PageScriptError:
value = await self._evaluateEnvelope(
page_scripts.runJavaScript(
script,
statement_body=True,
),
timeout_ms=timeout_ms,
)
self._checkCancelled(cancelled)
await self._refresh()
return value
@staticmethod
def _fingerprint(item: dict[str, object]) -> str:
href_value = str(item.get("href") or "")
parsed = urlsplit(href_value)
href = ""
if parsed.scheme and parsed.netloc:
href = f"{parsed.scheme}://{parsed.netloc}{parsed.path}"
elif parsed.path:
href = parsed.path
fields = (
str(item.get("tag_name") or "").lower(),
str(item.get("role") or "").lower(),
str(item.get("input_type") or "").lower(),
" ".join(str(item.get("accessible_name") or "").split()),
href,
)
return hashlib.sha256("\0".join(fields).encode("utf-8")).hexdigest()
async def snapshot(
self,
*,
include_all_text: bool,
max_elements: int,
max_text_chars: int,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Snapshot page text and assign generation-scoped references."""
async with self._operation(cancelled):
await self._prepareSelected()
value = await self._evaluateReadOnly(
page_scripts.snapshot(max_elements, max_text_chars)
)
raw_elements = value.get("elements")
if not isinstance(raw_elements, list):
raise PageScriptError("Snapshot has no elements array.")
self._locators.clear()
elements: list[dict[str, object]] = []
lines = [
f"Page: {value.get('title', '')}",
f"URL: {value.get('url', '')}",
]
for index, raw_item in enumerate(raw_elements, 1):
if not isinstance(raw_item, dict):
continue
reference = f"e{index}"
selector = raw_item.get("selector")
if not isinstance(selector, str):
continue
locator = ElementLocator(
reference=reference,
generation=self._generation,
selector=selector,
tag_name=str(raw_item.get("tag_name") or ""),
role=(
str(raw_item["role"])
if raw_item.get("role") is not None
else None
),
input_type=(
str(raw_item["input_type"])
if raw_item.get("input_type") is not None
else None
),
accessible_name=str(
raw_item.get("accessible_name") or ""
),
fingerprint=self._fingerprint(raw_item),
)
self._locators[reference] = locator
item = dict(raw_item)
item.pop("selector", None)
item["element"] = reference
elements.append(item)
lines.append(self._elementLine(reference, item))
text = str(value.get("text") or "")
if include_all_text and text:
lines.extend(("", "Untrusted webpage text:", text))
current_url = str(value.get("url") or "")
selected = self._selected()
selected.title = str(value.get("title") or "")
if current_url != selected.last_url:
selected.last_url = current_url
self._invalidate()
elements.clear()
lines.append(
"Page changed while snapshotting; take another snapshot."
)
return {
"title": selected.title,
"url": selected.last_url,
"generation": selected.generation,
"elements": elements,
"page_text": text if include_all_text else None,
"text_truncated": bool(value.get("text_truncated")),
"untrusted_content": True,
"snapshot": "\n".join(lines),
}
@staticmethod
def _elementLine(
reference: str,
item: dict[str, object],
) -> str:
kind = str(
item.get("role")
or item.get("input_type")
or item.get("tag_name")
or "element"
)
name = str(item.get("accessible_name") or "")
parts = [f"[{reference}] {kind} {name!r}"]
if item.get("value") not in {None, ""}:
parts.append(f"value={item['value']!r}")
if item.get("href"):
parts.append(f"href={item['href']!r}")
if item.get("accept"):
parts.append(f"accept={item['accept']!r}")
if item.get("disabled"):
parts.append("disabled")
return " ".join(parts)
def _locator(self, reference: str) -> ElementLocator:
locator = self._locators.get(reference)
if locator is None or locator.generation != self._generation:
raise StaleElementError(
f"Element {reference} no longer matches the page snapshot. "
"Take another snapshot."
)
return locator
async def _verify(self, reference: str) -> ElementLocator:
locator = self._locator(reference)
try:
current = await self._evaluateEnvelope(
page_scripts.describeElement(locator.selector)
)
except ElementNotFoundError as error:
raise StaleElementError(
f"Element {reference} no longer matches the page snapshot."
) from error
if self._fingerprint(current) != locator.fingerprint:
raise StaleElementError(
f"Element {reference} no longer matches the page snapshot."
)
return locator
async def _refresh(self) -> None:
selected = self._selected()
tabs = await self.rdp.listTabs()
current = next(
(tab for tab in tabs if tab.actor_id == selected.actor_id),
None,
)
if current is None:
self.selected_tab = None
self._locators.clear()
raise TabNotFoundError("The selected Firefox tab disappeared.")
selected.title = current.title
if current.url != selected.last_url:
console = await self.rdp.getConsoleActor(selected.actor_id)
selected.console_actor_id = console
selected.last_url = current.url
self._invalidate()
selected.connection_generation = self.rdp.generation
selected.target_generation = self.rdp.target_generation
async def click(
self,
element: str,
wait_after_ms: int,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Verify and invoke an element's synthetic native click method."""
async with self._operation(cancelled):
await self._prepareSelected()
locator = await self._verify(element)
result = await self._evaluateEnvelope(
page_scripts.click(locator.selector)
)
if wait_after_ms:
await asyncio.sleep(wait_after_ms / 1000)
await self._refresh()
return {
"element": element,
"synthetic": True,
"url": self._selected().last_url,
"page_result": result,
}
async def fill(
self,
element: str,
value: str,
submit: bool,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Replace a referenced text control through its native setter."""
async with self._operation(cancelled):
await self._prepareSelected()
locator = await self._verify(element)
result = await self._evaluateEnvelope(
page_scripts.fill(locator.selector, value, submit)
)
await self._refresh()
return {
"element": element,
"synthetic": True,
"submitted": submit,
"url": self._selected().last_url,
"page_result": result,
}
async def typeText(
self,
element: str,
text: str,
delay_ms: int,
cancelled: CancellationCheck,
) -> dict[str, object]:
"""Append text one code point at a time with bounded evaluations."""
async with self._operation(cancelled):
await self._prepareSelected()
locator = await self._verify(element)
for character in text:
self._checkCancelled(cancelled)
await self._evaluateEnvelope(
page_scripts.appendCharacter(
locator.selector,
character,
)
)
if delay_ms:
await asyncio.sleep(delay_ms / 1000)
await self._refresh()
return {
"element": element,
"characters_typed": len(text),
"synthetic": True,
"url": self._selected().last_url,
}
@staticmethod
def _checkCancelled(cancelled: CancellationCheck) -> None:
if cancelled():
raise CancelledError("The MCP request was cancelled.")
async def pressKey(
self,
element: str | None,
key: str,
cancelled: CancellationCheck = neverCancelled,
) -> dict[str, object]:
"""Send an allowlisted synthetic key to a reference or active node."""
async with self._operation(cancelled):
await self._prepareSelected()
selector: str | None = None
if element is not None:
selector = (await self._verify(element)).selector
result = await self._evaluateEnvelope(
page_scripts.pressKey(selector, key)
)
await self._refresh()
return {
"element": element,
"key": key,
"synthetic": True,
"page_result": result,
"url": self._selected().last_url,
}
async def navigate(
self,
url: str,
wait: str,
timeout_ms: int,
cancelled: CancellationCheck,
) -> dict[str, object]:
"""Navigate the selected top-level page and reacquire its target."""
async with self._operation(cancelled):
selected = await self._prepareSelected()
actor = selected.actor_id
self._invalidate()
try:
await self._evaluateEnvelope(page_scripts.navigate(url))
except ConnectionLostError:
pass
await self._reacquire(actor, timeout_ms, cancelled)
if wait != "none":
await self._waitReadiness(wait, timeout_ms, cancelled)
state = await self._evaluateReadOnly(page_scripts.pageState())
selected = self._selected()
selected.last_url = str(state.get("url") or url)
selected.title = str(state.get("title") or "")
return {
"url": selected.last_url,
"title": selected.title,
"ready_state": state.get("ready_state"),
"generation": selected.generation,
}
async def _reacquire(
self,
actor_id: str,
timeout_ms: int,
cancelled: CancellationCheck,
) -> None:
deadline = asyncio.get_running_loop().time() + timeout_ms / 1000
last_error: Exception | None = None
while asyncio.get_running_loop().time() <= deadline:
self._checkCancelled(cancelled)
try:
if not self.rdp.connected:
await self.rdp.connect()
self._invalidate()
tabs = await self.rdp.listTabs()
tab = next(
(item for item in tabs if item.actor_id == actor_id),
None,
)
if tab is None:
raise TabNotFoundError(
"The navigated Firefox tab disappeared."
)
console = await self.rdp.getConsoleActor(actor_id)
self.selected_tab = SelectedTab(
actor_id,
console,
self._generation,
tab.url,
tab.title,
self.rdp.generation,
self.rdp.target_generation,
)
return
except (ConnectionLostError, PageScriptError) as error:
last_error = error
await asyncio.sleep(0.1)
raise NavigationTimeoutError(
"Timed out reacquiring the Firefox target after navigation."
) from last_error
async def _waitReadiness(
self,
wait: str,
timeout_ms: int,
cancelled: CancellationCheck,
) -> None:
deadline = asyncio.get_running_loop().time() + timeout_ms / 1000
while asyncio.get_running_loop().time() <= deadline:
self._checkCancelled(cancelled)
try:
value = await self._evaluateReadOnly(page_scripts.pageState())
ready = value.get("ready_state")
if wait == "dom" and ready in {"interactive", "complete"}:
return
if wait == "complete" and ready == "complete":
return
except ConnectionLostError:
await self._reacquire(
self._selected().actor_id,
max(1, int((deadline -
asyncio.get_running_loop().time()) * 1000)),
cancelled,
)
await asyncio.sleep(0.1)
raise NavigationTimeoutError(
f"Page did not reach {wait} readiness before the timeout."
)
async def waitFor(
self,
condition: str,
value: str,
state: str,
timeout_ms: int,
poll_ms: int,
cancelled: CancellationCheck,
) -> dict[str, object]:
"""Poll one visible-text, selector, or URL condition."""
async with self._operation(cancelled):
await self._prepareSelected()
deadline = asyncio.get_running_loop().time() + timeout_ms / 1000
while True:
self._checkCancelled(cancelled)
result = await self._evaluateReadOnly(
page_scripts.waitCondition(condition, value, state)
)
if result.get("matched") is True:
return {
"condition": condition,
"state": state,
"matched": True,
"url": result.get("url"),
}
if asyncio.get_running_loop().time() >= deadline:
raise WaitTimeoutError(
"The requested page condition did not become true."
)
await asyncio.sleep(poll_ms / 1000)
def _openFiles(self, paths: list[str]) -> list[UploadFile]:
files: list[UploadFile] = []
total = 0
try:
for raw_path in paths:
if "\0" in raw_path or not Path(raw_path).is_absolute():
raise UploadFileInvalidError(
"Every upload path must be an absolute file path."
)
try:
resolved = Path(raw_path).resolve(strict=True)
before = resolved.stat()
except OSError as error:
raise UploadFileInvalidError(
"An upload path is not a readable regular file."
) from error
if not stat.S_ISREG(before.st_mode):
raise UploadFileInvalidError(
"An upload path is not a regular file."
)
try:
stream = resolved.open("rb")
after = os.fstat(stream.fileno())
except OSError as error:
raise UploadFileInvalidError(
"An upload file could not be opened."
) from error
if (before.st_dev, before.st_ino) != (
after.st_dev,
after.st_ino,
):
stream.close()
raise UploadFileInvalidError(
"An upload file changed during validation."
)
if after.st_size > self.config.max_file_bytes:
stream.close()
raise UploadTooLargeError(
"An upload file exceeds the per-file byte limit."
)
total += after.st_size
if total > self.config.max_upload_bytes:
stream.close()
raise UploadTooLargeError(
"Upload files exceed the total byte limit."
)
mime_type = mimetypes.guess_type(resolved.name)[0]
files.append(UploadFile(
stream=stream,
name=resolved.name,
size=after.st_size,
mime_type=mime_type or "application/octet-stream",
last_modified=int(after.st_mtime * 1000),
))
return files
except Exception:
for item in files:
item.stream.close()
raise
async def uploadFiles(
self,
element: str,
paths: list[str],
cancelled: CancellationCheck,
) -> dict[str, object]:
"""Transfer validated local files and assign them without submit."""
async with self._operation(cancelled):
await self._prepareSelected()
locator = await self._verify(element)
if (
locator.tag_name.lower() != "input"
or locator.input_type != "file"
):
raise UploadTargetInvalidError(
"The referenced element is not a file input."
)
files = self._openFiles(paths)
key = "__firefox_mcp_upload_" + secrets.token_hex(16)
metadata = [
{
"name": item.name,
"size": item.size,
"type": item.mime_type,
"last_modified": item.last_modified,
}
for item in files
]
initialized = False
try:
target = await self._evaluateEnvelope(
page_scripts.uploadInitialize(
locator.selector,
key,
metadata,
)
)
initialized = True
for file_index, item in enumerate(files):
encoded_length = 0
chunk_index = 0
while True:
self._checkCancelled(cancelled)
chunk = item.stream.read(
self.config.upload_chunk_bytes
)
if not chunk:
break
encoded = base64.b64encode(chunk).decode("ascii")
encoded_length += len(encoded)
result = await self._evaluateEnvelope(
page_scripts.uploadAppend(
key,
file_index,
encoded,
)
)
if (
result.get("chunk_index") != chunk_index
or result.get("encoded_length")
!= encoded_length
):
raise UploadTransferError(
"Page upload staging confirmation disagreed."
)
chunk_index += 1
assigned = await self._evaluateEnvelope(
page_scripts.uploadAssign(key)
)
initialized = False
result_files = [
{
"name": item["name"],
"size": item["size"],
"type": item["type"],
}
for item in metadata
]
return {
"element": element,
"files": result_files,
"accept": target.get("accept"),
"accept_warnings": self._acceptWarnings(
str(target.get("accept") or ""),
result_files,
),
"input_dispatched": assigned.get("input_dispatched"),
"change_dispatched": assigned.get("change_dispatched"),
"submitted": False,
}
finally:
for item in files:
item.stream.close()
if initialized and self.rdp.connected:
try:
await self._evaluateEnvelope(
page_scripts.uploadCleanup(key)
)
except FirefoxMcpError:
pass
@staticmethod
def _acceptWarnings(
accept: str,
files: list[dict[str, object]],
) -> list[str]:
if not accept.strip():
return []
tokens = [
token.strip().lower()
for token in accept.split(",")
if token.strip()
]
warnings: list[str] = []
for item in files:
name = str(item["name"])
mime_type = str(item["type"]).lower()
lower_name = name.lower()
matched = any(
(
token.startswith(".")
and lower_name.endswith(token)
)
or (
token.endswith("/*")
and mime_type.startswith(token[:-1])
)
or token == mime_type
for token in tokens
)
if not matched:
warnings.append(
f"{name} may not match the input accept attribute."
)
return warnings