From 70ff226b6ae404d34d66c188c2b84bad9377010f Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Sun, 9 Mar 2025 16:17:27 -0700 Subject: [PATCH] fix(library_client): ensure pending asyncio tasks like generator athrow are executed --- llama_stack/distribution/library_client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/llama_stack/distribution/library_client.py b/llama_stack/distribution/library_client.py index ab8ff60fa..5dc70bb67 100644 --- a/llama_stack/distribution/library_client.py +++ b/llama_stack/distribution/library_client.py @@ -163,6 +163,9 @@ class LlamaStackAsLibraryClient(LlamaStackClient): except StopAsyncIteration: pass finally: + pending = asyncio.all_tasks(loop) + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) loop.close() return sync_generator() @@ -383,7 +386,6 @@ class AsyncLlamaStackAsLibraryClient(AsyncLlamaStackClient): # Wrap the generator to preserve context across iterations wrapped_gen = preserve_headers_context_async_generator(gen()) - mock_response = httpx.Response( status_code=httpx.codes.OK, content=wrapped_gen,