Remove all observers on YChat.unobserve()

[dlqqq]

Summary

YChat.__init__() registers an observer (_on_messages_change, plus the state
observer _initialize) whose callback is a bound method holding a reference to
self. unobserve() did not remove these, so the YChat object was never
garbage collected after unobserve() — leaking the document and everything it
references.

This stores the messages subscription and adds a YChat.unobserve() override
that removes both __init__ observers (in addition to the observe()
subscriptions cleared by the base class). Adds a GC regression test.

Changes

• Store the messages subscription in __init__ as _ymessages_subscription.
• Add a YChat.unobserve() override: call super().unobserve(), then remove
  _ymessages_subscription and _ystate_subscription (guarded against
  double-unobserve by nulling out after removal).
• _initialize() now nulls _ystate_subscription after unobserving, so
  unobserve() doesn't try to remove it twice.
• observe() now calls super().unobserve() (instead of self.unobserve()) so
  re-registering observe() subscriptions no longer strips the __init__
  observers that the document needs for its lifetime.
• Add tests/test_ychat_gc.py with weakref-based GC regression tests covering:
  plain unobserve(), observe() + unobserve(), and teardown while still
  dirty.

[github-actions]

👈 Launch a Binder on branch dlqqq/jupyter-chat/20260707-chat-fix-ychat-unobserve%2Fjupyter-chat

[dlqqq]

@brichet This PR fixes a memory leak by making YChat.unobserve() match the expected behavior of YBaseDoc.unobserve(), which should remove all observers from the shared model. Without this, consumers have no way to remove that '_on_messages_change' observer added in YChat.__init__() , meaning that YChat objects get leaked in memory since that observer holds a reference to self: YChat.

Typically YBaseDoc.unobserve() should only get called when a collaborative room is shutting down (e.g. after the client closes the chat). There's no direct API for just removing one observer: this method should remove all observers, and it's implied that this only gets called when the parent YRoom is shutting down.

So even though this is technically a breaking change (calling unobserve() now removes the _on_messages_change observer permanently), it's doing so to match the expected behavior of shared models provided by Jupyter YDoc, and it fixes a memory leak.

Ideally, there would be some kind of YBaseDoc.dispose() method as well to distinguish 'removing an observer' v.s. 'prepare the shared model to be freed from memory', but that's another API-breaking change that'll have to land in Jupyter YDoc v5.

[brichet]

Thanks @dlqqq for working on this.

Perhaps this is expected at that stage, but the new tests fail with jupyter-collaboration >= 5.

jupyter-collaboration==4.4.1

$ pip freeze | grep jupyter-collaboration
jupyter-collaboration==4.4.1
jupyter-collaboration-ui==2.4.1

$ pytest
============================= test session starts ==============================
platform linux -- Python 3.14.6, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/brichet/projects/jupyter-chat
configfile: pyproject.toml
plugins: cov-7.1.0, anyio-4.14.1, asyncio-1.4.0
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 22 items                                                             

python/jupyterlab-chat/jupyterlab_chat/tests/test_imports.py .           [  4%]
python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py ............. [ 63%]
.....                                                                    [ 86%]
python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py ...        [100%]

============================== 22 passed in 1.01s ==============================

jupyter-collaboration==5.0.0a0

$ pip freeze | grep jupyter-collaboration
jupyter-collaboration==5.0.0a0
jupyter-collaboration-ui==3.0.0a0

$ pytest
============================= test session starts ==============================
platform linux -- Python 3.14.6, pytest-9.1.1, pluggy-1.6.0
rootdir: /home/brichet/projects/jupyter-chat
configfile: pyproject.toml
plugins: cov-7.1.0, anyio-4.14.1, asyncio-1.4.0
asyncio: mode=Mode.STRICT, debug=False, asyncio_default_fixture_loop_scope=None, asyncio_default_test_loop_scope=function
collected 22 items                                                             

