improve error message

Author: mlikasam-askui

src/askui/tools/askui/agent_os_server_manager.py

@@ -139,6 +139,19 @@ def get(self, computer_id: str) -> AgentOsServer:
         """
         return self._servers[self._index_of(computer_id)]
 
+    def get_by_session_guid(self, session_guid: str) -> AgentOsServer | None:
+        """
+        Return the registered server with the given `session_guid`, or `None` if
+        no such server is registered.
+
+        Intended for internal lookups (e.g. mapping a gRPC session GUID back to
+        its server during teardown). For user-facing selection, prefer `get`.
+        """
+        for server in self._servers:
+            if server.session_guid == session_guid:
+                return server
+        return None
+
     def switch(self, computer_id: str) -> AgentOsServer:
         """
         Set the active server by its `computer_id`.
@@ -161,10 +174,7 @@ def active(self) -> AgentOsServer | None:
         """The currently active server, or `None` if no servers are registered."""
         if self._active_session_guid is None:
             return None
-        for server in self._servers:
-            if server.session_guid == self._active_session_guid:
-                return server
-        return None
+        return self.get_by_session_guid(self._active_session_guid)
 
     def __len__(self) -> int:
         return len(self._servers)

src/askui/tools/askui/askui_controller.py

@@ -386,12 +386,20 @@ def _connect_server(self, server: AgentOsServer) -> None:
             stub.SetActiveDisplay(
                 controller_v1_pbs.Request_SetActiveDisplay(displayID=server.display)
             )
-        except Exception:
+        except Exception as e:
             try:
                 channel.close()
             finally:
                 if started_process:
                     server.stop()
+            if hasattr(e, "add_note"):
+                e.add_note(
+                    f"While connecting to Agent OS server {server.description!r} "
+                    f"(computer_id={server.computer_id!r}, "
+                    f"session_guid={server.session_guid}, "
+                    f"display={server.display}, "
+                    f"address={server.address})"
+                )
             raise
         self._connections[server.session_guid] = _Connection(
             channel=channel,
@@ -484,10 +492,7 @@ def _disconnect_server(self, session_guid: str) -> None:
         except Exception:  # noqa: BLE001
             logger.exception("Error closing channel for controller %s", session_guid)
         if conn.started_process:
-            server = next(
-                (s for s in self._manager.list() if s.session_guid == session_guid),
-                None,
-            )
+            server = self._manager.get_by_session_guid(session_guid)
             if server is None:
                 return
             try: