fix: getting_started notebook

llama_stack/distribution/utils/context.py

@@ -18,15 +18,19 @@ def preserve_contexts_async_generator(
     This is needed because we start a new asyncio event loop for each streaming request,
     and we need to preserve the context across the event loop boundary.
     """
+    # Capture initial context values
+    initial_context_values = {context_var.name: context_var.get() for context_var in context_vars}
 
     async def wrapper() -> AsyncGenerator[T, None]:
         while True:
             try:
-                item = await gen.__anext__()
-                context_values = {context_var.name: context_var.get() for context_var in context_vars}
-                yield item
+                # Restore context values before any await
                 for context_var in context_vars:
-                    _ = context_var.set(context_values[context_var.name])
+                    context_var.set(initial_context_values[context_var.name])
+
+                item = await gen.__anext__()
+                yield item
+
             except StopAsyncIteration:
                 break