@@ -1929,19 +1929,44 @@ async def _achat_completion(self, response, tools, reasoning_steps=False):
19291929
19301930 async def astart (self , prompt : str , ** kwargs ):
19311931 """Async version of start method"""
1932- return await self .achat (prompt , ** kwargs )
1932+ try :
1933+ result = await self .achat (prompt , ** kwargs )
1934+ return result
1935+ finally :
1936+ # Ensure proper cleanup of telemetry system to prevent hanging
1937+ self ._cleanup_telemetry ()
19331938
19341939 def run (self ):
19351940 """Alias for start() method"""
19361941 return self .start ()
19371942
1943+ def _cleanup_telemetry (self ):
1944+ """Clean up telemetry system to ensure proper program termination."""
1945+ try :
1946+ # Import here to avoid circular imports
1947+ from ..telemetry import get_telemetry
1948+
1949+ # Get the global telemetry instance and shut it down
1950+ telemetry = get_telemetry ()
1951+ if telemetry and hasattr (telemetry , 'shutdown' ):
1952+ telemetry .shutdown ()
1953+ except Exception as e :
1954+ # Log error but don't fail the execution
1955+ logging .debug (f"Error cleaning up telemetry: { e } " )
1956+
19381957 def start (self , prompt : str , ** kwargs ):
19391958 """Start the agent with a prompt. This is a convenience method that wraps chat()."""
1940- # Check if streaming is enabled and user wants streaming chunks
1941- if self .stream and kwargs .get ('stream' , True ):
1942- return self ._start_stream (prompt , ** kwargs )
1943- else :
1944- return self .chat (prompt , ** kwargs )
1959+ try :
1960+ # Check if streaming is enabled and user wants streaming chunks
1961+ if self .stream and kwargs .get ('stream' , True ):
1962+ result = self ._start_stream (prompt , ** kwargs )
1963+ return result
1964+ else :
1965+ result = self .chat (prompt , ** kwargs )
1966+ return result
1967+ finally :
1968+ # Ensure proper cleanup of telemetry system to prevent hanging
1969+ self ._cleanup_telemetry ()
19451970
19461971 def _start_stream (self , prompt : str , ** kwargs ):
19471972 """Generator method that yields streaming chunks from the agent."""
0 commit comments