From 7da9814ec08a7d08fcbdefc6eb1a54af91740ece Mon Sep 17 00:00:00 2001
From: Brad Barnett <127794626+bdbarnett@users.noreply.github.com>
Date: Tue, 27 Oct 2026 12:59:15 -0500
Subject: [PATCH] Roku displays close after multimer's keep-alive loop, not
 before

A timer-driven app's script body ends while the app is still running, and
multimer keeps it alive from its own exit hook. RokuDisplay and
RokuDisplayWrapper registered close() with atexit on their first frame, so
close() ran before that loop and quit the channel while the app was still
drawing. When the first frame came from a timer instead, the hook was
registered during exit and never ran at all.

close() now goes to multimer's stop hooks when multimer is keeping the
program alive, and to atexit otherwise. show() also returns early while
close() is running: a timer could run it inside close()'s blocking ECP
request, relaunch the channel and leave _launched set, so the next frame
server was never fetched.

Fixes #245
---
 lib/utils/roku_companion.py |  76 +++++++++++++-----
 tests/test_roku_close.py    | 148 ++++++++++++++++++++++++++++++++++++
 2 files changed, 204 insertions(+), 20 deletions(-)
 create mode 100644 tests/test_roku_close.py

diff --git a/lib/utils/roku_companion.py b/lib/utils/roku_companion.py
index 95675cf8..f56cb9af 100644
--- a/lib/utils/roku_companion.py
+++ b/lib/utils/roku_companion.py
@@ -333,6 +333,22 @@ def _at_exit(fn):
         pass
 
 
+def _at_stop(fn):
+    """Call fn when the program stops. Where multimer keeps the program alive
+    past its script body, that's after its keep-alive loop: an atexit hook
+    registered now would run first and close a timer-driven app while it is
+    still running. Otherwise, at exit."""
+    try:
+        from multimer import _hostloop
+
+        if _hostloop.strategy() is not None:
+            _hostloop.on_stop(fn)
+            return
+    except (ImportError, AttributeError):
+        pass
+    _at_exit(fn)
+
+
 def _session_url(url):
     """*url* with a token unique to this session. The channel appends
     ``t=<frame number>`` to every fetch, counting from 1 each session, and the
@@ -839,6 +855,7 @@ def __init__(
         self._server = None
         self._png = None
         self._launched = False
+        self._closing = False
         self._shown = 0
         self._t0 = None
         self.key_repeat = key_repeat
@@ -850,6 +867,11 @@ def __init__(
 
     def show(self, _timer=None):
         """Encode the current framebuffer to PNG and update the Roku stream."""
+        if self._closing:
+            # A timer ran this show() while close() waits on the TV: launching
+            # now would leave the channel closed and _launched set, so the
+            # next frame server would never be fetched.
+            return
         if self._png is None:
             self._png = _png_encoder()
             self._t0 = _now()
@@ -871,7 +893,7 @@ def show(self, _timer=None):
             print("[RokuDisplay] Showing frames from %s" % frame_url)
             self.tv.frames(frame_url)
             self._launched = True
-            _at_exit(self.close)
+            _at_stop(self.close)
 
         if hasattr(self._server, "poll"):
             self._server.poll()
@@ -979,15 +1001,19 @@ def close(self):
         app leaves the TV on its home screen rather than frozen on its last
         frame.
         """
-        if self._launched:
-            self._launched = False
-            try:
-                self.tv.close_channel()
-            except OSError:
-                pass
-        if self._server is not None:
-            self._server.close()
-            self._server = None
+        self._closing = True
+        try:
+            if self._launched:
+                self._launched = False
+                try:
+                    self.tv.close_channel()
+                except OSError:
+                    pass
+            if self._server is not None:
+                self._server.close()
+                self._server = None
+        finally:
+            self._closing = False
 
     def _deinit(self):
         """DisplayDriver's cleanup hook, run by quit() and deinit()."""
@@ -1024,6 +1050,7 @@ def __init__(self, display, tv, port=8090, my_ip=None):
         self._server = None
         self._png = None
         self._launched = False