python/jupyterlab-chat/jupyterlab_chat/tests/test_imports.py .           [  4%]
python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat.py ............. [ 63%]
.....                                                                    [ 86%]
python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py FFF        [100%]

=================================== FAILURES ===================================
_____________________ test_ychat_collected_after_unobserve _____________________

    def test_ychat_collected_after_unobserve():
        """YChat must be garbage-collected after unobserve().
    
        Observers registered in __init__ are bound methods that capture ``self``.
        If unobserve() does not remove them, the YChat is never collected (leak).
        """
        ychat = YChat()
        ychat.unobserve()
    
        ref = weakref.ref(ychat)
        del ychat
        gc.collect()
    
>       assert ref() is None, "YChat leaked after unobserve() (observers not removed)"
E       AssertionError: YChat leaked after unobserve() (observers not removed)
E       assert <jupyterlab_chat.ychat.YChat object at 0x7160197f3e90> is None
E        +  where <jupyterlab_chat.ychat.YChat object at 0x7160197f3e90> = <weakref at 0x7160193d4cc0; to 'jupyterlab_chat.ychat.YChat' at 0x7160197f3e90>()

python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py:26: AssertionError
______________ test_ychat_collected_after_observe_then_unobserve _______________

    def test_ychat_collected_after_observe_then_unobserve():
        """YChat must be collectable after the realistic observe()/unobserve() flow.
    
        A server registers a sync callback via ``observe()`` and later tears the
        document down with ``unobserve()``. Both the ``observe()`` subscriptions and
        the ``__init__`` observers must be removed for the YChat to be collected.
        """
        ychat = YChat()
        ychat.observe(lambda *args: None)
        ychat.unobserve()
    
        ref = weakref.ref(ychat)
        del ychat
        gc.collect()
    
>       assert ref() is None, "YChat leaked after observe() then unobserve()"
E       AssertionError: YChat leaked after observe() then unobserve()
E       assert <jupyterlab_chat.ychat.YChat object at 0x716019e05cd0> is None
E        +  where <jupyterlab_chat.ychat.YChat object at 0x716019e05cd0> = <weakref at 0x71601985c310; to 'jupyterlab_chat.ychat.YChat' at 0x716019e05cd0>()

python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py:44: AssertionError
_______________ test_ychat_collected_after_unobserve_while_dirty _______________

    def test_ychat_collected_after_unobserve_while_dirty():
        """YChat must be collectable even if torn down before any state event fires.
    
        When the document is still dirty, ``_initialize`` has not yet removed the
        state subscription, so ``unobserve()`` must remove it defensively.
        """
        ychat = YChat()
        assert ychat.dirty is True
        ychat.unobserve()
    
        ref = weakref.ref(ychat)
        del ychat
        gc.collect()
    
>       assert ref() is None, "YChat leaked after unobserve() while dirty"
E       AssertionError: YChat leaked after unobserve() while dirty
E       assert <jupyterlab_chat.ychat.YChat object at 0x7160193f81d0> is None
E        +  where <jupyterlab_chat.ychat.YChat object at 0x7160193f81d0> = <weakref at 0x71601982fb00; to 'jupyterlab_chat.ychat.YChat' at 0x7160193f81d0>()

python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py:61: AssertionError
=========================== short test summary info ============================
FAILED python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py::test_ychat_collected_after_unobserve - AssertionError: YChat leaked after unobserve() (observers not removed)
FAILED python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py::test_ychat_collected_after_observe_then_unobserve - AssertionError: YChat leaked after observe() then unobserve()
FAILED python/jupyterlab-chat/jupyterlab_chat/tests/test_ychat_gc.py::test_ychat_collected_after_unobserve_while_dirty - AssertionError: YChat leaked after unobserve() while dirty
========================= 3 failed, 19 passed in 1.22s =========================

[brichet]

Is this still necessary after #457 ?

[brichet]

Out of curiosity (not related to the change), because I did not realize that before this PR, does it mean that only one callback can be attached to the whole document change (probably the document room) ?
What would happen if another extension do ychat.observe(callback) ?