Generalize the custom contextmanager decorator to be an async decorator too

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.
 

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
         )
 

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: