diff --git a/changelog.d/tsk-4hzmaf-orb-agent-buddy-probe.md b/changelog.d/tsk-4hzmaf-orb-agent-buddy-probe.md new file mode 100644 index 000000000..d1304615b --- /dev/null +++ b/changelog.d/tsk-4hzmaf-orb-agent-buddy-probe.md @@ -0,0 +1,3 @@ +### Added + +- `tinyagentos.orb.discovery`: Agent Buddy UDP 4666 discovery probe (`probe_agent_buddy`, `parse_agent_buddy_reply`, `AGENT_BUDDY_PORT`, `AGENT_BUDDY_PROBE`) for ESP32 agent companion devices. \ No newline at end of file diff --git a/tests/orb/test_discovery.py b/tests/orb/test_discovery.py index 5a5c1d5e4..474f468a8 100644 --- a/tests/orb/test_discovery.py +++ b/tests/orb/test_discovery.py @@ -391,4 +391,453 @@ def test_service_browser_init_raises_error_propagates_close_called(fake_zeroconf assert str(e) == "browser init failed" assert fake_zeroconf._close_called else: - assert False, "Expected RuntimeError" \ No newline at end of file + assert False, "Expected RuntimeError" + + +# TestAgentBuddyDiscovery tests + +AGENT_BUDDY_PROBE = b"ESP32_AGENT_COMPANION_DISCOVER_V1" + + +def _make_valid_reply(device_id="ABCDEF1234567890", hostname="test-device", port=8080, boot=42): + """Create a valid Agent Buddy discovery reply JSON.""" + return { + "service": "esp32-agent-companion", + "protocol": 1, + "deviceId": device_id, + "hostname": hostname, + "port": port, + "boot": boot, + } + + +def _make_reply_bytes(reply_dict): + """Convert a reply dict to bytes.""" + import json + return json.dumps(reply_dict).encode("utf-8") + + +def test_parse_agent_buddy_reply_valid(): + """Test parse_agent_buddy_reply with a valid reply.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply(device_id="ABCDEF1234567890", hostname="test-device", port=8080, boot=42) + data = _make_reply_bytes(reply) + + result = parse_agent_buddy_reply(data, "192.168.1.100") + + assert result is not None + assert result["device_id"] == "abcdef1234567890" # lowercased + assert result["hostname"] == "test-device" + assert result["address"] == "192.168.1.100" + assert result["port"] == 8080 + assert result["boot"] == 42 + + +def test_parse_agent_buddy_reply_not_json(): + """Test parse_agent_buddy_reply returns None for non-JSON data.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + data = b"not json at all" + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_json_list(): + """Test parse_agent_buddy_reply returns None for JSON list.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + data = b"[1, 2, 3]" + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_wrong_service(): + """Test parse_agent_buddy_reply returns None for wrong service.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["service"] = "wrong-service" + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_protocol_2(): + """Test parse_agent_buddy_reply returns None for protocol 2.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["protocol"] = 2 + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_protocol_true(): + """Test parse_agent_buddy_reply returns None for protocol True (bool).""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["protocol"] = True + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_device_id_15_hex(): + """Test parse_agent_buddy_reply returns None for 15-char deviceId.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply(device_id="ABCDEF123456789") # 15 chars + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_device_id_non_hex(): + """Test parse_agent_buddy_reply returns None for non-hex deviceId.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply(device_id="ABCDEF123456789G") # G is not hex + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_hostname_not_str(): + """Test parse_agent_buddy_reply returns None for non-str hostname.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["hostname"] = 5 + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_port_0(): + """Test parse_agent_buddy_reply returns None for port 0.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply(port=0) + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_port_65536(): + """Test parse_agent_buddy_reply returns None for port 65536.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply(port=65536) + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_port_str(): + """Test parse_agent_buddy_reply returns None for port as string.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["port"] = "80" + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_port_true(): + """Test parse_agent_buddy_reply returns None for port True (bool).""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["port"] = True + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_boot_negative(): + """Test parse_agent_buddy_reply returns None for boot -1.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply(boot=-1) + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_boot_missing(): + """Test parse_agent_buddy_reply returns None for missing boot.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + del reply["boot"] + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_extra_key_ignored(): + """Test parse_agent_buddy_reply ignores extra unknown keys.""" + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + reply = _make_valid_reply() + reply["extra_unknown_key"] = "should be ignored" + data = _make_reply_bytes(reply) + result = parse_agent_buddy_reply(data, "192.168.1.100") + + assert result is not None + assert result["device_id"] == "abcdef1234567890" + assert result["hostname"] == "test-device" + assert result["address"] == "192.168.1.100" + assert result["port"] == 8080 + assert result["boot"] == 42 + assert "extra_unknown_key" not in result + + +def test_parse_agent_buddy_reply_deeply_nested_returns_none(): + """A datagram of 2000 '[' nests past the recursion limit. + + On python 3.11 json.loads raises RecursionError (3.12+ raise JSONDecodeError), + so this test is red only on 3.11 before the fix. The monkeypatched test below + covers the except clause on every interpreter. + """ + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + result = parse_agent_buddy_reply(b"[" * 2000, "192.168.1.100") + assert result is None + + +def test_parse_agent_buddy_reply_recursion_error_returns_none(monkeypatch): + """A RecursionError out of json.loads is a bad datagram, not a probe abort.""" + import json + from tinyagentos.orb.discovery import parse_agent_buddy_reply + + def raise_recursion(_text): + raise RecursionError("maximum recursion depth exceeded") + + monkeypatch.setattr(json, "loads", raise_recursion) + result = parse_agent_buddy_reply(b"{}", "192.168.1.100") + assert result is None + + +# Probe tests using real loopback UDP + + +def _start_responder(reply_bytes_list, port): + """Start a responder thread that replies to AGENT_BUDDY_PROBE with configured datagrams.""" + import socket + import threading + + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.bind(("127.0.0.1", port)) + stop_event = threading.Event() + received_probes = [] + + def run(): + while not stop_event.is_set(): + try: + sock.settimeout(0.1) + data, addr = sock.recvfrom(2048) + received_probes.append((data, addr)) + if data == AGENT_BUDDY_PROBE: + for reply in reply_bytes_list: + sock.sendto(reply, addr) + except socket.timeout: + continue + except OSError: + break + + thread = threading.Thread(target=run, daemon=True) + thread.start() + + def stop(): + stop_event.set() + sock.close() + thread.join(timeout=1.0) + + return stop, received_probes + + +def test_probe_agent_buddy_one_valid_reply(): + """Test probe_agent_buddy with one valid reply.""" + import time + from tinyagentos.orb.discovery import probe_agent_buddy, AGENT_BUDDY_PROBE + + reply = _make_valid_reply(device_id="ABCDEF1234567890", hostname="test-device", port=8080, boot=42) + reply_bytes = _make_reply_bytes(reply) + + # Bind to port 0 to get a free port + import socket + temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + temp_sock.bind(("127.0.0.1", 0)) + responder_port = temp_sock.getsockname()[1] + temp_sock.close() + + stop, received = _start_responder([reply_bytes], responder_port) + try: + start = time.monotonic() + result = probe_agent_buddy(timeout=0.5, addresses=["127.0.0.1"], port=responder_port) + elapsed = time.monotonic() - start + finally: + stop() + + assert elapsed < 1.0 # should return within timeout + 0.5s + assert len(result) == 1 + assert result[0]["device_id"] == "abcdef1234567890" + assert result[0]["hostname"] == "test-device" + assert result[0]["address"] == "127.0.0.1" + assert result[0]["port"] == 8080 + assert result[0]["boot"] == 42 + # Verify the probe was sent + assert any(data == AGENT_BUDDY_PROBE for data, _ in received) + + +def test_probe_agent_buddy_two_valid_replies_sorted(): + """Test probe_agent_buddy with two valid replies, sorted by hostname.""" + import time + from tinyagentos.orb.discovery import probe_agent_buddy, AGENT_BUDDY_PROBE + + reply1 = _make_valid_reply(device_id="AAAAAAAAAAAAAAAA", hostname="zebra-device", port=8080, boot=1) + reply2 = _make_valid_reply(device_id="BBBBBBBBBBBBBBBB", hostname="alpha-device", port=8081, boot=2) + # Add an invalid reply (wrong service) + invalid_reply = _make_valid_reply(device_id="CCCCCCCCCCCCCCCC", hostname="invalid", port=8082, boot=3) + invalid_reply["service"] = "wrong-service" + invalid_bytes = _make_reply_bytes(invalid_reply) + + import socket + temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + temp_sock.bind(("127.0.0.1", 0)) + responder_port = temp_sock.getsockname()[1] + temp_sock.close() + + stop, received = _start_responder([_make_reply_bytes(reply1), _make_reply_bytes(reply2), invalid_bytes], responder_port) + try: + start = time.monotonic() + result = probe_agent_buddy(timeout=0.5, addresses=["127.0.0.1"], port=responder_port) + elapsed = time.monotonic() - start + finally: + stop() + + assert elapsed < 1.0 + assert len(result) == 2 + # Should be sorted by hostname: alpha-device, then zebra-device + assert result[0]["hostname"] == "alpha-device" + assert result[0]["device_id"] == "bbbbbbbbbbbbbbbb" + assert result[0]["port"] == 8081 + assert result[1]["hostname"] == "zebra-device" + assert result[1]["device_id"] == "aaaaaaaaaaaaaaaa" + assert result[1]["port"] == 8080 + assert any(data == AGENT_BUDDY_PROBE for data, _ in received) + + +def test_probe_agent_buddy_same_device_id_replaces(): + """Test probe_agent_buddy with two replies for same device_id keeps second.""" + import time + from tinyagentos.orb.discovery import probe_agent_buddy, AGENT_BUDDY_PROBE + + # First reply with port 8080 + reply1 = _make_valid_reply(device_id="AAAAAAAAAAAAAAAA", hostname="test-device", port=8080, boot=1) + # Second reply with same device_id but different port + reply2 = _make_valid_reply(device_id="AAAAAAAAAAAAAAAA", hostname="test-device", port=9090, boot=2) + + import socket + temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + temp_sock.bind(("127.0.0.1", 0)) + responder_port = temp_sock.getsockname()[1] + temp_sock.close() + + stop, received = _start_responder([_make_reply_bytes(reply1), _make_reply_bytes(reply2)], responder_port) + try: + start = time.monotonic() + result = probe_agent_buddy(timeout=0.5, addresses=["127.0.0.1"], port=responder_port) + elapsed = time.monotonic() - start + finally: + stop() + + assert elapsed < 1.0 + assert len(result) == 1 + # Should keep the second reply's port + assert result[0]["device_id"] == "aaaaaaaaaaaaaaaa" + assert result[0]["port"] == 9090 + assert any(data == AGENT_BUDDY_PROBE for data, _ in received) + + +def test_probe_agent_buddy_ignores_wrong_probe(): + """Test probe_agent_buddy responder ignores non-probe datagrams.""" + import time + from tinyagentos.orb.discovery import probe_agent_buddy, AGENT_BUDDY_PROBE + + reply = _make_valid_reply(device_id="ABCDEF1234567890", hostname="test-device", port=8080, boot=42) + + import socket + temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + temp_sock.bind(("127.0.0.1", 0)) + responder_port = temp_sock.getsockname()[1] + temp_sock.close() + + stop, received = _start_responder([_make_reply_bytes(reply)], responder_port) + try: + # Send a wrong probe first + wrong_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + wrong_sock.sendto(b"WRONG_PROBE", ("127.0.0.1", responder_port)) + wrong_sock.close() + + start = time.monotonic() + result = probe_agent_buddy(timeout=0.5, addresses=["127.0.0.1"], port=responder_port) + elapsed = time.monotonic() - start + finally: + stop() + + assert elapsed < 1.0 + assert len(result) == 1 + # Responder should have received the wrong probe but only replied to the correct one + assert any(data == b"WRONG_PROBE" for data, _ in received) + assert any(data == AGENT_BUDDY_PROBE for data, _ in received) + + +def test_probe_agent_buddy_no_reply(): + """Test probe_agent_buddy with no responder reply returns empty list within timeout.""" + import time + from tinyagentos.orb.discovery import probe_agent_buddy + + import socket + temp_sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + temp_sock.bind(("127.0.0.1", 0)) + responder_port = temp_sock.getsockname()[1] + temp_sock.close() + + # Start responder but it won't reply to anything (empty reply list) + stop, _ = _start_responder([], responder_port) + try: + start = time.monotonic() + result = probe_agent_buddy(timeout=0.5, addresses=["127.0.0.1"], port=responder_port) + elapsed = time.monotonic() - start + finally: + stop() + + assert elapsed < 1.0 # timeout (0.5) + 0.5 = 1.0 + assert result == [] + + +def test_probe_agent_buddy_socket_error_raises_transport_unavailable(monkeypatch): + """Test probe_agent_buddy raises TransportUnavailable when socket creation fails.""" + import socket + from tinyagentos.orb.discovery import probe_agent_buddy, TransportUnavailable + + def raise_oserror(*args, **kwargs): + raise OSError("socket creation failed") + + monkeypatch.setattr(socket, "socket", raise_oserror) + + try: + probe_agent_buddy(timeout=0.5, addresses=["127.0.0.1"]) + except TransportUnavailable as e: + assert "Failed to create or bind UDP socket" in str(e) + # Check it's chained from OSError + assert e.__cause__ is not None + assert isinstance(e.__cause__, OSError) + else: + assert False, "Expected TransportUnavailable" \ No newline at end of file diff --git a/tinyagentos/orb/discovery.py b/tinyagentos/orb/discovery.py index db5584d67..c72fdea5d 100644 --- a/tinyagentos/orb/discovery.py +++ b/tinyagentos/orb/discovery.py @@ -6,6 +6,10 @@ ORB_SERVICE_TYPE = "_taos-orb._tcp.local." +# Agent Buddy UDP discovery +AGENT_BUDDY_PORT = 4666 +AGENT_BUDDY_PROBE = b"ESP32_AGENT_COMPANION_DISCOVER_V1" + def _parse_service_info(info) -> dict: """Parse a zeroconf ServiceInfo-like object into a dict.""" @@ -95,4 +99,116 @@ def update_service(self, zc, type_, name): results.sort(key=lambda x: x["name"]) return results finally: - zc.close() \ No newline at end of file + zc.close() + + +def parse_agent_buddy_reply(data: bytes, address: str) -> dict | None: + """Parse an Agent Buddy discovery reply. + + Returns a dict with device_id, hostname, address, port, boot on success. + Returns None on any decode/parse error or if validation fails. + """ + import json + import re + + try: + text = data.decode("utf-8") + parsed = json.loads(text) + except (UnicodeDecodeError, json.JSONDecodeError, RecursionError): + return None + + if not isinstance(parsed, dict): + return None + + # Validate required fields + if parsed.get("service") != "esp32-agent-companion": + return None + + protocol = parsed.get("protocol") + if not isinstance(protocol, int) or isinstance(protocol, bool) or protocol != 1: + return None + + device_id = parsed.get("deviceId") + if not isinstance(device_id, str) or not re.fullmatch(r"[0-9a-fA-F]{16}", device_id): + return None + + hostname = parsed.get("hostname") + if not isinstance(hostname, str): + return None + + port = parsed.get("port") + if not isinstance(port, int) or isinstance(port, bool) or not (1 <= port <= 65535): + return None + + boot = parsed.get("boot") + if not isinstance(boot, int) or isinstance(boot, bool) or boot < 0: + return None + + return { + "device_id": device_id.lower(), + "hostname": hostname, + "address": address, + "port": port, + "boot": boot, + } + + +def probe_agent_buddy(timeout: float = 1.2, addresses=None, port: int = AGENT_BUDDY_PORT) -> list[dict]: + """Probe for Agent Buddy devices via UDP broadcast. + + Sends a discovery probe to each address and collects replies for `timeout` seconds. + Returns a list of device dicts sorted by (hostname, device_id). + + Async callers must use asyncio.to_thread. + """ + import socket + import time + + if addresses is None: + addresses = ["255.255.255.255"] + + sock = None + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1) + sock.bind(("", 0)) + except OSError as e: + if sock is not None: + sock.close() + raise TransportUnavailable("Failed to create or bind UDP socket") from e + + try: + # Send probe to each address + for addr in addresses: + try: + sock.sendto(AGENT_BUDDY_PROBE, (addr, port)) + except OSError: + # Skip failed sends, continue with other addresses + pass + + # Receive replies until timeout + deadline = time.monotonic() + timeout + results = {} + + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + sock.settimeout(remaining) + try: + data, (remote_ip, _) = sock.recvfrom(2048) + except socket.timeout: + break + except OSError: + break + + parsed = parse_agent_buddy_reply(data, remote_ip) + if parsed is not None: + results[parsed["device_id"]] = parsed + + # Sort by (hostname, device_id) + sorted_results = sorted(results.values(), key=lambda x: (x["hostname"], x["device_id"])) + return sorted_results + finally: + if sock is not None: + sock.close() \ No newline at end of file