From cc9bb81bc9a0473dbda12c15d3f03f7e799b909d Mon Sep 17 00:00:00 2001
From: Brian Krabach <brkrabac@microsoft.com>
Date: Wed, 15 Apr 2026 11:45:06 -0700
Subject: [PATCH] feat: add _lookup_remote_by_device_id helper for federation
 proxy

Implements _lookup_remote_by_device_id(device_id: str) -> dict | None
in muxplex/main.py, inserted before the federation WebSocket proxy section.

Logic:
- Load settings and get remote_instances list
- Primary: iterate remotes, return first where remote.get('device_id') == device_id
- Fallback: if device_id parses as integer, use index-based lookup for
  transition compatibility (0 <= idx < len(remotes))
- Return None if not found

Tests added:
- test_lookup_remote_by_device_id_found: two remotes with device_ids
  'aaa-111' and 'bbb-222'; lookup 'bbb-222' returns Desktop remote
- test_lookup_remote_by_device_id_not_found: one remote with device_id
  'aaa-111'; lookup 'zzz-999' returns None

Task: task-8
---
 muxplex/main.py           | 37 ++++++++++++++++++
 muxplex/tests/test_api.py | 81 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 118 insertions(+)

diff --git a/muxplex/main.py b/muxplex/main.py
index 07fe54e..aa8bf1b 100644
--- a/muxplex/main.py
+++ b/muxplex/main.py
@@ -1005,6 +1005,43 @@ async def terminal_ws_proxy(websocket: WebSocket) -> None:
             pass
 
 
+# ---------------------------------------------------------------------------
+# Federation helper utilities
+# ---------------------------------------------------------------------------
+
+
+def _lookup_remote_by_device_id(device_id: str) -> dict | None:
+    """Return the first remote instance whose ``device_id`` matches *device_id*.
+
+    Primary lookup: iterate ``remote_instances`` and return the first entry
+    where ``remote.get('device_id') == device_id``.
+
+    Fallback (transition compatibility): if *device_id* looks like an integer
+    (i.e. ``int(device_id)`` succeeds) treat it as a 0-based index into the
+    ``remote_instances`` list and return the remote at that position, provided
+    the index is in range.
+
+    Returns ``None`` if no match is found.
+    """
+    settings = load_settings()
+    remotes: list[dict] = settings.get("remote_instances", [])
+
+    # Primary: match by device_id field
+    for remote in remotes:
+        if remote.get("device_id") == device_id:
+            return remote
+
+    # Fallback: index-based lookup for transition compatibility
+    try:
+        idx = int(device_id)
+        if 0 <= idx < len(remotes):
+            return remotes[idx]
+    except (ValueError, TypeError):
+        pass
+
+    return None
+
+
 # ---------------------------------------------------------------------------
 # Federation WebSocket proxy — bridges browser to a remote instance's ttyd
 # ---------------------------------------------------------------------------
diff --git a/muxplex/tests/test_api.py b/muxplex/tests/test_api.py
index e019e0a..aec13f8 100644
--- a/muxplex/tests/test_api.py
+++ b/muxplex/tests/test_api.py
@@ -3284,3 +3284,84 @@ def test_fetch_remote_marks_unreachable_after_grace_period(
     assert any(s.get("status") == "unreachable" for s in host_entries), (
         f"After exceeding grace period, device must be marked 'unreachable'. Got: {host_entries}"
     )
+
+
+# ---------------------------------------------------------------------------
+# Tests for _lookup_remote_by_device_id helper
+# ---------------------------------------------------------------------------
+
+
+def test_lookup_remote_by_device_id_found(tmp_path, monkeypatch):
+    """_lookup_remote_by_device_id returns the remote dict matching the given device_id."""
+    import json
+
+    import muxplex.settings as settings_mod
+
+    from muxplex.main import _lookup_remote_by_device_id
+
+    settings_path = tmp_path / "settings.json"
+    monkeypatch.setattr(settings_mod, "SETTINGS_PATH", settings_path)
+
+    settings_path.write_text(
+        json.dumps(
+            {
+                "remote_instances": [
+                    {
+                        "url": "http://laptop:8088",
+                        "key": "key-aaa",
+                        "name": "Laptop",
+                        "device_id": "aaa-111",
+                    },
+                    {
+                        "url": "http://desktop:8088",
+                        "key": "key-bbb",
+                        "name": "Desktop",
+                        "device_id": "bbb-222",
+                    },
+                ]
+            }
+        )
+    )
+
+    result = _lookup_remote_by_device_id("bbb-222")
+
+    assert result is not None, "Expected a remote dict, got None"
+    assert result.get("name") == "Desktop", (
+        f"Expected 'Desktop' remote, got: {result!r}"
+    )
+    assert result.get("device_id") == "bbb-222", (
+        f"Expected device_id 'bbb-222', got: {result.get('device_id')!r}"
+    )
+
+
+def test_lookup_remote_by_device_id_not_found(tmp_path, monkeypatch):
+    """_lookup_remote_by_device_id returns None when no remote matches the given device_id."""
+    import json
+
+    import muxplex.settings as settings_mod
+
+    from muxplex.main import _lookup_remote_by_device_id
+
+    settings_path = tmp_path / "settings.json"
+    monkeypatch.setattr(settings_mod, "SETTINGS_PATH", settings_path)
+
+    settings_path.write_text(
+        json.dumps(
+            {
+                "remote_instances": [
+                    {
+                        "url": "http://laptop:8088",
+                        "key": "key-aaa",
+                        "name": "Laptop",
+                        "device_id": "aaa-111",
+                    },
+                ]
+            }
+        )
+    )
+
+    result = _lookup_remote_by_device_id("zzz-999")
+
+    assert result is None, (
+        f"Expected None for unknown device_id 'zzz-999', got: {result!r}"
+    )