+        self._closing = False
         print(
             "[RokuDisplayWrapper] Initialized with display %s, stream host %s:%d"
             % (display, self.local_ip, port)
@@ -1031,6 +1058,11 @@ def __init__(self, display, tv, port=8090, my_ip=None):
 
     def show(self, *args, **kwargs):
         """Invoke underlying display show(), then stream the frame to the Roku TV."""
+        if self._closing:
+            # A timer ran this show() while close() waits on the TV: launching
+            # now would leave the channel closed and _launched set, so the
+            # next frame server would never be fetched.
+            return
         try:
             self.display.show(*args, **kwargs)
         except AttributeError:
@@ -1062,7 +1094,7 @@ def show(self, *args, **kwargs):
             print("[RokuDisplayWrapper] Showing frames from %s" % frame_url)
             self.tv.frames(frame_url)
             self._launched = True
-            _at_exit(self.close)
+            _at_stop(self.close)
 
         if hasattr(self._server, "poll"):
             self._server.poll()
@@ -1073,15 +1105,19 @@ def format(self):
 
     def close(self):
         """Close the TV's Companion channel, the frame server and the wrapped display."""
-        if self._launched:
-            self._launched = False
-            try:
-                self.tv.close_channel()
-            except OSError:
-                pass
-        if self._server is not None:
-            self._server.close()
-            self._server = None
+        self._closing = True
+        try:
+            if self._launched:
+                self._launched = False
+                try:
+                    self.tv.close_channel()
+                except OSError:
+                    pass
+            if self._server is not None:
+                self._server.close()
+                self._server = None
+        finally:
+            self._closing = False
         if hasattr(self.display, "close"):
             self.display.close()
 
diff --git a/tests/test_roku_close.py b/tests/test_roku_close.py
new file mode 100644
index 00000000..417e0aaf
--- /dev/null
+++ b/tests/test_roku_close.py
@@ -0,0 +1,148 @@
+"""When the Roku displays close: after multimer's keep-alive loop, and never twice over.
+
+A timer-driven app's script body ends while the app is still running, and
+multimer keeps it alive from its own exit hook. The displays used to close
+from ``atexit``, which runs the most recently registered hook first, so the
+channel closed while the app was still drawing. While ``close()`` waited on
+the TV, a timer could also run ``show()`` inside it and relaunch the channel.
+
+Run with: python tests/test_roku_close.py
+"""
+
+from pathlib import Path
+import sys
+import types
+import unittest
+from unittest import mock
+
+import _env  # noqa: F401
+
+_LIB = Path(__file__).resolve().parent.parent / "lib"
+if str(_LIB) not in sys.path:
+    sys.path.insert(0, str(_LIB))
+
+
+class _Fake:
+    def __init__(self, *args, **kwargs):
+        pass
+
+
+def _import_roku():
+    stub = None
+    try:
+        import framebuf  # noqa: F401
+    except ImportError:
+        stub = types.ModuleType("framebuf")
+        stub.RGB565 = 1
+        stub.FrameBuffer = _Fake
+        sys.modules["framebuf"] = stub
+    try:
+        from utils import roku_companion
+    finally:
+        if stub is not None:
+            sys.modules.pop("framebuf", None)
+    return roku_companion
+
+
+roku_companion = _import_roku()
+
+
+class _Encoder:
+    def encode(self, buf, w, h):
+        return b"png"
+
+
+_log = []
+
+
+class _Server:
+    def __init__(self, host="", port=0):
+        self.closed = False
+        _log.append("bind")
+
+    def update_frame(self, png):
+        pass
+
+    def close(self):
+        self.closed = True
+
+
+class _Inner:
+    width = 4
+    height = 2
+    buffer = bytearray(16)
+
+    def show(self):
+        pass
+
+
+class _TV:
+    """Records ECP calls; ``during_close`` runs inside close_channel(), as a
+    timer would while the request blocks."""
+
+    host = "127.0.0.1"
+
+    def __init__(self):
+        self.during_close = None
+
+    def frames(self, url):
+        _log.append("launch")
+
+    def close_channel(self):
+        _log.append("quit")
+        if self.during_close is not None:
+            self.during_close()
+
+
+class TestRokuClose(unittest.TestCase):
+    def setUp(self):
+        del _log[:]
+        patches = (
+            mock.patch.object(roku_companion, "_png_encoder", _Encoder),
+            mock.patch.object(roku_companion, "_FrameServer", _Server),
+        )
+        for p in patches:
+            p.start()
+            self.addCleanup(p.stop)
+
+    def _wrapper(self):
+        tv = _TV()
+        return roku_companion.RokuDisplayWrapper(_Inner(), tv, my_ip="127.0.0.1"), tv
+
+    def test_close_waits_for_multimer_keep_alive_loop(self):
+        from multimer import _hostloop
+
+        with mock.patch.object(_hostloop, "strategy", return_value="exit_hook"), mock.patch.object(
+            _hostloop, "on_stop"
+        ) as on_stop, mock.patch("atexit.register") as register:
+            display, _tv = self._wrapper()
+            display.show()
+        on_stop.assert_called_once_with(display.close)
+        register.assert_not_called()
+
+    def test_close_falls_back_to_atexit_without_a_keep_alive_loop(self):
+        from multimer import _hostloop
+
+        with mock.patch.object(_hostloop, "strategy", return_value=None), mock.patch.object(
+            _hostloop, "on_stop"
+        ) as on_stop, mock.patch("atexit.register") as register:
+            display, _tv = self._wrapper()
+            display.show()
+        register.assert_called_once_with(display.close)
+        on_stop.assert_not_called()
+
+    def test_show_inside_close_does_not_relaunch(self):
+        with mock.patch("atexit.register"):
+            display, tv = self._wrapper()
+            display.show()
+            tv.during_close = display.show
+            display.close()
+            # Once closed, the next frame binds a new server and launches the
+            # channel at it, the way the first frame did.
+            tv.during_close = None
+            display.show()
+        self.assertEqual(_log, ["bind", "launch", "quit", "bind", "launch"])
+
+
+if __name__ == "__main__":
+    unittest.main()