aiturk-hermes-ide/tests/gateway/test_browser_control_local_lane.py

"""Local desktop browser lane: the AITURK Chrome extension behind the IDE.

One controller registered from the loopback desktop transport serves every
desktop session of the backend; each session still gets an exact scope.
"""

import base64
import json
import threading
from types import SimpleNamespace

import pytest

from gateway.browser_control_broker import (
    LOCAL_DESKTOP_PRINCIPAL,
    LOCAL_DESKTOP_TRANSPORT_FAMILY,
    BrowserControlBroker,
    ControllerCancelled,
    ControllerScope,
    get_browser_control_broker,
)
from hermes_cli import web_server
from tools.browser_extension_router import route_browser_tool
from tui_gateway import server

CAPS = frozenset({"browser_navigate", "browser_snapshot", "browser_vision"})


def _template(controller_id="chrome-1", capabilities=CAPS):
    return ControllerScope(
        principal_id=LOCAL_DESKTOP_PRINCIPAL,
        profile_id="default",
        session_id=None,
        controller_id=controller_id,
        browser_profile_id="chrome-profile",
        transport_family=LOCAL_DESKTOP_TRANSPORT_FAMILY,
        capabilities=capabilities,
    )


class _Writer:
    def __init__(self):
        self.frames = []
        self.sent = threading.Event()

    def send_for(self, session_id):
        def send(frame):
            self.frames.append((session_id, frame))
            self.sent.set()

        return send


def _lookup(broker, session_id):
    return broker.scope_for_session(
        session_id=session_id,
        principal_id=LOCAL_DESKTOP_PRINCIPAL,
        transport_family=LOCAL_DESKTOP_TRANSPORT_FAMILY,
    )


def _dispatch_async(broker, scope, action="browser_navigate"):
    outcome = {}

    def run():
        try:
            outcome["result"] = broker.dispatch(
                scope, action=action, arguments={"url": "https://example.com"}, tool_call_id="call-1"
            )
        except Exception as exc:  # noqa: BLE001 - asserted by the tests
            outcome["error"] = exc

    thread = threading.Thread(target=run)
    thread.start()
    return thread, outcome


def test_lane_template_serves_every_session_with_its_own_exact_scope():
    broker = BrowserControlBroker(command_timeout=5)
    writer = _Writer()
    owner = object()
    assert broker.lane_registered(
        session_id="a", principal_id=LOCAL_DESKTOP_PRINCIPAL, transport_family=LOCAL_DESKTOP_TRANSPORT_FAMILY
    ) is False
    broker.attach_lane_template(_template(), writer.send_for, owner=owner)

    first, second = _lookup(broker, "session-a"), _lookup(broker, "session-b")
    assert first.session_id == "session-a" and second.session_id == "session-b"
    assert _lookup(broker, "session-a") == first
    assert broker.select(first, "browser_navigate") is not None
    assert broker.select(first, "browser_click") is None
    assert broker.lane_registered(
        session_id="anything", principal_id=LOCAL_DESKTOP_PRINCIPAL, transport_family=LOCAL_DESKTOP_TRANSPORT_FAMILY
    )

    thread, outcome = _dispatch_async(broker, second)
    assert writer.sent.wait(2)
    session_id, frame = writer.frames[-1]
    assert session_id == "session-b"
    command_id = frame["params"]["command_id"]
    assert broker.complete_for_owner(command_id, owner=object(), ok=True, result="x") is False
    assert broker.complete_for_owner(command_id, owner=owner, ok=True, result={"success": True}) is True
    thread.join(2)
    assert outcome == {"result": {"success": True}}


def test_other_lanes_never_see_the_local_template():
    broker = BrowserControlBroker()
    broker.attach_lane_template(_template(), _Writer().send_for, owner=object())
    assert broker.scope_for_session(
        session_id="session-a", principal_id="principal:dashboard:x", transport_family="cloud-ticket-ws"
    ) is None
    assert broker.scope_for_session(
        session_id="session-a", principal_id=LOCAL_DESKTOP_PRINCIPAL, transport_family="cloud-ticket-ws"
    ) is None


def test_losing_the_transport_removes_the_lane_and_cancels_its_work():
    broker = BrowserControlBroker(command_timeout=5)
    writer, owner = _Writer(), object()
    broker.attach_lane_template(_template(), writer.send_for, owner=owner)
    scope = _lookup(broker, "session-a")
    thread, outcome = _dispatch_async(broker, scope)
    assert writer.sent.wait(2)

    broker.disconnect_owner(owner)
    thread.join(2)
    assert isinstance(outcome.get("error"), ControllerCancelled)
    assert _lookup(broker, "session-a") is None
    assert broker.lane_registered(
        session_id="session-a", principal_id=LOCAL_DESKTOP_PRINCIPAL, transport_family=LOCAL_DESKTOP_TRANSPORT_FAMILY
    ) is False


