55# the root directory of this source tree.
66
77import asyncio
8+ import atexit
9+ import concurrent .futures
810import inspect
911import json
1012import logging # allow-direct-logging
1113import os
14+ import queue
1215import sys
16+ import threading
1317import typing
1418from collections .abc import AsyncGenerator , Generator
1519from enum import Enum
5458
5559T = TypeVar ("T" )
5660
61+ _INIT_TIMEOUT : float = 60.0
62+ _SHUTDOWN_TIMEOUT : float = 10.0
63+ _CLEANUP_TIMEOUT : float = 5.0
64+ _HANG_GUARD_TIMEOUT : float = 600.0
65+ _STREAM_HEARTBEAT_INTERVAL : float = 1.0
66+
5767
5868def convert_pydantic_to_json_value (value : Any ) -> Any :
5969 """Recursively convert Pydantic models, enums, and nested structures to JSON-serializable values.
@@ -156,7 +166,11 @@ def __init__(self, response: FastAPIResponse) -> None:
156166
157167
158168class OGXAsLibraryClient (OgxClient ):
159- """Synchronous client that runs a OGX distribution in-process as a library."""
169+ """Synchronous client that runs a OGX distribution in-process as a library.
170+
171+ This is a sync-on-async implementation wrapping `AsyncOGXAsLibraryClient` class,
172+ starting a daemon loop thread, which will be shut down when the main thread exits.
173+ """
160174
161175 def __init__ (
162176 self ,
@@ -170,24 +184,39 @@ def __init__(
170184 config_path_or_distro_name , custom_provider_registry , provider_data , skip_logger_removal
171185 )
172186 self .provider_data = provider_data
187+ self ._shutdown_lock : threading .Lock = threading .Lock ()
188+ self ._shutdown = False
173189
190+ # stick with one loop and run it in a dedicated daemon thread
174191 self .loop = asyncio .new_event_loop ()
192+ self .loop_thread = threading .Thread (
193+ target = self ._run_event_loop , daemon = True , name = "ogx-lib-sync-client-event-loop"
194+ )
195+ self .loop_thread .start ()
175196
176- # use a new event loop to avoid interfering with the main event loop
177- loop = asyncio .new_event_loop ()
178- asyncio .set_event_loop (loop )
179197 try :
180- loop .run_until_complete (self .async_client .initialize ())
198+ future = asyncio .run_coroutine_threadsafe (self .async_client .initialize (), self .loop )
199+ future .result (timeout = _INIT_TIMEOUT ) # Block until initialization completes + timeout if hangs
200+ except Exception :
201+ self .loop .call_soon_threadsafe (self .loop .stop )
202+ self .loop_thread .join (timeout = _CLEANUP_TIMEOUT )
203+ raise
204+
205+ atexit .register (self .shutdown ) # Safety net: if the user forgets to shutdown properly
206+
207+ def _run_event_loop (self ) -> None :
208+ """Runs forever in the background thread."""
209+ asyncio .set_event_loop (self .loop )
210+ try :
211+ self .loop .run_forever ()
181212 finally :
182- asyncio . set_event_loop ( None )
213+ self . loop . close () # Close the loop when the thread is instructed to stop
183214
184215 def initialize (self ) -> None :
185- """
186- Deprecated method for backward compatibility.
187- """
216+ """Deprecated method for backward compatibility."""
188217 pass
189218
190- def shutdown (self ) -> None :
219+ def shutdown (self , timeout : float = _SHUTDOWN_TIMEOUT ) -> None :
191220 """Shutdown the client and release all resources.
192221
193222 This method should be called when you're done using the client to properly
@@ -197,18 +226,40 @@ def shutdown(self) -> None:
197226
198227 This method is idempotent and can be called multiple times safely.
199228
229+ Args:
230+ timeout: Maximum seconds to wait for graceful shutdown before forcing close.
231+
232+ **IMPORTANT!** `shutdown()` is not safe to call concurrently with requests!
233+ Use the client as a context manager to assure proper shutdown.
234+
200235 Example:
201- client = OGXAsLibraryClient("starter")
202- # ... use the client ...
203- client.shutdown()
236+ with OGXAsLibraryClient("starter") as client:
237+ # ... use the client ...
204238 """
205- loop = self .loop
206- asyncio .set_event_loop (loop )
239+ # Guard against calling shutdown before init finishes, or multiple times
240+ with self ._shutdown_lock :
241+ if self ._shutdown :
242+ return
243+ self ._shutdown = True
244+ if not self .loop .is_running ():
245+ return
246+
247+ future = asyncio .run_coroutine_threadsafe (self .async_client .shutdown (), self .loop )
207248 try :
208- loop .run_until_complete (self .async_client .shutdown ())
209- finally :
210- loop .close ()
211- asyncio .set_event_loop (None )
249+ future .result (timeout = timeout )
250+ except concurrent .futures .TimeoutError :
251+ logger .warning ("Async client shutdown timed out" , timeout = timeout )
252+ future .cancel ()
253+ except Exception as e :
254+ logger .warning ("Unexpected error during async client shutdown" , exception = e )
255+
256+ # Safely instruct the background loop to stop
257+ self .loop .call_soon_threadsafe (self .loop .stop )
258+
259+ # Wait for the thread to actually exit
260+ self .loop_thread .join (timeout = timeout )
261+ if self .loop_thread .is_alive ():
262+ logger .error ("Background event loop thread failed to join (zombie thread)" )
212263
213264 def __enter__ (self ) -> "OGXAsLibraryClient" :
214265 """Enter the context manager.
@@ -227,33 +278,77 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
227278 self .shutdown ()
228279
229280 def request (self , * args : Any , ** kwargs : Any ) -> Any :
230- loop = self .loop
231- asyncio .set_event_loop (loop )
232-
281+ # Route streaming vs non-streaming
233282 if kwargs .get ("stream" ):
283+ return self ._stream_request (* args , ** kwargs )
234284
235- def sync_generator () -> Generator [Any , None , None ]:
236- try :
237- async_stream = loop .run_until_complete (self .async_client .request (* args , ** kwargs ))
238- while True :
239- chunk = loop .run_until_complete (async_stream .__anext__ ())
240- yield chunk
241- except StopAsyncIteration :
242- pass
243- finally :
244- pending = asyncio .all_tasks (loop )
245- if pending :
246- loop .run_until_complete (asyncio .gather (* pending , return_exceptions = True ))
247-
248- return sync_generator ()
249- else :
285+ coro = self .async_client .request (* args , ** kwargs )
286+ future = asyncio .run_coroutine_threadsafe (coro , self .loop )
287+ # the giant timeout here is to prevent it from hanging forever:
288+ return future .result (timeout = _HANG_GUARD_TIMEOUT )
289+
290+ def _stream_request (self , * args : Any , ** kwargs : Any ) -> Generator [Any , None , None ]:
291+ """Thread-safe synchronous generator wrapper around an async generator."""
292+ # 32 chunks of buffering. LLM token rate makes OOM from unbounded queue unlikely
293+ # but a bound prevents runaway memory if the consumer stalls.
294+ q : queue .Queue = queue .Queue (maxsize = 32 )
295+
296+ async def _consume () -> None :
297+ async_gen = None
250298 try :
251- result = loop .run_until_complete (self .async_client .request (* args , ** kwargs ))
299+ async_gen = await self .async_client .request (* args , ** kwargs )
300+ async for chunk in async_gen :
301+ while True :
302+ try :
303+ q .put_nowait (("chunk" , chunk ))
304+ break
305+ except queue .Full :
306+ await asyncio .sleep (0.01 )
307+
308+ while True :
309+ try :
310+ q .put_nowait (("done" , None ))
311+ break
312+ except queue .Full :
313+ await asyncio .sleep (0.01 )
314+
315+ except asyncio .CancelledError :
316+ pass
317+ except Exception as err :
318+ while True :
319+ try :
320+ q .put_nowait (("error" , err ))
321+ break
322+ except queue .Full :
323+ await asyncio .sleep (0.01 )
252324 finally :
253- pending = asyncio .all_tasks (loop )
254- if pending :
255- loop .run_until_complete (asyncio .gather (* pending , return_exceptions = True ))
256- return result
325+ if async_gen is not None :
326+ await async_gen .aclose ()
327+
328+ future = asyncio .run_coroutine_threadsafe (_consume (), self .loop )
329+
330+ try :
331+ while True :
332+ try :
333+ # Timeout prevents the sync thread from hanging forever if the loop dies or shutdown is called.
334+ msg_type , payload = q .get (timeout = 1.0 )
335+ except queue .Empty as err :
336+ with self ._shutdown_lock :
337+ if self ._shutdown :
338+ raise RuntimeError ("Client was shut down during streaming" ) from err
339+
340+ if not self .loop .is_running ():
341+ raise RuntimeError ("Event loop crashed during streaming" ) from err
342+ continue
343+
344+ if msg_type == "chunk" :
345+ yield payload
346+ elif msg_type == "error" :
347+ raise payload
348+ elif msg_type == "done" :
349+ break
350+ finally :
351+ future .cancel ()
257352
258353
259354class AsyncOGXAsLibraryClient (AsyncOgxClient ):
0 commit comments