"""Empirical test: contextvars + async generator suspend/resume.
Demonstrates that a token created in an async generator before yield can be
invalid in that generator's finally after the consumer raises; reset(token)
fails with ValueError ("Token was created in a different Context"). This
justifies not relying on contextvars inside async generators that are
consumed across suspend/resume (e.g. chat_completion_create).
Consumer (evaluator) scenario: when the consumer sets context and does
async for over a generator that raises, our tests show the consumer's token
remains valid in the consumer's finally. So we do not reproduce a consumer-side
failure; avoiding start_as_current_span in evaluate is defensive (same
mechanism class; generator-side fix is safer).
Run from repo root:
uv run python internal_docs/vignettes/otel-contextvars-async/contextvars_async_gen_demo.py
No pytest: uses asyncio + assert only, to avoid test-framework confounders.
"""
import asyncio
from contextlib import aclosing
from contextvars import ContextVar
cv = ContextVar[float]("x", default=42)
async def f():
# Set value and get token in Context A
token = cv.set(3.14)
print(f"Generator set cv to: {cv.get()}")
try:
yield 1
finally:
# We're now in Context B (due to exception from main)
print(f"Finally block inside generator: {cv.get()}")
try:
cv.reset(token) # This fails
except ValueError as exc:
print(f"✗ Token reset failed: {exc}")
async def main():
print(f"Initial default: {cv.get()}")
try:
async for _ in f():
raise RuntimeError()
except RuntimeError:
print("Handling RuntimeError")
# Value persists because token reset failed
print(f"After generator: {cv.get()}")
async def f_normal():
"""Generator that yields twice and exits normally."""
token = cv.set(100.0)
print(f" [gen] set cv to: {cv.get()}")
try:
yield 1
print(f" [gen] after first yield, cv.get() = {cv.get()}")
yield 2
print(f" [gen] after second yield, cv.get() = {cv.get()}")
finally:
print(f" [gen] finally: cv.get() = {cv.get()}")
try:
cv.reset(token)
print(" [gen] reset(token) OK")
except ValueError as exc:
print(f" [gen] ✗ reset failed: {exc}")
async def main_normal():
"""Normal async for (no exception). Does context survive suspend/resume?"""
print("--- Normal async for (no exception) ---")
print(f"Before loop: {cv.get()}")
async for x in f_normal():
print(f" [main] got {x}, cv.get() = {cv.get()}")
print(f"After loop: {cv.get()}")
async def gen_that_yields_then_raises():
"""Simulates chat_completion_create: yields once, then raises (e.g. API error)."""
yield 1
raise RuntimeError("simulated client error")
async def consumer_with_context():
"""
Simulates the EVALUATOR: consumer sets context (like start_as_current_span),
then does async for over a generator that raises. When we exit (due to exception),
does our token reset in finally still work?
"""
token = cv.set("evaluator_span")
try:
async for chunk in gen_that_yields_then_raises():
pass
finally:
# This is like the __exit__ of start_as_current_span (detach current span).
try:
cv.reset(token)
print(" [consumer] reset(token) OK")
except ValueError as exc:
print(f" [consumer] ✗ reset(token) failed: {exc}")
async def main_consumer_raises():
"""
Prove: consumer (evaluator) sets context, then async for over generator that raises.
Consumer's finally runs when exception propagates. Is token still valid?
"""
print("--- Consumer sets context, async for over generator that raises ---")
print(f"Before: {cv.get()}")
try:
await consumer_with_context()
except RuntimeError as e:
print(f" Caught: {e}")
print(f"After: {cv.get()}")
async def gen_yield_sleep_raise():
"""Yields, then sleep(0), then raises — more context switches before raise."""
yield 1
await asyncio.sleep(0)
yield 2
raise RuntimeError("client error after second chunk")
async def consumer_with_sleep():
"""Consumer sets context, async for over gen that yields twice then raises."""
token = cv.set("evaluator_span")
try:
async for chunk in gen_yield_sleep_raise():
await asyncio.sleep(0) # force context switch in consumer too
finally:
try:
cv.reset(token)
print(" [consumer] reset(token) OK")
except ValueError as exc:
print(f" [consumer] ✗ reset(token) failed: {exc}")
async def main_consumer_with_sleep():
print("--- Consumer sets context, async for (with sleep), generator raises ---")
print(f"Before: {cv.get()}")
try:
await consumer_with_sleep()
except RuntimeError as e:
print(f" Caught: {e}")
print(f"After: {cv.get()}")
# ---- Empirical assertions (no pytest: asyncio + assert only) ----
_gen_reset_error: list[BaseException] = []
_gen_reset_ok_with_aclosing: list[bool] = []
_consumer_reset_ok: list[bool] = []
_gen_reset_ok_normal_exit: list[bool] = []
_chain_inner_reset_error: list[BaseException] = []
async def _gen_sets_token_consumer_raises():
"""Generator sets contextvar token, yields once; consumer will raise."""
token = cv.set(3.14)
try:
yield 1
finally:
try:
cv.reset(token)
except ValueError as e:
_gen_reset_error.append(e)
async def _gen_sets_token_consumer_raises_record_both():
"""Same as above but records both success and failure of reset (for aclosing)."""
token = cv.set(3.14)
try:
yield 1
finally:
try:
cv.reset(token)
_gen_reset_ok_with_aclosing.append(True)
except ValueError as e:
_gen_reset_error.append(e)
async def _consumer_sets_token_gen_raises():
"""Consumer sets context, async for over generator that yields then raises."""
token = cv.set("evaluator_span")
try:
async for _ in gen_that_yields_then_raises():
pass
finally:
try:
cv.reset(token)
_consumer_reset_ok.append(True)
except ValueError:
_consumer_reset_ok.append(False)
async def _consumer_sets_token_gen_raises_with_sleep():
token = cv.set("evaluator_span")
try:
async for _ in gen_yield_sleep_raise():
await asyncio.sleep(0)
finally:
try:
cv.reset(token)
_consumer_reset_ok.append(True)
except ValueError:
_consumer_reset_ok.append(False)
async def _gen_sets_token_normal_exit():
"""Generator sets token, yields, exits normally (no consumer raise). Scenario 4."""
token = cv.set(3.14)
try:
yield 1
yield 2
finally:
try:
cv.reset(token)
_gen_reset_ok_normal_exit.append(True)
except ValueError as e:
_gen_reset_ok_normal_exit.append(False)
_gen_reset_error.append(e)
async def _inner_b_chain():
"""Inner generator B: sets token, yields. Used inside outer A with aclosing(B)."""
token = cv.set(3.14)
try:
yield 1
finally:
try:
cv.reset(token)
except ValueError as e:
_chain_inner_reset_error.append(e)
async def _outer_a_chain():
"""
Outer generator A (subscription): uses aclosing(inner B), yields B's chunks.
No aclosing(A) at call site → when consumer raises, A is finalized → A.aclose()
runs in finalizer task → B.aclose() runs there → B's finally in wrong context.
"""
inner = _inner_b_chain()
async with aclosing(inner) as stream:
async for x in stream:
yield x
async def run_empirical_tests():
"""
Run all six scenarios and assert expected outcomes. Maps to analysis §4.1 (table)
and §4.4 (chain of generators). No pytest: plain asyncio + assert to avoid confounders.
"""
# 1. Generator sets token, consumer raises (no aclosing) -> generator's reset fails (§4.1 row 1)
_gen_reset_error.clear()
try:
async for _ in _gen_sets_token_consumer_raises():
raise RuntimeError("consumer raises")
except RuntimeError:
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(_gen_reset_error) == 1, f"expected 1 reset error, got {_gen_reset_error}"
assert "different Context" in str(_gen_reset_error[0])
# 2. Generator+aclosing, consumer raises -> generator's reset succeeds (§4.1 row 2)
_gen_reset_error.clear()
_gen_reset_ok_with_aclosing.clear()
try:
async with aclosing(_gen_sets_token_consumer_raises_record_both()) as agen:
async for _ in agen:
raise RuntimeError("consumer raises")
except RuntimeError:
pass
await asyncio.sleep(0)
assert _gen_reset_ok_with_aclosing == [True]
assert len(_gen_reset_error) == 0
# 3. Consumer sets token, generator raises -> consumer's reset succeeds (§4.1 row 3)
_consumer_reset_ok.clear()
try:
await _consumer_sets_token_gen_raises()
except RuntimeError:
pass
assert _consumer_reset_ok == [True]
# 4. Same with sleep in loop -> consumer's reset still succeeds (§4.1 row 3 variant)
_consumer_reset_ok.clear()
try:
await _consumer_sets_token_gen_raises_with_sleep()
except RuntimeError:
pass
assert _consumer_reset_ok == [True]
# 5. Normal exit (no exception) -> generator's reset succeeds (§4.1 scenario 4)
_gen_reset_ok_normal_exit.clear()
_gen_reset_error.clear()
async for _ in _gen_sets_token_normal_exit():
pass
assert _gen_reset_ok_normal_exit == [True]
assert len(_gen_reset_error) == 0
# 6. Chain of generators (§4.4): outer A not closed by caller, A uses aclosing(B).
# Consumer raises -> A finalized -> A.aclose() in finalizer task -> B.aclose()
# there -> B's finally runs in wrong context -> B's reset fails.
_chain_inner_reset_error.clear()
try:
async for _ in _outer_a_chain():
raise RuntimeError("consumer raises")
except RuntimeError:
pass
await asyncio.sleep(0)
await asyncio.sleep(0)
assert len(_chain_inner_reset_error) == 1
assert "different Context" in str(_chain_inner_reset_error[0])
if __name__ == "__main__":
asyncio.run(main())
print()
asyncio.run(main_normal())
print()
asyncio.run(main_consumer_raises())
print()
asyncio.run(main_consumer_with_sleep())
print()
print("--- Empirical assertions (no pytest) ---")
asyncio.run(run_empirical_tests())
print("All assertions passed.")