def test_same_controller_refresh_keeps_work_but_a_new_controller_replaces_it():
    broker = BrowserControlBroker(command_timeout=5)
    writer, owner = _Writer(), object()
    broker.attach_lane_template(_template(), writer.send_for, owner=owner)
    scope = _lookup(broker, "session-a")
    thread, outcome = _dispatch_async(broker, scope)
    assert writer.sent.wait(2)
    command_id = writer.frames[-1][1]["params"]["command_id"]

    refreshed_writer = _Writer()
    broker.attach_lane_template(
        _template(capabilities=CAPS | {"browser_click"}), refreshed_writer.send_for, owner=owner
    )
    refreshed = _lookup(broker, "session-a")
    assert broker.select(refreshed, "browser_click") is not None
    assert broker.complete_for_owner(command_id, owner=owner, ok=True, result="done") is True
    thread.join(2)
    assert outcome == {"result": "done"}

    thread, outcome = _dispatch_async(broker, refreshed)
    assert refreshed_writer.sent.wait(2)
    broker.attach_lane_template(_template(controller_id="chrome-2"), _Writer().send_for, owner=owner)
    thread.join(2)
    assert isinstance(outcome.get("error"), ControllerCancelled)
    assert _lookup(broker, "session-a").controller_id == "chrome-2"


def test_router_uses_local_lane_without_the_remote_feature_flag(monkeypatch):
    monkeypatch.setattr("gateway.browser_control_broker.browser_control_enabled", lambda config=None: False)
    monkeypatch.setattr("gateway.browser_control_broker.browser_control_local_enabled", lambda config=None: True)
    broker = BrowserControlBroker(command_timeout=5)
    calls = []

    def fallback():
        calls.append("legacy")
        return "legacy"

    identity = dict(
        session_id="session-a",
        principal_id=LOCAL_DESKTOP_PRINCIPAL,
        transport_family=LOCAL_DESKTOP_TRANSPORT_FAMILY,
    )
    # No controller registered yet: the normal browser keeps working.
    assert route_browser_tool("browser_navigate", {}, fallback=fallback, broker=broker, enabled=True, **identity) == "legacy"

    owner = object()

    def send_for(session_id):
        def send(frame):
            broker.complete_for_owner(
                frame["params"]["command_id"], owner=owner, ok=True, result={"success": True, "screenshot": "eA=="}
            )

        return send

    broker.attach_lane_template(_template(), send_for, owner=owner)
    seen = []
    result = route_browser_tool(
        "browser_vision", {"question": "q"}, fallback=fallback, broker=broker, enabled=True,
        transform=lambda value: seen.append(value) or {"envelope": True}, **identity,
    )
    assert result == {"envelope": True}
    assert seen == [{"success": True, "screenshot": "eA=="}]
    assert calls == ["legacy"]


def test_lane_switch_off_keeps_the_normal_browser(monkeypatch):
    from tools import browser_extension_router as router

    monkeypatch.setattr("gateway.browser_control_broker.browser_control_local_enabled", lambda config=None: False)
    assert router.lane_enabled(LOCAL_DESKTOP_TRANSPORT_FAMILY) is False
    monkeypatch.setattr("gateway.browser_control_broker.browser_control_local_enabled", lambda config=None: True)
    assert router.lane_enabled(LOCAL_DESKTOP_TRANSPORT_FAMILY) is True


def test_local_switch_defaults_on_and_respects_explicit_off():
    from gateway.browser_control_broker import browser_control_local_enabled

    assert browser_control_local_enabled({}) is True
    assert browser_control_local_enabled({"browser": {"extension_control": {"enabled": False}}}) is True
    assert browser_control_local_enabled({"browser": {"extension_control": {"local_desktop": False}}}) is False


def test_only_the_loopback_token_marks_a_local_desktop_transport():
    token_ws = SimpleNamespace(
        query_params={"token": web_server._SESSION_TOKEN}, headers={},
        client=SimpleNamespace(host="127.0.0.1"), url=SimpleNamespace(path="/api/ws"),
    )
    previous = getattr(web_server.app.state, "auth_required", False)
    web_server.app.state.auth_required = False
    try:
        assert web_server._ws_auth_ok(token_ws) is True
        assert token_ws._hermes_local_desktop is True
        wrong = SimpleNamespace(
            query_params={"token": "not-the-token"}, headers={},
            client=SimpleNamespace(host="127.0.0.1"), url=SimpleNamespace(path="/api/ws"),
        )
        assert web_server._ws_auth_ok(wrong) is False
        assert getattr(wrong, "_hermes_local_desktop", False) is False
    finally:
        web_server.app.state.auth_required = previous


