diff --git a/llama_stack/distribution/utils/context.py b/llama_stack/distribution/utils/context.py
index fcc72161d..c34079ac6 100644
--- a/llama_stack/distribution/utils/context.py
+++ b/llama_stack/distribution/utils/context.py
@@ -29,6 +29,11 @@ def preserve_contexts_async_generator(
                     context_var.set(initial_context_values[context_var.name])
 
                 item = await gen.__anext__()
+
+                # Update our tracked values with any changes made during this iteration
+                for context_var in context_vars:
+                    initial_context_values[context_var.name] = context_var.get()
+
                 yield item
 
             except StopAsyncIteration:
diff --git a/llama_stack/distribution/utils/tests/test_context.py b/tests/unit/distribution/test_context.py
similarity index 100%
rename from llama_stack/distribution/utils/tests/test_context.py
rename to tests/unit/distribution/test_context.py