browser-cli/browser_cli/sdk/routing.py

"""Multi-browser routing mixin for :class:`~browser_cli.BrowserCLI`.

When no specific browser is selected and more than one browser (local or
remote) is active, list/count commands fan out to every target and aggregate
the results. This mixin holds that fan-out machinery plus a few small response
helpers; single-browser mode falls straight through to ``_cmd``.
"""
from __future__ import annotations

import importlib
import sys
from collections.abc import Callable, Iterable
from typing import TYPE_CHECKING, Protocol, cast

from browser_cli.client import BrowserTarget
from browser_cli.errors import BrowserNotConnected
from browser_cli.models import BrowserCounts, Tab

if TYPE_CHECKING:
  from browser_cli.sdk.tabs import TabsNS

class _RoutingClient(Protocol):
  _browser: str | None
  _remote: str | None
  _key: str | None
  tabs: "TabsNS"

  def dispatch(self, command: str, args: dict | None = None): ...

# send_command / active_browser_targets / remote_browser_targets are resolved
# through the ``browser_cli`` package namespace at call time, not bound here at
# import, so tests patching ``browser_cli.send_command`` still take effect —
# without a module-level import cycle back into browser_cli.__init__.

_UNSET = object()

def _browser_cli_package():
  return sys.modules.get("browser_cli") or importlib.import_module("browser_cli")

def _with_profile_display(targets: list[BrowserTarget]) -> list[BrowserTarget]:
  """Use profile-only labels when a command is already scoped to one remote."""
  return [
    BrowserTarget(
      profile=target.profile,
      display_name=target.profile if target.remote else target.display_name,
      socket_path=target.socket_path,
      remote=target.remote,
      browser_name=target.browser_name,
      display_group=None,
    )
    for target in targets
  ]

class RoutingMixin:
  """Fan-out + aggregation across active browsers, mixed into ``BrowserCLI``.

  Relies on the client exposing ``_browser``/``_remote``/``_key``, ``_cmd``,
  and a ``tabs`` namespace.
  """

  @property
  def _client(self) -> _RoutingClient:
    return cast(_RoutingClient, cast(object, self))

  def _multi_browser_targets(self) -> list[BrowserTarget]:
    client = self._client
    package = _browser_cli_package()
    if client._browser is not None and not client._remote:
      targets = package.remote_targets_for_alias(client._browser, key=client._key)
      if len(targets) <= 1:
        return []
      targets = _with_profile_display(targets)
    elif client._browser is not None:
      return []
    elif client._remote:
      targets = _with_profile_display(package.remote_browser_targets(client._remote, key=client._key))
    else:
      targets = package.active_browser_targets()
    if len(targets) <= 1 and not any(target.remote for target in targets):
      return []
    return targets

  def _collect_multi_browser(self, command: str, args: dict | None = None):
    results = []
    targets = self._multi_browser_targets()
    for target in targets:
      try:
        if target.remote:
          data = _browser_cli_package().send_command(
            command, args, profile=target.profile, remote=target.remote, key=self._client._key
          )
        else:
          data = _browser_cli_package().send_command(command, args, profile=target.profile)
      except (BrowserNotConnected, RuntimeError):
        continue
      results.append((target, data))
    if results:
      return results
    if targets:
      raise BrowserNotConnected(
        "Cannot resolve a browser socket automatically.\n"
        "Make sure the browser is running with the browser-cli extension enabled,\n"
        "or pass --browser <alias> / set BROWSER_CLI_PROFILE to a known alias."
      )
    return []

  @staticmethod
  def _field(result, key, default=None, *, fallback=_UNSET):
    """Pull *key* out of a dict response, with a non-dict fallback.

    Returns ``result[key]`` (or *default*) when *result* is a dict. When it
    is not a dict, returns *fallback* if given, else *default*.
    """
    if isinstance(result, dict):
      return result.get(key, default)
    return default if fallback is _UNSET else fallback

  def toggle_tab(self, command: str, tab_id: int | None) -> int:
    """Run a tab toggle command (mute/pin/...) and return the target tab ID."""
    result = self._client.dispatch(command, {"tabId": tab_id})
    return self._field(result, "tabId", tab_id, fallback=int(tab_id or 0))

  def multi_count(self, command: str, args: dict | None = None) -> "int | BrowserCounts":
    """Count command that aggregates into :class:`BrowserCounts` in multi-browser mode."""
    multi_results = self._collect_multi_browser(command, args or {})
    if not multi_results:
      return self._client.dispatch(command, args or {})
    by_browser = {target.display_name: int(count or 0) for target, count in multi_results}
    browser_groups = {
      target.display_name: target.display_group or "local"
      for target, _count in multi_results
      if target.display_group or target.remote is None
    }
    return BrowserCounts(total=sum(by_browser.values()), by_browser=by_browser, browser_groups=browser_groups)

  def multi_list(self, command: str, args: dict | None, mapper):
    """List command, flattening per-browser results in multi-browser mode.

    *mapper* is ``(item, target) -> mapped`` where ``target`` is the source
    :class:`BrowserTarget` in multi mode, or ``None`` in single-browser mode.
    """
    multi_results = self._collect_multi_browser(command, args or {})
    if multi_results:
      return [
        mapper(item, target)
        for target, items in multi_results
        for item in (items or [])
      ]
    return [mapper(item, None) for item in (self._client.dispatch(command, args or {}) or [])]

  def apply_tab_filter(self, filter_fn: Callable[[Tab], bool] | Callable[[list[Tab]], Iterable[Tab]]) -> list[Tab]:
    tabs = self._client.tabs.list()

    try:
      transformed = filter_fn(tabs)
    except (AttributeError, TypeError):
      return [tab for tab in tabs if filter_fn(tab)]

    if isinstance(transformed, list):
      return transformed
    if isinstance(transformed, tuple):
      return list(transformed)
    if isinstance(transformed, set):
      return list(transformed)
    if transformed is tabs:
      return tabs
    if isinstance(transformed, bool):
      return [tab for tab in tabs if filter_fn(tab)]

    try:
      return list(transformed)
    except TypeError:
      return [tab for tab in tabs if filter_fn(tab)]