def test_desktop_session_context_binds_the_local_lane_identity():
    from gateway.session_context import clear_session_vars, get_session_env

    server._sessions["local-context-fixture"] = {
        "transport": SimpleNamespace(auth_identity=None, local_desktop=True),
        "session_key": "stored-local-context",
        "profile": None,
        "agent": SimpleNamespace(session_id="durable-local-session"),
    }
    tokens = []
    try:
        tokens = server._set_session_context("stored-local-context")
        assert get_session_env("HERMES_SESSION_ID") == "durable-local-session"
        assert get_session_env("HERMES_BROWSER_CONTROL_PRINCIPAL") == LOCAL_DESKTOP_PRINCIPAL
        assert get_session_env("HERMES_BROWSER_CONTROL_TRANSPORT_FAMILY") == LOCAL_DESKTOP_TRANSPORT_FAMILY
    finally:
        clear_session_vars(tokens)
        server._sessions.pop("local-context-fixture", None)


def _rpc(method, params, transport):
    return server.dispatch({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}, transport)


class _Transport:
    def __init__(self, local=True):
        self.auth_identity = None
        self.local_desktop = local
        self.frames = []
        self.sent = threading.Event()

    def write(self, frame):
        self.frames.append(frame)
        self.sent.set()
        return True


REGISTER = {
    "protocol_version": 1,
    "controller_id": "chrome-extension",
    "browser_profile_id": "chrome-default",
    "capabilities": ["browser_navigate", "browser_snapshot", "browser_cdp", "unknown.method"],
}


def test_local_registration_is_refused_off_the_desktop_transport(monkeypatch):
    monkeypatch.setattr("gateway.browser_control_broker.browser_control_local_enabled", lambda config=None: True)
    get_browser_control_broker().reset()
    assert _rpc("browser.controller.local.register", REGISTER, _Transport(local=False))["error"]["code"] == 4403
    monkeypatch.setattr("gateway.browser_control_broker.browser_control_local_enabled", lambda config=None: False)
    assert _rpc("browser.controller.local.register", REGISTER, _Transport())["error"]["code"] == 4403
    assert _rpc("browser.controller.local.register", {**REGISTER, "capabilities": ["browser_cdp"]}, _Transport())["error"]["code"] == 4403


def test_local_gateway_round_trip_and_detach(monkeypatch):
    monkeypatch.setattr("gateway.browser_control_broker.browser_control_local_enabled", lambda config=None: True)
    broker = get_browser_control_broker()
    broker.reset()
    transport = _Transport()
    try:
        registered = _rpc("browser.controller.local.register", REGISTER, transport)["result"]
        assert registered["capabilities"] == ["browser_navigate", "browser_snapshot"]

        scope = _lookup(broker, "durable-session")
        thread, outcome = _dispatch_async(broker, scope)
        assert transport.sent.wait(2)
        event = transport.frames[-1]
        assert event["method"] == "event"
        assert event["params"]["type"] == "browser.controller.command"
        assert event["params"]["session_id"] == "durable-session"
        command_id = event["params"]["payload"]["command_id"]

        foreign = _rpc("browser.controller.local.result", {"command_id": command_id, "ok": True}, _Transport())
        assert foreign["result"]["accepted"] is False
        done = _rpc("browser.controller.local.result", {"command_id": command_id, "ok": True, "result": {"success": True}}, transport)
        assert done["result"]["accepted"] is True
        thread.join(2)
        assert outcome == {"result": {"success": True}}

        assert _rpc("browser.controller.local.detach", {}, transport)["result"]["detached"] is True
        assert _lookup(broker, "durable-session") is None
    finally:
        broker.reset()


def test_vision_from_controller_saves_the_jpeg_and_reuses_the_analysis(monkeypatch, tmp_path):
    from tools import browser_tool

    monkeypatch.setattr("hermes_constants.get_hermes_dir", lambda *parts: tmp_path)
    captured = {}

    def analyze(path, question, annotate, result, *, mime_type="image/png", fallback_warning=None):
        captured.update(path=path, question=question, mime=mime_type, bytes=path.read_bytes())
        return {"native": True}

    monkeypatch.setattr(browser_tool, "_analyze_browser_screenshot", analyze)
    jpeg = b"\xff\xd8\xff\xe0fixture"
    out = browser_tool._browser_vision_from_controller(
        {"success": True, "screenshot": base64.b64encode(jpeg).decode(), "mime": "image/jpeg"}, "Ne görünüyor?"
    )
    assert out == {"native": True}
    assert captured["bytes"] == jpeg and captured["mime"] == "image/jpeg"
    assert captured["path"].suffix == ".jpg" and captured["question"] == "Ne görünüyor?"

    bad = json.loads(browser_tool._browser_vision_from_controller({"success": True, "screenshot": "%%%"}, "q"))
    assert bad["success"] is False
    failed = json.loads(browser_tool._browser_vision_from_controller({"success": False, "error": "Sekme kapandı"}, "q"))
    assert failed == {"success": False, "error": "Sekme kapandı"}