Fix logic when looking up pn.state.curdoc (#6254)

panel/io/state.py

@@ -968,16 +968,16 @@ def curdoc(self) -> Document | None:
         """
         Returns the Document that is currently being executed.
         """
+        curdoc = self._curdoc.get()
+        if curdoc:
+            return curdoc
         try:
             doc = curdoc_locked()
             pyodide_session = self._is_pyodide and 'pyodide_kernel' not in sys.modules
             if doc and (doc.session_context or pyodide_session):
                 return doc
         except Exception:
-            pass
-        curdoc = self._curdoc.get()
-        if curdoc:
-            return curdoc
+            return None
 
     @curdoc.setter
     def curdoc(self, doc: Document) -> None:

panel/tests/conftest.py

@@ -21,6 +21,7 @@
 
 from bokeh.client import pull_session
 from bokeh.document import Document
+from bokeh.io.doc import curdoc, set_curdoc as set_bkdoc
 from pyviz_comms import Comm
 
 from panel import config, serve
@@ -159,6 +160,18 @@ def server_document():
         yield doc
     doc._session_context = None
 
+@pytest.fixture
+def bokeh_curdoc():
+    old_doc = curdoc()
+    doc = Document()
+    session_context = unittest.mock.Mock()
+    doc._session_context = lambda: session_context
+    set_bkdoc(doc)
+    try:
+        yield doc
+    finally:
+        set_bkdoc(old_doc)
+
 @pytest.fixture
 def comm():
     return Comm()

panel/tests/test_server.py

@@ -155,16 +155,15 @@ async def cb(event, count=[0]):
     wait_until(lambda: len(counts) > 0 and max(counts) > 1)
 
 
-def test_server_async_local_state():
+def test_server_async_local_state(bokeh_curdoc):
     docs = {}
 
     async def task():
-        curdoc = state.curdoc
         await asyncio.sleep(0.5)
-        docs[curdoc] = []
+        docs[state.curdoc] = []
         for _ in range(5):
             await asyncio.sleep(0.1)
-            docs[curdoc].append(state.curdoc)
+            docs[state.curdoc].append(state.curdoc)
 
     def app():
         state.execute(task)
@@ -177,18 +176,17 @@ def app():
     wait_until(lambda: all([len(set(docs)) == 1 and docs[0] is doc for doc, docs in docs.items()]))
 
 
-def test_server_async_local_state_nested_tasks():
+def test_server_async_local_state_nested_tasks(bokeh_curdoc):
     docs = {}
 
     async def task(depth=1):
-        curdoc = state.curdoc
         await asyncio.sleep(0.5)
         if depth > 0:
             asyncio.ensure_future(task(depth-1))
-        docs[curdoc] = []
+        docs[state.curdoc] = []
         for _ in range(10):
             await asyncio.sleep(0.1)
-            docs[curdoc].append(state.curdoc)
+            docs[state.curdoc].append(state.curdoc)
 
     def app():
         state.execute(task)

panel/tests/ui/io/test_state.py

@@ -1,3 +1,5 @@
+import asyncio
+
 import pytest
 
 pytest.importorskip("playwright")
@@ -6,7 +8,8 @@
 
 from panel.io.state import state
 from panel.pane import Markdown
-from panel.tests.util import serve_component
+from panel.tests.util import serve_component, wait_until
+from panel.widgets import Button
 
 pytestmark = pytest.mark.ui
 
@@ -23,3 +26,37 @@ def cb():
     serve_component(page, app)
 
     expect(page.locator('.markdown').locator("div")).to_have_text('Loaded\n')
+
+
+def test_server_async_local_state_button_click(page, bokeh_curdoc):
+    docs = {}
+    buttons = {}
+
+    async def task(event):
+        assert buttons[event.obj] is state.curdoc
+        await asyncio.sleep(0.5)
+        docs[state.curdoc] = []
+        for _ in range(10):
+            await asyncio.sleep(0.1)
+            docs[state.curdoc].append(state.curdoc)
+
+    def app():
+        button = Button(on_click=task)
+        buttons[button] = state.curdoc
+        return button
+
+    _, port = serve_component(page, app)
+
+    page.click('.bk-btn')
+
+    page.goto(f"http://localhost:{port}")
+
+    page.click('.bk-btn')
+
+    page.goto(f"http://localhost:{port}")
+
+    page.click('.bk-btn')
+
+    # Ensure state.curdoc was consistent despite asyncio context switching
+    wait_until(lambda: len(docs) == 3)
+    wait_until(lambda: all([len(set(docs)) == 1 and docs[0] is doc for doc, docs in docs.items()]))