Carry the runtime Context into thread pool workers

src/nu/lang/runtime/runtime.py

@@ -57,6 +57,18 @@
 _RT_CTX: contextvars.ContextVar[Context] = contextvars.ContextVar("nu_rt_ctx")
 
 
+def _carry_ctx() -> Callable[..., object]:
+    """Return a runner that calls a function inside a copy of the caller's context.
+
+    A worker thread starts with an empty contextvars context, so `_RT_CTX` is
+    unset there and `Runtime.ctx` raises `LookupError`. Taking the copy on the
+    calling side and submitting `copy.run` keeps the Context resolvable on the
+    worker. A fresh copy per branch keeps a `.set()` inside that branch local to
+    it, the same copy-on-write rule an asyncio.Task gets.
+    """
+    return contextvars.copy_context().run
+
+
 class Runtime:
     """Per-drive Runtime. Owns a Program, a per-task Context, and a Budget."""
 
@@ -118,7 +130,7 @@ def eval_parallel(self, nids: Iterable[int]) -> list:
         if self.budget.max_parallel == 1 or self.budget.thread_pool is None:
             return self.eval_each(nids)
         pool = self.budget.thread_pool
-        futures = [pool.submit(self.eval, n) for n in nids]
+        futures = [pool.submit(_carry_ctx(), self.eval, n) for n in nids]
         return [f.result() for f in futures]
 
     def _drive_async(self, nids: list[int]) -> list:
@@ -147,7 +159,7 @@ async def place(n: int) -> object:
             async with sem:
                 if on_loop_col[n]:
                     return await self.aeval(n)
-                return await loop.run_in_executor(pool, self.eval, n)
+                return await loop.run_in_executor(pool, _carry_ctx(), self.eval, n)
 
         return [place(n) for n in nids]
 
@@ -261,7 +273,7 @@ def drain(n: int) -> None:
             finally:
                 q.put(_DONE)
 
-        futures = [pool.submit(drain, n) for n in nids]
+        futures = [pool.submit(_carry_ctx(), drain, n) for n in nids]
         remaining = len(futures)
         try:
             while remaining > 0:
@@ -306,20 +318,21 @@ def in_thread(self, fn: Callable, *args: object, **kwargs: object) -> Future:
         if self.budget.thread_pool is None:
             msg = "in_thread requires max_parallel > 1"
             raise RuntimeError(msg)
-        return self.budget.thread_pool.submit(fn, *args, **kwargs)
+        return self.budget.thread_pool.submit(_carry_ctx(), fn, *args, **kwargs)
 
     async def a_in_thread(self, fn: Callable, *args: object, **kwargs: object) -> object:
         """Await a blocking call on the Budget's thread pool."""
         if self.budget.thread_pool is None:
             msg = "a_in_thread requires max_parallel > 1"
             raise RuntimeError(msg)
         loop = asyncio.get_running_loop()
+        run = _carry_ctx()
         if kwargs:
             return await loop.run_in_executor(
                 self.budget.thread_pool,
-                lambda: fn(*args, **kwargs),
+                lambda: run(fn, *args, **kwargs),
             )
-        return await loop.run_in_executor(self.budget.thread_pool, fn, *args)
+        return await loop.run_in_executor(self.budget.thread_pool, run, fn, *args)
 
     # --- sentinel-propagating evaluation -----------------------------------
 
@@ -417,7 +430,7 @@ async def run_child(n: int) -> None:
                             await q.put(v)
                 else:
                     async with sem:
-                        await loop.run_in_executor(pool, _drain_sync, n, loop)
+                        await loop.run_in_executor(pool, _carry_ctx(), _drain_sync, n, loop)
             finally:
                 await q.put(_DONE)
 

tests/nu/lang/runtime/test_runtime.py

@@ -512,6 +512,67 @@ def test_ctx_attrs_writes_are_visible_through_runtime() -> None:
     assert ctx.attrs["y"] == "hello"
 
 
+# --- context across the thread boundary -----------------------------------
+#
+# A pool worker starts with an empty contextvars context, so every hand-off to
+# a thread has to carry the caller's copy or `rt.ctx` raises LookupError there.
+# One case per dispatch site.
+
+
+def test_ctx_resolves_on_eval_parallel_workers() -> None:
+    program = _fake_program(thunks=[lambda rt: rt.ctx])
+    ctx = Context()
+    with Budget(max_parallel=2) as budget:
+        rt = Runtime(program, ctx, budget=budget)
+        assert rt.eval_parallel([0, 0]) == [ctx, ctx]
+
+
+async def test_ctx_resolves_on_async_placement_workers() -> None:
+    program = _fake_program(thunks=[lambda rt: rt.ctx], athunks=[None], on_loop=[False])
+    ctx = Context()
+    with Budget(max_parallel=2, async_mode=True) as budget:
+        rt = Runtime(program, ctx, budget=budget)
+        assert await rt.aeval_parallel([0, 0]) == [ctx, ctx]
+
+
+def test_ctx_resolves_on_merge_workers() -> None:
+    def gen(rt: Runtime) -> object:
+        return iter([rt.ctx])
+
+    program = _fake_program(thunks=[gen, gen])
+    ctx = Context()
+    with Budget(max_parallel=2) as budget:
+        rt = Runtime(program, ctx, budget=budget)
+        assert list(rt.merge([0, 1])) == [ctx, ctx]
+
+
+async def test_ctx_resolves_on_amerge_workers() -> None:
+    def gen(rt: Runtime) -> object:
+        return iter([rt.ctx])
+
+    program = _fake_program(thunks=[gen], athunks=[None], on_loop=[False])
+    ctx = Context()
+    with Budget(max_parallel=2, async_mode=True) as budget:
+        rt = Runtime(program, ctx, budget=budget)
+        assert [v async for v in rt.amerge([0])] == [ctx]
+
+
+def test_ctx_resolves_inside_in_thread() -> None:
+    program = compile(Literal(1))
+    ctx = Context()
+    with Budget(max_parallel=2) as budget:
+        rt = Runtime(program, ctx, budget=budget)
+        assert rt.in_thread(lambda: rt.ctx).result() is ctx
+
+
+async def test_ctx_resolves_inside_a_in_thread() -> None:
+    program = compile(Literal(1))
+    ctx = Context()
+    with Budget(max_parallel=2, async_mode=True) as budget:
+        rt = Runtime(program, ctx, budget=budget)
+        assert await rt.a_in_thread(lambda: rt.ctx) is ctx
+
+
 # --- budget lifecycle -----------------------------------------------------