fix: Only treat the document lock as held by the context that acquired it

@@ -38,6 +38,9 @@ from bokeh.embed.wrappers import wrap_in_script_tag from bokeh.server.contexts import _RequestProxy as BkRequestProxy from bokeh.server.server import Server as BokehServer+from bokeh.server.session import (+    ServerSession, _needs_document_lock, _needs_document_lock_on_loop,+) from bokeh.server.urls import per_app_patterns, toplevel_patterns from bokeh.server.views.autoload_js_handler import (     AutoloadJsHandler as BkAutoloadJsHandler,@@ -104,7 +107,6 @@     from bokeh.core.types import ID     from bokeh.document.document import DocJson     from bokeh.embed.bundle import Bundle-    from bokeh.server.session import ServerSession     from jinja2 import Template      from .application import TViewableFuncOrPath@@ -734,6 +736,30 @@ def validate_static_dirs(static_dirs: Mapping[str, str]) -> dict[str, str]:     return static  +def _track_document_lock(func):+    if inspect.iscoroutinefunction(func):+        @wraps(func)+        async def wrapper(self, *args, **kwargs):+            with state._own_document_lock(self):+                return await func(self, *args, **kwargs)+    else:+        @wraps(func)+        def wrapper(self, *args, **kwargs):+            with state._own_document_lock(self):+                return func(self, *args, **kwargs)+    return wrapper++# Record which context holds the document lock so state._unblocked can+# tell the holder apart from other threads while the lock is taken.+ServerSession.with_document_locked = _needs_document_lock(  # type: ignore[method-assign]+    _track_document_lock(ServerSession.with_document_locked.__wrapped__)  # type: ignore[attr-defined]+)+for _method in ('_handle_push', '_handle_patch'):+    setattr(ServerSession, _method, _needs_document_lock_on_loop(+        _track_document_lock(getattr(ServerSession, _method).__wrapped__)+    ))++ # Patch Server to attach task factory to asyncio loop and handle Admin server context class Server(BokehServer): 

@@ -157,6 +157,12 @@ class _state(param.Parameterized):     # Holds temporary curdoc overrides per thread     _curdoc: ContextVar[Document | None] = ContextVar('curdoc', default=None) +    # A token identifying the document lock acquisition the current+    # context runs under, matched against the token of the acquisition+    # each session is currently held by+    _document_lock: ContextVar[object | None] = ContextVar('document_lock', default=None)+    _document_lock_owners: t.ClassVar[WeakKeyDictionary[ServerSession, object]] = WeakKeyDictionary()+     # Whether to hold comm events     _hold: bool = False @@ -341,10 +347,30 @@ def _document_lock_held(self) -> bool:         is held for the duration. Mutating Bokeh models is safe in that         case since no other code can be touching them concurrently, but         raw socket writes are NOT (see `_on_loop_thread`).++        Only the context the lock was acquired for counts. Another+        callback holding the lock does not make writes from this+        thread safe, so ownership is tracked by `_own_document_lock`.         """         session = getattr(getattr(self.curdoc, 'session_context', None), 'session', None)-        session_lock = getattr(session, '_lock', None)-        return bool(session_lock is not None and session_lock.locked())+        owner = self._document_lock.get()+        return owner is not None and session is not None and self._document_lock_owners.get(session) is owner++    @contextmanager+    def _own_document_lock(self, session: ServerSession) -> Iterator[None]:+        """+        Marks the current context as running under the session's+        document lock. Must only be entered while the lock is held.+        """+        owner = object()+        self._document_lock_owners[session] = owner+        token = self._document_lock.set(owner)+        try:+            yield+        finally:+            self._document_lock.reset(token)+            if self._document_lock_owners.get(session) is owner:+                del self._document_lock_owners[session]      def _install_thread_pool(self, loop: asyncio.AbstractEventLoop | None = None) -> None:         """

@@ -380,6 +380,57 @@ async def trigger():     assert 'on_loop' in seen  +@pytest.mark.xdist_group(name="server")+def test_unblocked_only_for_document_lock_holder():+    """+    While a locked callback holds the document lock on a worker thread,+    only that thread may write to the Document directly. Any other+    thread has to schedule its changes, otherwise Bokeh raises because+    the writes race with the holder releasing the lock.+    """+    slider = IntSlider()++    serve_and_request(slider)+    wait_until(lambda: bool(slider._documents))++    doc = list(slider._documents)[0]+    session = doc.session_context.session+    asyncio_loop = doc.session_context.server_context.application_context.io_loop.asyncio_loop+    state._connected[doc] = True++    entered, release = threading.Event(), threading.Event()+    seen = {}++    def holder():+        with set_curdoc(doc):+            seen['holder'] = state._unblocked(doc)+        entered.set()+        release.wait(timeout=5)++    def other():+        with set_curdoc(doc):+            seen['other'] = state._unblocked(doc)++    async def trigger():+        await session.with_document_locked(holder)++    future = asyncio.run_coroutine_threadsafe(trigger(), asyncio_loop)+    try:+        assert entered.wait(timeout=5)+        thread = threading.Thread(target=other)+        thread.start()+        thread.join()+    finally:+        release.set()+    future.result(timeout=5)++    assert seen == {'holder': True, 'other': False}++    # Once released the holder's context no longer counts as owning the lock+    with set_curdoc(doc):+        assert not state._document_lock_held++ @pytest.mark.xdist_group(name="server") def test_threaded_hold_does_not_boomerang_python_changes():     """

@@ -114,7 +114,7 @@ def test_fast_list_template_both_sizing_modes_fill_height(page, sizing_mode):     card = page.locator(f"#main .card-margin.{sizing_mode}")     expect(card).to_have_count(1)     wait_until(lambda: card.bounding_box()['height'] > 300, page)-    +  ACTIVE_TOOL_COLORS = """ () => {