browser-cli/browser_cli/client/core.py

"""
Local IPC client — sends commands to native host relay endpoint.
Used by both CLI and public Python API.

Profile selection order:
  1. Explicit `profile` argument to send_command()
  2. BROWSER_CLI_PROFILE environment variable
  3. First entry in runtime registry
  4. Otherwise, no browser can be resolved automatically
"""
import asyncio
from typing import Any

from browser_cli import local_transport
from browser_cli.client import auth, messages, targets as target_discovery
from browser_cli.client.targets import BrowserTarget
from browser_cli.remote import registry as remote_registry
from browser_cli.errors import BrowserNotConnected
from browser_cli.endpoints import _remote_display_name, display_browser_name
from browser_cli.registry import load_registry
from browser_cli.remote.transport import _send_remote, _send_remote_async

def _run_concurrent(factories: list) -> list:
  """Run async thunks concurrently, returning results in order.

  Each item in *factories* is a zero-arg callable returning a coroutine. The
  return list mirrors the input order; a thunk that raises yields its exception
  object in that slot (callers filter as they would in a sequential loop). Falls
  back to sequential execution if an event loop is already running on this
  thread (e.g. inside the async serve handler), where ``asyncio.run`` is illegal.
  """
  if not factories:
    return []

  async def _gather():
    return await asyncio.gather(*(factory() for factory in factories), return_exceptions=True)

  try:
    asyncio.get_running_loop()
  except RuntimeError:
    return asyncio.run(_gather())

  # An event loop is already running on this thread (e.g. the async serve
  # handler), where asyncio.run is illegal. Run the gather on a worker thread
  # that has no loop of its own, preserving concurrency and result order.
  import concurrent.futures

  with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
    return executor.submit(lambda: asyncio.run(_gather())).result()

def _remote_target_items(endpoint: str, items: list[dict] | None) -> list[BrowserTarget]:
  targets: list[BrowserTarget] = []
  for item in items or []:
    profile = str(item.get("profile") or "default")
    display = str(item.get("displayName") or profile)
    display_name = _remote_display_name(endpoint, profile, display)
    browser_name = item.get("browserName") or item.get("name")
    version = item.get("version")
    extension_version = item.get("extensionVersion")
    targets.append(
      BrowserTarget(
        profile=profile,
        display_name=display_name,
        socket_path="",
        remote=endpoint,
        browser_name=str(browser_name) if browser_name else None,
        display_group=display_name.rsplit(":", 1)[0],
        version=str(version) if version else None,
        extension_version=str(extension_version) if extension_version else None,
      )
    )
  return targets

def remote_browser_targets(endpoint: str, key=None, *, suppress_pq_warning: bool = False) -> list[BrowserTarget]:
  """Return browser targets advertised by a single remote endpoint."""
  kwargs = {"suppress_pq_warning": True} if suppress_pq_warning else {}
  return _remote_target_items(
    endpoint,
    send_command("browser-cli.targets", remote=endpoint, key=key, **kwargs),
  )

def _remote_browser_targets(key=None, *, suppress_pq_warning: bool = False) -> list[BrowserTarget]:
  endpoints = list(remote_registry.load_remotes())
  if not endpoints:
    return []
  results = _run_concurrent([
    (lambda ep=ep: asyncio.to_thread(remote_browser_targets, ep, key=key, suppress_pq_warning=suppress_pq_warning))
    for ep in endpoints
  ])
  targets: list[BrowserTarget] = []
  for result in results:
    if isinstance(result, (BrowserNotConnected, RuntimeError)):
      continue
    if isinstance(result, BaseException):
      raise result
    targets.extend(result)
  return targets

def remote_targets_for_alias(alias: str | None, key=None) -> list[BrowserTarget]:
  """Return remote targets matching a user-facing alias.

  Exact browser aliases such as ``host:profile`` return one target. Endpoint
  aliases such as ``host`` or ``host:8765`` may return multiple targets, which
  lets read/list SDK commands fan out while command dispatch can still reject
  the ambiguous target.
  """
  if not alias:
    return []
  targets = _remote_browser_targets(key=key) if key is not None else _remote_browser_targets()
  for target in targets:
    endpoint_profile = f"{target.remote}:{target.profile}" if target.remote else None
    if alias in {target.display_name, endpoint_profile}:
      return [target]

  endpoint_matches = []
  for target in targets:
    if not target.remote:
      continue
    remote_host, sep, _remote_port = target.remote.rpartition(":")
    if alias == target.remote or (sep and alias == remote_host):
      endpoint_matches.append(target)
  return endpoint_matches

