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():
while True:
# Set context before each anext() call
token = _provider_data_var.set(context_value)
_ = _provider_data_var.set(context_value)
try:
item = await gen.__anext__()
yield item
except StopAsyncIteration:
break
finally:
# Restore the previous value
_provider_data_var.reset(token)
return wrapper()