From b0a6bcdad7dba4f0b0938909f4970063fbc8e1a1 Mon Sep 17 00:00:00 2001
From: "Michael J. Sullivan" <sully@msully.net>
Date: Fri, 10 Jul 2026 12:02:32 -0700
Subject: [PATCH] Generalize the custom contextmanager decorator to be an async
 decorator too

Allows a bit more ergonomics because it lets you mix them in one async
with block.

(I'm probably about to add another
---
 src/ai/types/messages.py     |  2 +-
 src/ai/util.py               | 34 +++++++++++++++++++++++++++-------
 tests/types/test_messages.py |  7 ++++++-
 3 files changed, 34 insertions(+), 9 deletions(-)

diff --git a/src/ai/types/messages.py b/src/ai/types/messages.py
index 0a848440..5acd340b 100644
--- a/src/ai/types/messages.py
+++ b/src/ai/types/messages.py
@@ -48,7 +48,7 @@ def _resolve_random(source: RandomSource) -> random.Random:
     return source if isinstance(source, random.Random) else source()
 
 
-@util.contextmanager_with_async_decorator
+@util.contextmanager_any_sync
 def use_random(source: RandomSource) -> Iterator[None]:
     """Draw message/part ids from ``source`` within this context.
 
diff --git a/src/ai/util.py b/src/ai/util.py
index ddc9ef8c..733db487 100644
--- a/src/ai/util.py
+++ b/src/ai/util.py
@@ -22,7 +22,7 @@
     from types import TracebackType
 
 
-class ContextManagerWithAsyncDecorator[T](Protocol):
+class ContextManagerAnySync[T](Protocol):
     def __enter__(self) -> T: ...
 
     def __exit__(
@@ -32,6 +32,15 @@ def __exit__(
         traceback: TracebackType | None,
     ) -> bool | None: ...
 
+    async def __aenter__(self) -> T: ...
+
+    async def __aexit__(
+        self,
+        typ: type[BaseException] | None,
+        value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None: ...
+
     @overload
     def __call__[**P, R](
         self, func: Callable[P, Coroutine[Any, Any, R]]
@@ -141,9 +150,20 @@ async def __aexit__(
         return False
 
 
-class _GeneratorContextManagerWithAsyncDecorator[T](
+class _GeneratorContextManagerAnySync[T](
     contextlib._GeneratorContextManager[T]
 ):
+    async def __aenter__(self) -> T:
+        return self.__enter__()
+
+    async def __aexit__(
+        self,
+        typ: type[BaseException] | None,
+        value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> bool | None:
+        return self.__exit__(typ, value, traceback)
+
     def __call__[_F: Callable[..., Any]](self, func: _F) -> _F:
         if inspect.iscoroutinefunction(func):
 
@@ -161,16 +181,16 @@ def inner(*args: Any, **kwds: Any) -> Any:
         return cast("_F", inner)
 
 
-def contextmanager_with_async_decorator[**P, T](
+def contextmanager_any_sync[**P, T](
     func: Callable[P, Iterator[T]],
-) -> Callable[P, ContextManagerWithAsyncDecorator[T]]:
-    """@contextmanager decorator but the result can be a decorator for async."""
+) -> Callable[P, ContextManagerAnySync[T]]:
+    """@contextmanager decorator but the result is also usable in async."""
 
     @functools.wraps(func)
     def helper(
         *args: P.args, **kwds: P.kwargs
-    ) -> _GeneratorContextManagerWithAsyncDecorator[T]:
-        return _GeneratorContextManagerWithAsyncDecorator(
+    ) -> _GeneratorContextManagerAnySync[T]:
+        return _GeneratorContextManagerAnySync(
             cast("Callable[..., Generator[T, None, None]]", func), args, kwds
         )
 
diff --git a/tests/types/test_messages.py b/tests/types/test_messages.py
index 117ccc69..f07325a2 100644
--- a/tests/types/test_messages.py
+++ b/tests/types/test_messages.py
@@ -455,10 +455,15 @@ def test_use_random_overrides_and_restores() -> None:
     assert messages.generate_id("msg").startswith("msg_")
 
 
-async def test_use_random_decorator_handles_async_functions() -> None:
+async def test_use_random_overrides_and_restores_async() -> None:
     with messages.use_random(random.Random(0)):
         expected = messages.generate_id("msg")
 
+    # Works across an await...
+    async with messages.use_random(random.Random(0)):
+        await asyncio.sleep(0)
+        assert messages.generate_id("msg") == expected
+
     # Works as a decorator on an async fn, resolving the factory per call.
     @messages.use_random(lambda: random.Random(0))
     async def build() -> str: