[core] Isolate contextvars per task in the coroutine runner (#19238)

esphome/coroutine.py

@@ -45,6 +45,7 @@ the last `yield` expression defines what is returned.
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable, Generator, Iterator
+import contextvars
 import enum
 import functools
 import heapq
@@ -277,14 +278,22 @@ class _Task:
         id_number: int,
         iterator: Iterator[None],
         original_function: Any,
+        context: contextvars.Context,
     ):
         self.priority = priority
         self.id_number = id_number
         self.iterator = iterator
         self.original_function = original_function
+        self.context = context
 
     def with_priority(self, priority: float) -> _Task:
-        return _Task(priority, self.id_number, self.iterator, self.original_function)
+        return _Task(
+            priority,
+            self.id_number,
+            self.iterator,
+            self.original_function,
+            self.context,
+        )
 
     @property
     def _cmp_tuple(self) -> tuple[float, int]:
@@ -321,7 +330,10 @@ class FakeEventLoop:
             coro = coroutine(func)
             gen = coro(*args, **kwargs)
         prio = getattr(coro, "priority", 0.0)
-        task = _Task(prio, self._task_counter, gen, func)
+        # Each task gets its own copy of the current context, isolating any
+        # contextvars it sets from other tasks the scheduler interleaves it with
+        # (mirrors what asyncio.Task does internally).
+        task = _Task(prio, self._task_counter, gen, func, contextvars.copy_context())
         self._task_counter += 1
         heapq.heappush(self._pending_tasks, task)
 
@@ -352,7 +364,7 @@ class FakeEventLoop:
             )
 
             try:
-                next(task.iterator)
+                task.context.run(next, task.iterator)
                 # Decrease priority over time, so that if this task is blocked
                 # due to a dependency others will clear the dependency
                 # This could be improved with a less naive approach

tests/unit_tests/test_coroutine.py

@@ -1,5 +1,7 @@
 """Tests for the coroutine module."""
 
+import contextvars
+
 import pytest
 
 from esphome.coroutine import CoroPriority, FakeEventLoop, coroutine_with_priority
@@ -217,3 +219,46 @@ def test_custom_priority_between_enum_values() -> None:
 
     # Check execution order
     assert execution_order == ["core", "custom", "diagnostics"]
+
+
+def test_context_isolated_between_interleaved_tasks() -> None:
+    """Test that a contextvar set in one task does not leak into another task that the scheduler interleaves with it."""
+    my_var: contextvars.ContextVar[str] = contextvars.ContextVar("my_var")
+    seen: dict[str, str] = {}
+
+    def task_a():
+        my_var.set("a")
+        yield  # suspend so task_b can run before task_a resumes
+        seen["a"] = my_var.get()
+
+    def task_b():
+        my_var.set("b")
+        yield
+        seen["b"] = my_var.get()
+
+    loop = FakeEventLoop()
+    loop.add_job(task_a)
+    loop.add_job(task_b)
+    loop.flush_tasks()
+
+    assert seen == {"a": "a", "b": "b"}
+
+
+def test_context_inherits_ambient_value_at_schedule_time() -> None:
+    """Test that a job sees whatever contextvar value was set before it was scheduled."""
+    my_var: contextvars.ContextVar[str] = contextvars.ContextVar("my_var")
+    token = my_var.set("ambient")
+    seen: dict[str, str] = {}
+
+    def task():
+        seen["value"] = my_var.get()
+        yield
+
+    try:
+        loop = FakeEventLoop()
+        loop.add_job(task)
+        loop.flush_tasks()
+    finally:
+        my_var.reset(token)
+
+    assert seen == {"value": "ambient"}