From 0fdb15bcc74f236cd4e6ac4291a361a08b6bf1b3 Mon Sep 17 00:00:00 2001
From: Dinesh Yeduguru <yvdinesh@gmail.com>
Date: Wed, 12 Mar 2025 13:26:23 -0700
Subject: [PATCH] fix: fix build error in context.py (#1595)

# What does this PR do?
This fixes the build error


## Test Plan
pre-commit run --all-files
check for merge
conflicts................................................Passed
trim trailing
whitespace.................................................Passed
check for added large
files..............................................Passed
fix end of
files.........................................................Passed
Insert license in
comments...............................................Passed

ruff.....................................................................Passed

ruff-format..............................................................Passed

blacken-docs.............................................................Passed

uv-lock..................................................................Passed

uv-export................................................................Passed

mypy.....................................................................Passed
Distribution Template
Codegen............................................Passed
---
 llama_stack/distribution/utils/context.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llama_stack/distribution/utils/context.py b/llama_stack/distribution/utils/context.py
index 107ce7127..2f32afba2 100644
--- a/llama_stack/distribution/utils/context.py
+++ b/llama_stack/distribution/utils/context.py
@@ -19,7 +19,7 @@ def preserve_contexts_async_generator(
     and we need to preserve the context across the event loop boundary.
     """
 
-    async def wrapper():
+    async def wrapper() -> AsyncGenerator[T, None]:
         while True:
             try:
                 item = await gen.__anext__()