Merge pull request #1000 from MervinPraison/claude/fix-telemetry-termination-20250718

Fix: Comprehensive telemetry cleanup to prevent agent termination hang

Author: MervinPraison

src/praisonai-agents/praisonaiagents/agent/agent.py

@@ -1958,12 +1958,10 @@ def _cleanup_telemetry(self):
         """Clean up telemetry system to ensure proper program termination."""
         try:
             # Import here to avoid circular imports
-            from ..telemetry import get_telemetry
+            from ..telemetry import force_shutdown_telemetry
 
-            # Get the global telemetry instance and shut it down
-            telemetry = get_telemetry()
-            if telemetry and hasattr(telemetry, 'shutdown'):
-                telemetry.shutdown()
+            # Force shutdown of telemetry system with comprehensive cleanup
+            force_shutdown_telemetry()
         except Exception as e:
             # Log error but don't fail the execution
             logging.debug(f"Error cleaning up telemetry: {e}")

src/praisonai-agents/praisonaiagents/telemetry/__init__.py

@@ -24,6 +24,7 @@
     'get_telemetry',
     'enable_telemetry',
     'disable_telemetry',
+    'force_shutdown_telemetry',
     'MinimalTelemetry',
     'TelemetryCollector',  # For backward compatibility
 ]
@@ -47,6 +48,12 @@ def disable_telemetry():
     _disable_telemetry()
 
 
+def force_shutdown_telemetry():
+    """Force shutdown of telemetry system with comprehensive cleanup."""
+    from .telemetry import force_shutdown_telemetry as _force_shutdown_telemetry
+    _force_shutdown_telemetry()
+
+
 # Auto-instrumentation and cleanup setup
 _initialized = False
 _atexit_registered = False

src/praisonai-agents/praisonaiagents/telemetry/telemetry.py

@@ -318,6 +318,7 @@ def flush(self):
     def shutdown(self):
         """
         Shutdown telemetry and ensure all events are sent.
+        Forces proper cleanup of background threads to prevent hanging.
         """
         if not self.enabled:
             return
@@ -330,8 +331,55 @@ def shutdown(self):
             try:
                 # Force a synchronous flush before shutdown
                 self._posthog.flush()
+                
+                # Get the PostHog client's internal thread pool for cleanup
+                if hasattr(self._posthog, '_thread_pool'):
+                    thread_pool = self._posthog._thread_pool
+                    if thread_pool:
+                        try:
+                            # Stop accepting new tasks
+                            thread_pool.shutdown(wait=False)
+                            # Wait for threads to finish with timeout
+                            thread_pool.shutdown(wait=True)
+                        except:
+                            pass
+                
+                # Force shutdown of any remaining threads
+                if hasattr(self._posthog, '_consumer'):
+                    try:
+                        self._posthog._consumer.flush()
+                        self._posthog._consumer.shutdown()
+                    except:
+                        pass
+                
+                # Standard shutdown
                 self._posthog.shutdown()
-            except:
+                
+                # Additional cleanup - force thread termination
+                import threading
+                import time
+                
+                # Wait up to 2 seconds for threads to terminate
+                max_wait = 2.0
+                start_time = time.time()
+                
+                while time.time() - start_time < max_wait:
+                    # Check for any PostHog related threads
+                    posthog_threads = [
+                        t for t in threading.enumerate() 
+                        if t != threading.current_thread() 
+                        and not t.daemon
+                        and ('posthog' in t.name.lower() or 'analytics' in t.name.lower())
+                    ]
+                    
+                    if not posthog_threads:
+                        break
+                        
+                    time.sleep(0.1)
+                
+            except Exception as e:
+                # Log the error but don't fail shutdown
+                self.logger.debug(f"Error during PostHog shutdown: {e}")
                 pass
 
 
@@ -361,6 +409,41 @@ def disable_telemetry():
         _telemetry_instance = MinimalTelemetry(enabled=False)
 
 
+def force_shutdown_telemetry():
+    """
+    Force shutdown of telemetry system with comprehensive cleanup.
+    This function ensures proper termination of all background threads.
+    """
+    global _telemetry_instance
+    if _telemetry_instance:
+        _telemetry_instance.shutdown()
+        
+        # Additional cleanup - wait for all threads to finish
+        import threading
+        import time
+        
+        # Wait up to 3 seconds for any remaining threads to finish
+        max_wait = 3.0
+        start_time = time.time()
+        
+        while time.time() - start_time < max_wait:
+            # Check for any analytics/telemetry related threads
+            analytics_threads = [
+                t for t in threading.enumerate() 
+                if t != threading.current_thread() 
+                and not t.daemon
+                and any(keyword in t.name.lower() for keyword in ['posthog', 'analytics', 'telemetry', 'consumer'])
+            ]
+            
+            if not analytics_threads:
+                break
+                
+            time.sleep(0.1)
+        
+        # Reset the global instance
+        _telemetry_instance = None
+
+
 def enable_telemetry():
     """Programmatically enable telemetry (if not disabled by environment)."""
     global _telemetry_instance