def remote_target_for_alias(alias: str | None) -> BrowserTarget | None:
  """Resolve a user-facing remote alias such as 'host:profile' to a target."""
  matches = remote_targets_for_alias(alias)
  if len(matches) == 1:
    return matches[0]
  if len(matches) > 1:
    aliases = [target.profile for target in matches]
    endpoint = matches[0].remote or alias or "remote"
    examples = "\n".join(
      f"  browser-cli --remote {endpoint} --browser {a} ..."
      for a in aliases
    )
    display_aliases = [target.display_name for target in matches]
    shorthand_examples = "\n".join(
      f"  browser-cli --browser {a} ..."
      for a in display_aliases
    )
    raise BrowserNotConnected(
      f"Multiple remote browser instances are active on {alias}: {', '.join(aliases)}\n"
      f"Use --browser <alias> with --remote to select one:\n{examples}\n"
      f"Or use the full remote browser alias:\n{shorthand_examples}"
    )
  return None

def active_browser_targets(*, include_remotes: bool = True, key=None, suppress_pq_warning: bool = False) -> list[BrowserTarget]:
  targets = target_discovery.active_local_browser_targets()
  if include_remotes:
    targets.extend(_remote_browser_targets(key=key, suppress_pq_warning=suppress_pq_warning))
  return targets

def _cached_client_row(target: BrowserTarget) -> dict | None:
  """Build a clients row from a target's discovery data, skipping a roundtrip.

  Returns None when the remote didn't advertise its version (older serve), so
  callers fall back to an explicit ``clients.list`` query.
  """
  if target.version is None and target.extension_version is None:
    return None
  return {
    "profile": target.display_name,
    "profileGroup": target.display_group,
    "name": target.browser_name or "",
    "version": target.version or "",
    "extensionVersion": target.extension_version or "",
  }

def _rows_from_result(result, label: str, profile_group: str | None) -> list[dict]:
  rows = []
  for item in result or []:
    row = dict(item)
    row["profile"] = label
    if profile_group:
      row["profileGroup"] = profile_group
    rows.append(row)
  return rows

async def _client_rows_async(
  label: str,
  *,
  profile: str | None = None,
  remote: str | None = None,
  key=None,
  suppress_pq_warning: bool = False,
  profile_group: str | None = None,
) -> list[dict]:
  """Return display-ready clients.list rows for one browser target."""
  kwargs = {"suppress_pq_warning": True} if suppress_pq_warning else {}
  result = await asyncio.to_thread(
    send_command,
    "clients.list",
    profile=profile,
    remote=remote,
    key=key,
    **kwargs,
  )
  return _rows_from_result(result, label, profile_group)

def collect_browser_clients(
  *,
  browser_alias: str | None = None,
  remote: str | None = None,
  key=None,
  registry_path=None,
) -> list[dict]:
  """Return display-ready browser client rows for CLI/SDK consumers.

  Rows preserve the CLI-facing shape: ``profile``, optional ``profileGroup``,
  ``name``, ``version``, and ``extensionVersion``.
  """
  rows: list[dict] = []

  if not remote and browser_alias:
    resolved = remote_target_for_alias(browser_alias)
    if not resolved:
      return rows
    targets = remote_browser_targets(resolved.remote)
    uncached = []
    for target in targets:
      cached = _cached_client_row(target)
      if cached is not None:
        rows.append(cached)
      else:
        uncached.append(target)
    results = _run_concurrent([
      (lambda t=t: _client_rows_async(
        t.display_name,
        profile=t.profile,
        remote=resolved.remote,
        key=key,
        profile_group=t.display_group,
      ))
      for t in uncached
    ])
    for result in results:
      if isinstance(result, (BrowserNotConnected, RuntimeError)):
        continue
      if isinstance(result, BaseException):
        raise result
      rows.extend(result)
    return rows

  if remote:
    result = send_command("clients.list", profile=browser_alias, remote=remote, key=key)
    for item in result or []:
      row = dict(item)
      row["profile"] = row.get("profile") or browser_alias or "remote"
      rows.append(row)
    return rows

  path = registry_path or target_discovery.REGISTRY_PATH
  profiles: dict[str, str] = load_registry(path) if path.exists() else {}
  local_items = list(profiles.items())

  remote_targets = []
  cached_remote_rows = []  # deferred so local profiles still render first
  for target in active_browser_targets(suppress_pq_warning=True):
    if target.remote is None:
      continue
    cached = _cached_client_row(target)
    if cached is not None:
      cached_remote_rows.append(cached)  # discovery already carried version/extVersion — no extra roundtrip
    else:
      remote_targets.append(target)

  factories = [
    (lambda name=name, sock=sock: _client_rows_async(
      display_browser_name(name, sock), profile=name, profile_group="local",
    ))
    for name, sock in local_items
  ] + [
    (lambda t=t: _client_rows_async(
      t.display_name,
      profile=t.profile,
      remote=t.remote,
      suppress_pq_warning=True,
      profile_group=t.display_group,
    ))
    for t in remote_targets
  ]
  results = _run_concurrent(factories)

  for (name, sock), result in zip(local_items, results[:len(local_items)]):
    if isinstance(result, (BrowserNotConnected, RuntimeError)):
      rows.append({
        "profile": display_browser_name(name, sock),
        "profileGroup": "local",
        "name": "—",
        "version": "—",
        "extensionVersion": "disconnected",
      })
    elif isinstance(result, BaseException):
      raise result
    else:
      rows.extend(result)

  for result in results[len(local_items):]:
    if isinstance(result, (BrowserNotConnected, RuntimeError)):
      continue
    if isinstance(result, BaseException):
      raise result
    rows.extend(result)
  rows.extend(cached_remote_rows)
  return rows

