fix(python): retain returned effect cleanups

sdk/python/packages/flet/src/flet/components/component.py

@@ -341,7 +341,7 @@ def _run_render_effects(self):
                     or hook.deps != hook.prev_deps
                 )
                 if deps_changed:
-                    if callable(hook.cleanup):
+                    if callable(hook.cleanup) or callable(hook._returned_cleanup):
                         self._schedule_effect(hook, is_cleanup=True)
                     self._schedule_effect(hook, is_cleanup=False)
 
@@ -354,7 +354,7 @@ def _run_unmount_effects(self):
             logger.debug("%s._run_unmount_effects()", self)
         for hook in self._state.hooks:
             # all effects are running on unmount
-            if isinstance(hook, EffectHook) and callable(hook.cleanup):
+            if isinstance(hook, EffectHook) and hook.has_lifecycle_work():
                 self._schedule_effect(hook, is_cleanup=True)
         self._state.hooks.clear()
 

sdk/python/packages/flet/src/flet/components/hooks/use_effect.py

@@ -20,13 +20,21 @@ class EffectHook(Hook):
 
     setup: Callable[[], Any | Awaitable[Any]]
     cleanup: Callable[[], Any | Awaitable[Any]] | None = None
+    _returned_cleanup: Callable[[], Any | Awaitable[Any]] | None = None
     deps: list[Any] | None = None
     prev_deps: list[Any] | None = None
 
     # runtime
     _setup_task: asyncio.Task | None = None  # last scheduled setup task
     _cleanup_task: asyncio.Task | None = None  # last scheduled cleanup task
 
