make sure library client is cool

This commit is contained in:
Ashwin Bharambe 2025-03-08 22:50:40 -08:00
parent d89ef35151
commit 1bd002f39a

View file

@ -52,15 +52,12 @@ def preserve_headers_context_async_generator(gen: AsyncGenerator[T, None]) -> As
async def wrapper(): async def wrapper():
while True: while True:
# Set context before each anext() call # Set context before each anext() call
token = _provider_data_var.set(context_value) _ = _provider_data_var.set(context_value)
try: try:
item = await gen.__anext__() item = await gen.__anext__()
yield item yield item
except StopAsyncIteration: except StopAsyncIteration:
break break
finally:
# Restore the previous value
_provider_data_var.reset(token)
return wrapper() return wrapper()