def _auto_route_remote(endpoint: str, key=None) -> str | None:
  targets = remote_browser_targets(endpoint, key=key)
  if len(targets) == 1:
    return targets[0].profile
  if len(targets) > 1:
    aliases = [target.profile for target in targets]
    examples = "\n".join(f"  browser-cli --remote {endpoint} --browser {a} ..." for a in aliases)
    raise BrowserNotConnected(
      f"Multiple remote browser instances are active: {', '.join(aliases)}\n"
      f"Use --browser <alias> to select one:\n{examples}"
    )
  return None

def send_command(
  command: str,
  args: dict | None = None,
  profile: str | None = None,
  remote: str | None = None,
  key: "Path | None" = None,
  *,
  suppress_pq_warning: bool = False,
) -> Any:
  """Send a command to the browser and return the response data."""
  requested_profile, remote_endpoint = messages.requested_target(profile, remote)
  if not remote_endpoint and requested_profile and not target_discovery.is_active_local_profile(requested_profile):
    if remote_alias_target := remote_target_for_alias(requested_profile):
      remote_endpoint = remote_alias_target.remote
      requested_profile = remote_alias_target.profile

  msg = messages.base_message(command, args)
  private_key = None
  if remote_endpoint:
    if suppress_pq_warning:
      msg["_suppress_pq_warning"] = True
    private_key = auth.add_remote_auth_fields(msg, command, requested_profile, remote_endpoint, key, _auto_route_remote)

  try:
    payload = messages.encode_payload(msg)
    response = (
      _send_remote(remote_endpoint, msg, private_key)
      if remote_endpoint
      else local_transport.send_local_sync(requested_profile, payload, target_discovery.resolve_socket)
    )
  except (FileNotFoundError, ConnectionRefusedError, OSError):
    raise messages.remote_connection_error(remote_endpoint) if remote_endpoint else messages.local_connection_error(requested_profile)

  return messages.decode_response(response)

async def remote_browser_targets_async(endpoint: str, key=None, *, suppress_pq_warning: bool = False) -> list[BrowserTarget]:
  """Async variant of :func:`remote_browser_targets`."""
  kwargs = {"suppress_pq_warning": True} if suppress_pq_warning else {}
  return _remote_target_items(
    endpoint,
    await send_command_async("browser-cli.targets", remote=endpoint, key=key, **kwargs),
  )

async def _auto_route_remote_async(endpoint: str, key=None) -> str | None:
  targets = await remote_browser_targets_async(endpoint, key=key)
  if len(targets) == 1:
    return targets[0].profile
  if len(targets) > 1:
    aliases = [target.profile for target in targets]
    examples = "\n".join(f"  browser-cli --remote {endpoint} --browser {a} ..." for a in aliases)
    raise BrowserNotConnected(
      f"Multiple remote browser instances are active: {', '.join(aliases)}\n"
      f"Use --browser <alias> to select one:\n{examples}"
    )
  return None

async def send_command_async(
  command: str,
  args: dict | None = None,
  profile: str | None = None,
  remote: str | None = None,
  key: "Path | None" = None,
  *,
  suppress_pq_warning: bool = False,
) -> Any:
  """Async variant of :func:`send_command` with native async socket/TCP paths."""
  requested_profile, remote_endpoint = messages.requested_target(profile, remote)
  if not remote_endpoint and requested_profile and not await asyncio.to_thread(target_discovery.is_active_local_profile, requested_profile):
    if remote_alias_target := await asyncio.to_thread(remote_target_for_alias, requested_profile):
      remote_endpoint = remote_alias_target.remote
      requested_profile = remote_alias_target.profile

  msg = messages.base_message(command, args)
  private_key = None
  if remote_endpoint:
    if suppress_pq_warning:
      msg["_suppress_pq_warning"] = True
    private_key = await auth.add_remote_auth_fields_async(msg, command, requested_profile, remote_endpoint, key, _auto_route_remote_async)

  try:
    payload = messages.encode_payload(msg)
    response = (
      await _send_remote_async(remote_endpoint, msg, private_key)
      if remote_endpoint
      else await local_transport.send_local_async(requested_profile, payload, target_discovery.resolve_socket)
    )
  except (FileNotFoundError, ConnectionRefusedError, OSError):
    raise messages.remote_connection_error(remote_endpoint) if remote_endpoint else messages.local_connection_error(requested_profile)

  return messages.decode_response(response)