+    def has_lifecycle_work(self) -> bool:
+        return (
+            callable(self.cleanup)
+            or callable(self._returned_cleanup)
+            or (self._setup_task is not None and not self._setup_task.done())
+        )
+
     def cancel(self):
         """
         Cancels currently scheduled setup and cleanup tasks, if still pending.

sdk/python/packages/flet/src/flet/messaging/session.py

@@ -735,14 +735,17 @@ async def __updates_scheduler(self):
                                 hook._setup_task = asyncio.create_task(hook.setup())
                             else:
                                 res = hook.setup()
-                            if callable(res):
-                                hook.cleanup = res
-                        elif hook and hook.cleanup and is_cleanup:
+                            hook._returned_cleanup = res if callable(res) else None
+                        elif hook and is_cleanup:
                             hook.cancel()
-                            if inspect.iscoroutinefunction(hook.cleanup):
-                                hook._cleanup_task = asyncio.create_task(hook.cleanup())
-                            else:
-                                hook.cleanup()
+                            cleanup = hook.cleanup or hook._returned_cleanup
+                            hook._returned_cleanup = None
+                            if callable(cleanup) and inspect.iscoroutinefunction(
+                                cleanup
+                            ):
+                                hook._cleanup_task = asyncio.create_task(cleanup())
+                            elif callable(cleanup):
+                                cleanup()
                     except Exception as ex:
                         tb = traceback.format_exc()
                         self.error(f"Exception in effect: {ex}\n{tb}")

sdk/python/packages/flet/tests/test_component_effects.py

@@ -1,3 +1,8 @@
+import asyncio
+
+import pytest
+
+import flet.components.hooks.use_effect as use_effect_module
 from flet.components.component import Component
 from flet.components.hooks.use_effect import EffectHook
 
@@ -12,7 +17,7 @@ def test_render_effect_changed_deps_runs_cleanup_then_setup_once_each():
         EffectHook(
             c,
             setup=lambda: None,
-            cleanup=lambda: None,
+            _returned_cleanup=lambda: None,
             deps=[2],
             prev_deps=[1],
         )
@@ -60,3 +65,56 @@ def test_will_unmount_releases_cached_render_references():
     assert c._state.last_args == ()
     assert c._state.last_kwargs == {}
     assert c._contexts == {}
+
+
+def test_use_effect_retains_setup_returned_cleanup_across_renders(monkeypatch):
+    c = Component(fn=lambda: None, args=(), kwargs={})
+
+    def cleanup():
+        return None
+
+    def setup():
+        return cleanup
+
+    monkeypatch.setattr(use_effect_module, "current_component", lambda: c)
+
+    use_effect_module.use_effect(setup, dependencies=[])
+    hook = c._state.hooks[0]
+    hook._returned_cleanup = cleanup
+
+    c._state.hook_cursor = 0
+    use_effect_module.use_effect(setup, dependencies=[])
+
+    assert hook._returned_cleanup is cleanup
+
+
+def test_unmount_schedules_returned_cleanup():
+    c = Component(fn=lambda: None, args=(), kwargs={})
+    c._state.mounted = True
+    calls: list[bool] = []
+    c._schedule_effect = lambda hook, is_cleanup=False: calls.append(is_cleanup)  # type: ignore[method-assign]
+    hook = EffectHook(c, setup=lambda: None, _returned_cleanup=lambda: None)
+    c._state.hooks = [hook]
+
+    c._run_unmount_effects()
+
+    assert calls == [True]
+
+
+@pytest.mark.asyncio
+async def test_unmount_schedules_cleanup_for_pending_setup_task():
+    c = Component(fn=lambda: None, args=(), kwargs={})
+    c._state.mounted = True
+    calls: list[bool] = []
+    c._schedule_effect = lambda hook, is_cleanup=False: calls.append(is_cleanup)  # type: ignore[method-assign]
+    task = asyncio.create_task(asyncio.sleep(60))
+    hook = EffectHook(c, setup=lambda: None)
+    hook._setup_task = task
+    c._state.hooks = [hook]
+
+    c._run_unmount_effects()
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    assert calls == [True]

sdk/python/packages/flet/tests/test_session_disconnect_buffering.py

@@ -1,5 +1,8 @@
+import asyncio
 from datetime import datetime, timezone
 
+import pytest
+
 from flet.components.component import Component
 from flet.components.hooks.use_effect import EffectHook
 from flet.messaging.connection import Connection
@@ -70,3 +73,56 @@ def test_attach_connection_restores_state_and_flushes_buffer():
     assert session.expires_at is None
     assert session._Session__send_buffer == []
     assert new_conn.messages == [buffered_message]
+
+
+@pytest.mark.asyncio
+async def test_cleanup_returned_by_effect_setup_runs_on_unmount():
+    conn = Connection()
+    conn.pubsubhub = PubSubHub()
+    session = Session(conn)
+    component = Component(fn=lambda: None, args=(), kwargs={})
+    component._state.mounted = True
+    cleanup_ran = asyncio.Event()
+
+    def cleanup():
+        cleanup_ran.set()
+
+    hook = EffectHook(component, setup=lambda: None, _returned_cleanup=cleanup)
+    component._state.hooks = [hook]
+    component._schedule_effect = session.schedule_effect  # type: ignore[method-assign]
+    component._run_unmount_effects()
+    session.start_updates_scheduler()
+
+    await asyncio.wait_for(cleanup_ran.wait(), timeout=1)
+    session._Session__closed = True
+    session._Session__updates_ready.set()
+    await session._Session__updates_task
+
+
+@pytest.mark.asyncio
+async def test_unmount_cancels_pending_async_effect_setup():
+    conn = Connection()
+    conn.pubsubhub = PubSubHub()
+    session = Session(conn)
+    component = Component(fn=lambda: None, args=(), kwargs={})
+    component._state.mounted = True
+    started = asyncio.Event()
+
+    async def setup():
+        started.set()
+        await asyncio.Event().wait()
+
+    task = asyncio.create_task(setup())
+    await started.wait()
+    hook = EffectHook(component, setup=setup)
+    hook._setup_task = task
+    component._state.hooks = [hook]
+    component._schedule_effect = session.schedule_effect  # type: ignore[method-assign]
+    component._run_unmount_effects()
+    session.start_updates_scheduler()
+
+    with pytest.raises(asyncio.CancelledError):
+        await asyncio.wait_for(task, timeout=1)
+    session._Session__closed = True
+    session._Session__updates_ready.set()
+    await session._Session__updates_task