| 1 | from __future__ import annotations |
| 2 | |
| 3 | import asyncio |
| 4 | from pathlib import Path |
| 5 | import uuid |
| 6 | |
| 7 | from plugins._a0_connector.api.v1 import launcher_gateway_control |
| 8 | from plugins._a0_connector.api.v1.capabilities import _feature_list |
| 9 | from plugins._a0_connector.api.v1.launcher_gateway_status import LauncherGatewayStatus |
| 10 | from plugins._a0_connector.api.ws_connector import WS_FEATURES |
| 11 | from plugins._a0_connector.helpers import ws_runtime |
| 12 | |
| 13 | |
| 14 | def _sid(label: str) -> str: |
| 15 | return f"gateway-{label}-{uuid.uuid4()}" |
| 16 | |
| 17 | |
| 18 | def _gateway( |
| 19 | gateway_id: str, |
| 20 | *, |
| 21 | files: bool = True, |
| 22 | file_write: bool | None = None, |
| 23 | ) -> dict: |
| 24 | return { |
| 25 | "version": 1, |
| 26 | "kind": "launcher", |
| 27 | "id": gateway_id, |
| 28 | "host_label": "Test host", |
| 29 | "state": "connected", |
| 30 | "master_enabled": True, |
| 31 | "scopes": { |
| 32 | "files": files, |
| 33 | "file_write": files if file_write is None else file_write, |
| 34 | "code_execution": True, |
| 35 | "browser": True, |
| 36 | "computer_use": True, |
| 37 | }, |
| 38 | } |
| 39 | |
| 40 | |
| 41 | def test_launcher_gateway_features_are_negotiated_on_http_and_websocket() -> None: |
| 42 | assert "launcher_gateway" in _feature_list() |
| 43 | assert "launcher_gateway_file_write" in _feature_list() |
| 44 | assert "launcher_gateway_control" in WS_FEATURES |
| 45 | assert LauncherGatewayStatus.requires_auth() is True |
| 46 | |
| 47 | |
| 48 | def test_launcher_gateway_has_no_agent_zero_webui_controls() -> None: |
| 49 | root = Path(__file__).parents[1] |
| 50 | extension = ( |
| 51 | root |
| 52 | / "plugins" |
| 53 | / "_a0_connector" |
| 54 | / "extensions" |
| 55 | / "webui" |
| 56 | / "sync-status-end" |
| 57 | / "launcher-gateway.html" |
| 58 | ) |
| 59 | store = ( |
| 60 | root |
| 61 | / "plugins" |
| 62 | / "_a0_connector" |
| 63 | / "webui" |
| 64 | / "launcher-gateway-store.js" |
| 65 | ) |
| 66 | commands_source = ( |
| 67 | root / "plugins" / "_commands" / "webui" / "commands-slash-store.js" |
| 68 | ).read_text(encoding="utf-8") |
| 69 | |
| 70 | assert not extension.exists() |
| 71 | assert not store.exists() |
| 72 | assert "a0LauncherHost" not in commands_source |
| 73 | assert "launcher-gateway-store.js" not in commands_source |
| 74 | assert 'type === "computer_use"' in commands_source |
| 75 | |
| 76 | |
| 77 | def test_launcher_gateway_is_fallback_after_context_bound_cli() -> None: |
| 78 | context_id = f"ctx-{uuid.uuid4()}" |
| 79 | cli_sid = _sid("cli") |
| 80 | gateway_sid = _sid("launcher") |
| 81 | ws_runtime.register_sid(cli_sid) |
| 82 | ws_runtime.register_sid(gateway_sid) |
| 83 | ws_runtime.subscribe_sid_to_context(cli_sid, context_id) |
| 84 | ws_runtime.store_sid_launcher_gateway_metadata(gateway_sid, _gateway("installation-a")) |
| 85 | try: |
| 86 | assert ws_runtime.remote_tool_sids_for_context(context_id)[:2] == [ |
| 87 | cli_sid, |
| 88 | gateway_sid, |
| 89 | ] |
| 90 | finally: |
| 91 | ws_runtime.unregister_sid(cli_sid) |
| 92 | ws_runtime.unregister_sid(gateway_sid) |
| 93 | |
| 94 | |
| 95 | def test_distinct_launcher_gateways_fail_closed() -> None: |
| 96 | first_sid = _sid("first") |
| 97 | second_sid = _sid("second") |
| 98 | ws_runtime.register_sid(first_sid) |
| 99 | ws_runtime.register_sid(second_sid) |
| 100 | ws_runtime.store_sid_launcher_gateway_metadata(first_sid, _gateway("installation-a")) |
| 101 | ws_runtime.store_sid_launcher_gateway_metadata(second_sid, _gateway("installation-b")) |
| 102 | try: |
| 103 | status = ws_runtime.launcher_gateway_status() |
| 104 | assert status["state"] == "multiple_hosts" |
| 105 | assert status["connected"] is False |
| 106 | assert ws_runtime.active_launcher_gateway_sid() is None |
| 107 | assert first_sid not in ws_runtime.remote_tool_sids_for_context("unbound") |
| 108 | assert second_sid not in ws_runtime.remote_tool_sids_for_context("unbound") |
| 109 | finally: |
| 110 | ws_runtime.unregister_sid(first_sid) |
| 111 | ws_runtime.unregister_sid(second_sid) |
| 112 | |
| 113 | |
| 114 | def test_duplicate_gateway_identity_replaces_stale_socket() -> None: |
| 115 | stale_sid = _sid("stale") |
| 116 | fresh_sid = _sid("fresh") |
| 117 | ws_runtime.register_sid(stale_sid) |
| 118 | ws_runtime.register_sid(fresh_sid) |
| 119 | ws_runtime.store_sid_launcher_gateway_metadata(stale_sid, _gateway("installation-a")) |
| 120 | ws_runtime.store_sid_launcher_gateway_metadata(fresh_sid, _gateway("installation-a")) |
| 121 | try: |
| 122 | assert ws_runtime.active_launcher_gateway_sid() == fresh_sid |
| 123 | assert ws_runtime.store_sid_launcher_gateway_metadata( |
| 124 | stale_sid, |
| 125 | _gateway("installation-a"), |
| 126 | ) is None |
| 127 | finally: |
| 128 | ws_runtime.unregister_sid(stale_sid) |
| 129 | ws_runtime.unregister_sid(fresh_sid) |
| 130 | |
| 131 | |
| 132 | def test_gateway_scope_dependencies_keep_reads_separate_from_writes() -> None: |
| 133 | sid = _sid("scope") |
| 134 | ws_runtime.register_sid(sid) |
| 135 | ws_runtime.store_sid_launcher_gateway_metadata(sid, _gateway("installation-a", files=False)) |
| 136 | try: |
| 137 | gateway = ws_runtime.launcher_gateway_status()["gateway"] |
| 138 | assert gateway["scopes"]["files"] is False |
| 139 | assert gateway["scopes"]["file_write"] is False |
| 140 | assert gateway["scopes"]["code_execution"] is False |
| 141 | finally: |
| 142 | ws_runtime.unregister_sid(sid) |
| 143 | |
| 144 | sid = _sid("read-only") |
| 145 | ws_runtime.register_sid(sid) |
| 146 | ws_runtime.store_sid_launcher_gateway_metadata( |
| 147 | sid, |
| 148 | _gateway("installation-a", file_write=False), |
| 149 | ) |
| 150 | try: |
| 151 | gateway = ws_runtime.launcher_gateway_status()["gateway"] |
| 152 | assert gateway["scopes"]["files"] is True |
| 153 | assert gateway["scopes"]["file_write"] is False |
| 154 | assert gateway["scopes"]["code_execution"] is False |
| 155 | finally: |
| 156 | ws_runtime.unregister_sid(sid) |
| 157 | |
| 158 | |
| 159 | def test_legacy_gateway_files_scope_keeps_read_write_behavior() -> None: |
| 160 | sid = _sid("legacy") |
| 161 | payload = _gateway("installation-a") |
| 162 | payload["scopes"].pop("file_write") |
| 163 | ws_runtime.register_sid(sid) |
| 164 | ws_runtime.store_sid_launcher_gateway_metadata(sid, payload) |
| 165 | try: |
| 166 | assert ws_runtime.launcher_gateway_status()["gateway"]["scopes"]["file_write"] is True |
| 167 | finally: |
| 168 | ws_runtime.unregister_sid(sid) |
| 169 | |
| 170 | |
| 171 | def test_gateway_status_metadata_is_bounded() -> None: |
| 172 | sid = _sid("bounded") |
| 173 | payload = _gateway("installation-a") |
| 174 | payload["status"] = { |
| 175 | "browser": { |
| 176 | "message": "x" * 4000, |
| 177 | "available_browsers": [{"browser_id": f"browser-{index}"} for index in range(100)], |
| 178 | }, |
| 179 | "computer_use": { |
| 180 | "capabilities": {"elements": {"tree_backends": ["ax", "at-spi"]}} |
| 181 | }, |
| 182 | } |
| 183 | ws_runtime.register_sid(sid) |
| 184 | ws_runtime.store_sid_launcher_gateway_metadata(sid, payload) |
| 185 | try: |
| 186 | status = ws_runtime.launcher_gateway_status()["gateway"]["status"] |
| 187 | assert len(status["browser"]["message"]) == 2048 |
| 188 | assert len(status["browser"]["available_browsers"]) == 64 |
| 189 | assert status["computer_use"]["capabilities"]["elements"]["tree_backends"] == [ |
| 190 | "ax", |
| 191 | "at-spi", |
| 192 | ] |
| 193 | finally: |
| 194 | ws_runtime.unregister_sid(sid) |
| 195 | |
| 196 | |
| 197 | def test_gateway_control_requires_csrf_and_waits_for_ack(monkeypatch) -> None: |
| 198 | sid = _sid("control") |
| 199 | ws_runtime.register_sid(sid) |
| 200 | ws_runtime.store_sid_launcher_gateway_metadata(sid, _gateway("installation-a")) |
| 201 | |
| 202 | class FakeManager: |
| 203 | async def emit_to(self, namespace, target_sid, event, data, **kwargs): |
| 204 | assert namespace == "/ws" |
| 205 | assert target_sid == sid |
| 206 | assert event == "connector_gateway_control" |
| 207 | updated = _gateway("installation-a") |
| 208 | updated["master_enabled"] = False |
| 209 | updated["state"] = "paused" |
| 210 | ws_runtime.resolve_pending_gateway_control( |
| 211 | data["request_id"], |
| 212 | sid=sid, |
| 213 | payload={ |
| 214 | "request_id": data["request_id"], |
| 215 | "ok": True, |
| 216 | "gateway": updated, |
| 217 | }, |
| 218 | ) |
| 219 | |
| 220 | monkeypatch.setattr(launcher_gateway_control, "get_shared_ws_manager", lambda: FakeManager()) |
| 221 | handler = launcher_gateway_control.LauncherGatewayControl(None, None) |
| 222 | try: |
| 223 | result = asyncio.run(handler.process({"action": "set_master", "enabled": False}, None)) |
| 224 | assert handler.requires_auth() is True |
| 225 | assert handler.requires_csrf() is True |
| 226 | assert result["ok"] is True |
| 227 | assert result["status"]["gateway"]["master_enabled"] is False |
| 228 | finally: |
| 229 | ws_runtime.unregister_sid(sid) |
| 230 | |
| 231 | |
| 232 | def test_gateway_scope_ack_updates_file_routing_before_follow_up_hello(monkeypatch) -> None: |
| 233 | sid = _sid("scope-transition") |
| 234 | ws_runtime.register_sid(sid) |
| 235 | ws_runtime.store_sid_launcher_gateway_metadata(sid, _gateway("installation-a")) |
| 236 | ws_runtime.store_sid_remote_file_metadata( |
| 237 | sid, |
| 238 | {"enabled": True, "write_enabled": True, "mode": "read_write"}, |
| 239 | ) |
| 240 | ws_runtime.store_sid_remote_exec_metadata(sid, {"enabled": True}) |
| 241 | |
| 242 | class FakeManager: |
| 243 | async def emit_to(self, _namespace, target_sid, _event, data, **_kwargs): |
| 244 | assert target_sid == sid |
| 245 | updated = _gateway("installation-a", file_write=False) |
| 246 | ws_runtime.resolve_pending_gateway_control( |
| 247 | data["request_id"], |
| 248 | sid=sid, |
| 249 | payload={ |
| 250 | "request_id": data["request_id"], |
| 251 | "ok": True, |
| 252 | "gateway": updated, |
| 253 | }, |
| 254 | ) |
| 255 | |
| 256 | monkeypatch.setattr(launcher_gateway_control, "get_shared_ws_manager", lambda: FakeManager()) |
| 257 | handler = launcher_gateway_control.LauncherGatewayControl(None, None) |
| 258 | try: |
| 259 | result = asyncio.run( |
| 260 | handler.process( |
| 261 | { |
| 262 | "action": "replace_scopes", |
| 263 | "scopes": _gateway("installation-a", file_write=False)["scopes"], |
| 264 | }, |
| 265 | None, |
| 266 | ) |
| 267 | ) |
| 268 | assert result["ok"] is True |
| 269 | assert ws_runtime.select_remote_file_target_sid("ctx", require_writes=False) == sid |
| 270 | assert ws_runtime.select_remote_file_target_sid("ctx", require_writes=True) is None |
| 271 | assert ws_runtime.select_remote_exec_target_sid("ctx") is None |
| 272 | finally: |
| 273 | ws_runtime.unregister_sid(sid) |
| 274 | |
| 275 | |
| 276 | def test_gateway_scope_control_requires_explicit_file_write() -> None: |
| 277 | handler = launcher_gateway_control.LauncherGatewayControl(None, None) |
| 278 | result = asyncio.run( |
| 279 | handler.process( |
| 280 | { |
| 281 | "action": "replace_scopes", |
| 282 | "scopes": { |
| 283 | "files": True, |
| 284 | "code_execution": True, |
| 285 | "browser": False, |
| 286 | "computer_use": False, |
| 287 | }, |
| 288 | }, |
| 289 | None, |
| 290 | ) |
| 291 | ) |
| 292 | assert result.status_code == 400 |
| 293 | |
| 294 | |
| 295 | def test_gateway_control_acknowledgement_timeout(monkeypatch) -> None: |
| 296 | sid = _sid("timeout") |
| 297 | ws_runtime.register_sid(sid) |
| 298 | ws_runtime.store_sid_launcher_gateway_metadata(sid, _gateway("installation-a")) |
| 299 | |
| 300 | class SilentManager: |
| 301 | async def emit_to(self, *_args, **_kwargs): |
| 302 | return None |
| 303 | |
| 304 | monkeypatch.setattr(launcher_gateway_control, "get_shared_ws_manager", lambda: SilentManager()) |
| 305 | monkeypatch.setattr(launcher_gateway_control, "_CONTROL_TIMEOUT_SECONDS", 0.01) |
| 306 | handler = launcher_gateway_control.LauncherGatewayControl(None, None) |
| 307 | try: |
| 308 | result = asyncio.run(handler.process({"action": "set_master", "enabled": False}, None)) |
| 309 | assert result.status_code == 504 |
| 310 | finally: |
| 311 | ws_runtime.unregister_sid(sid) |
| 312 | |
| 313 | |
| 314 | def test_gateway_emergency_disconnect_returns_acknowledged_disconnected_state(monkeypatch) -> None: |
| 315 | sid = _sid("emergency") |
| 316 | ws_runtime.register_sid(sid) |
| 317 | ws_runtime.store_sid_launcher_gateway_metadata(sid, _gateway("installation-a")) |
| 318 | |
| 319 | class FakeManager: |
| 320 | async def emit_to(self, _namespace, target_sid, event, data, **_kwargs): |
| 321 | assert target_sid == sid |
| 322 | assert event == "connector_gateway_control" |
| 323 | assert data["action"] == "emergency_disconnect" |
| 324 | updated = _gateway("installation-a") |
| 325 | updated["state"] = "disconnected" |
| 326 | ws_runtime.resolve_pending_gateway_control( |
| 327 | data["request_id"], |
| 328 | sid=sid, |
| 329 | payload={ |
| 330 | "request_id": data["request_id"], |
| 331 | "ok": True, |
| 332 | "gateway": updated, |
| 333 | }, |
| 334 | ) |
| 335 | |
| 336 | monkeypatch.setattr(launcher_gateway_control, "get_shared_ws_manager", lambda: FakeManager()) |
| 337 | handler = launcher_gateway_control.LauncherGatewayControl(None, None) |
| 338 | try: |
| 339 | result = asyncio.run(handler.process({"action": "emergency_disconnect"}, None)) |
| 340 | assert result["ok"] is True |
| 341 | assert result["status"]["state"] == "disconnected" |
| 342 | finally: |
| 343 | ws_runtime.unregister_sid(sid) |