Preserve async shell context across cells

Author: fff122

ipykernel/ipkernel.py

@@ -438,7 +438,8 @@ async def run_cell(*args, **kwargs):
                     **do_execute_args,
                 )
 
-                coro_future = asyncio.ensure_future(coro)
+                coro_future = asyncio.current_task()
+                assert coro_future is not None
 
                 cm = (
                     self._cancel_on_sigint
@@ -448,7 +449,7 @@ async def run_cell(*args, **kwargs):
                 with cm(coro_future):
                     res = None
                     try:
-                        res = await coro_future
+                        res = await coro
                     finally:
                         shell.events.trigger("post_execute")
                         if not silent:

ipykernel/utils.py

@@ -52,27 +52,27 @@ def _async_in_context(
     if context is None:
         context = copy_context()
 
-    if sys.version_info >= (3, 11):
-
-        @wraps(f)
-        async def run_in_context(*args, **kwargs):
-            coro = f(*args, **kwargs)
-            return await asyncio.create_task(coro, context=context)
-
-        return run_in_context
-
-    # don't need this backport when we require 3.11
-    # context_holder so we have a modifiable container for later calls
-    context_holder = [context]  # type: ignore[unreachable]
+    # Keep the context created by one shell request for subsequent requests. This is
+    # needed because an async cell runs in a task context, and changes made after an
+    # await would otherwise be discarded when that task finishes.
+    context_holder = [context]
 
     async def preserve_context(f, *args, **kwargs):
-        """call a coroutine, preserving the context after it is called"""
+        """Call a coroutine in a persistent ContextVar context."""
         try:
             return await f(*args, **kwargs)
         finally:
-            # persist changes to the context for future calls
             context_holder[0] = copy_context()
 
+    if sys.version_info >= (3, 11):
+
+        @wraps(f)
+        async def run_in_context(*args, **kwargs):
+            coro = preserve_context(f, *args, **kwargs)
+            return await asyncio.create_task(coro, context=context_holder[0])
+
+        return run_in_context
+
     @wraps(f)
     async def run_in_context_pre311(*args, **kwargs):
         ctx = context_holder[0]

tests/test_kernel.py

@@ -908,9 +908,15 @@ def test_context_vars():
         stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
         assert stdout.strip() == "set"
 
+        msg_id, _ = execute(
+            kc=kc,
+            code="async def produce_result():\n    return 'set after await'\nresult = await produce_result(); ctxvar.set(result)",
+        )
+        stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
+
         msg_id, _ = execute(
             kc=kc,
             code="print(ctxvar.get())",
         )
         stdout, _ = assemble_output(kc.get_iopub_msg, parent_msg_id=msg_id)
-        assert stdout.strip() == "set"
+        assert stdout.strip() == "set after await"