fix: Add atexit handlers to flush callbacks for async completions (#16487)

Fixes #16486

## Problem
Callbacks configured via litellm.success_callback (e.g., PostHog, LangSmith)
were not being invoked for litellm.acompletion() in short-lived scripts.
The callbacks worked correctly for synchronous completions but async
completions would queue callbacks that were lost when the script exited
before GLOBAL_LOGGING_WORKER could process them.

Root cause: asyncio.run() closes the event loop immediately after the
async function completes, preventing the background worker from processing
queued callbacks.

## Solution
Implemented a two-level atexit handler approach:

1. GLOBAL_LOGGING_WORKER atexit handler (logging_worker.py):
   - Processes remaining callbacks from queue before exit
   - Creates new event loop to run pending coroutines synchronously
   - Applies time and iteration limits to prevent blocking shutdown

2. Integration-specific atexit handlers (posthog.py as example):
   - Flushes internal queue to external service
   - Uses synchronous HTTP client for reliable delivery
   - Each integration needs its own handler due to varying sync APIs

## Changes
- litellm/litellm_core_utils/logging_worker.py:
  - Added _flush_on_exit() method
  - Registered atexit handler in __init__
  - Processes up to MAX_ITERATIONS_TO_CLEAR_QUEUE events
  - Time-limited to MAX_TIME_TO_CLEAR_QUEUE seconds

- litellm/integrations/posthog.py:
  - Added _flush_on_exit() method
  - Registered atexit handler in __init__
  - Groups events by credentials for batch sending
  - Uses sync_client for reliable HTTP delivery

- tests/logging_callback_tests/test_posthog.py:
  - Added test_async_callback_atexit_handler_exists()
  - Added test_posthog_atexit_flushes_internal_queue()
  - Added test_sync_callback_not_affected_by_atexit()

## Testing
- All existing tests pass
- Manual end-to-end testing confirms fix:
  - Async events now arrive in PostHog
  - Sync events continue working (no regression)
- Unit tests verify atexit handlers registered and functional

## Impact
- Fixes async callback delivery for ALL integrations using GLOBAL_LOGGING_WORKER
- No breaking changes - only adds missing functionality
- Sync path unchanged - no performance impact

litellm/integrations/posthog.py

@@ -10,6 +10,7 @@ For batching specific details see CustomBatchLogger class
 """
 
 import asyncio
+import atexit
 import os
 from typing import Any, Dict, Optional, Tuple
 
@@ -55,7 +56,10 @@ class PostHogLogger(CustomBatchLogger):
             self._async_initialized = False
             self.flush_lock = None
             self.log_queue = []
-            
+
+            # Register cleanup handler to flush internal queue on exit
+            atexit.register(self._flush_on_exit)
+
             super().__init__(
                 **kwargs, flush_lock=None, batch_size=POSTHOG_MAX_BATCH_SIZE
             )
@@ -377,3 +381,58 @@ class PostHogLogger(CustomBatchLogger):
         if obj is None or not hasattr(obj, 'get'):
             return default
         return obj.get(key, default)
+
+    def _flush_on_exit(self):
+        """
+        Flush remaining events from internal log_queue before process exit.
+        Called automatically via atexit handler.
+
+        This works in conjunction with GLOBAL_LOGGING_WORKER's atexit handler:
+        1. GLOBAL_LOGGING_WORKER atexit invokes pending callbacks
+        2. Callbacks add events to this logger's internal log_queue
+        3. This atexit handler flushes the internal queue to PostHog
+        """
+        if not self.log_queue:
+            return
+
+        verbose_logger.debug(
+            f"PostHog: Flushing {len(self.log_queue)} remaining events on exit"
+        )
+
+        try:
+            # Group events by credentials (same logic as async_send_batch)
+            batches_by_credentials: Dict[Tuple[str, str], list] = {}
+            for item in self.log_queue:
+                key = (item["api_key"], item["api_url"])
+                if key not in batches_by_credentials:
+                    batches_by_credentials[key] = []
+                batches_by_credentials[key].append(item["event"])
+
+            # Send each batch synchronously using sync_client
+            for (api_key, api_url), events in batches_by_credentials.items():
+                headers = {
+                    "Content-Type": "application/json",
+                }
+
+                payload = self._create_posthog_payload(events, api_key)
+                capture_url = f"{api_url.rstrip('/')}/batch/"
+
+                response = self.sync_client.post(
+                    url=capture_url,
+                    json=payload,
+                    headers=headers,
+                )
+                response.raise_for_status()
+
+                if response.status_code != 200:
+                    verbose_logger.error(
+                        f"PostHog: Failed to flush on exit - status {response.status_code}"
+                    )
+
+            verbose_logger.debug(
+                f"PostHog: Successfully flushed {len(self.log_queue)} events on exit"
+            )
+            self.log_queue.clear()
+
+        except Exception as e:
+            verbose_logger.error(f"PostHog: Error flushing events on exit: {str(e)}")

litellm/litellm_core_utils/logging_worker.py

@@ -1,4 +1,5 @@
 import asyncio
+import atexit
 import contextlib
 import contextvars
 from typing import Coroutine, Optional
@@ -43,6 +44,9 @@ class LoggingWorker:
         self._queue: Optional[asyncio.Queue[LoggingTask]] = None
         self._worker_task: Optional[asyncio.Task] = None
 
+        # Register cleanup handler to flush remaining events on exit
+        atexit.register(self._flush_on_exit)
+
     def _ensure_queue(self) -> None:
         """Initialize the queue if it doesn't exist."""
         if self._queue is None:
@@ -154,6 +158,61 @@ class LoggingWorker:
             except asyncio.QueueEmpty:
                 break
 
+    def _flush_on_exit(self):
+        """
+        Flush remaining events synchronously before process exit.
+        Called automatically via atexit handler.
+
+        This ensures callbacks queued by async completions are processed
+        even when the script exits before the worker loop can handle them.
+        """
+        if self._queue is None:
+            verbose_logger.debug("[LoggingWorker] atexit: No queue initialized")
+            return
+
+        if self._queue.empty():
+            verbose_logger.debug("[LoggingWorker] atexit: Queue is empty")
+            return
+
+        queue_size = self._queue.qsize()
+        verbose_logger.info(f"[LoggingWorker] atexit: Flushing {queue_size} remaining events...")
+
+        # Create a new event loop since the original is closed
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+
+        try:
+            # Process remaining queue items with time limit
+            processed = 0
+            start_time = loop.time()
+
+            while not self._queue.empty() and processed < self.MAX_ITERATIONS_TO_CLEAR_QUEUE:
+                if loop.time() - start_time >= self.MAX_TIME_TO_CLEAR_QUEUE:
+                    verbose_logger.warning(
+                        f"[LoggingWorker] atexit: Reached time limit ({self.MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush"
+                    )
+                    break
+
+                try:
+                    task = self._queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    break
+
+                # Run the coroutine synchronously in new loop
+                # Note: We run the coroutine directly, not via create_task,
+                # since we're in a new event loop context
+                try:
+                    loop.run_until_complete(task["coroutine"])
+                    processed += 1
+                except Exception as e:
+                    # Silent failure to not break user's program
+                    verbose_logger.debug(f"[LoggingWorker] atexit: Error flushing callback: {e}")
+
+            verbose_logger.info(f"[LoggingWorker] atexit: Successfully flushed {processed} events!")
+
+        finally:
+            loop.close()
+
 
 # Global instance for backward compatibility
 GLOBAL_LOGGING_WORKER = LoggingWorker()

tests/logging_callback_tests/test_posthog.py

@@ -304,3 +304,114 @@ async def test_dynamic_credentials():
     api_key, api_url = posthog_logger._get_credentials_for_request(kwargs)
     assert api_key == "test_key"  # falls back to env var
     assert api_url == "https://another.posthog.com"
+
+
+def test_async_callback_atexit_handler_exists():
+    """
+    Test that atexit handlers are properly registered.
+
+    This test verifies that both GLOBAL_LOGGING_WORKER and PostHogLogger
+    register atexit handlers for flushing pending events.
+
+    The actual functionality is validated by end-to-end tests (test_async_only.py)
+    since unit testing atexit behavior across event loop boundaries is complex.
+    """
+    import atexit
+    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+
+    # Verify GLOBAL_LOGGING_WORKER has _flush_on_exit method
+    assert hasattr(GLOBAL_LOGGING_WORKER, '_flush_on_exit'), \
+        "GLOBAL_LOGGING_WORKER should have _flush_on_exit method"
+
+    # Verify PostHogLogger has _flush_on_exit method
+    posthog_logger = PostHogLogger()
+    assert hasattr(posthog_logger, '_flush_on_exit'), \
+        "PostHogLogger should have _flush_on_exit method"
+
+    # Verify method can be called without crashing (with empty queue)
+    # This tests the early return paths
+    GLOBAL_LOGGING_WORKER._flush_on_exit()
+    posthog_logger._flush_on_exit()
+
+
+@pytest.mark.asyncio
+async def test_posthog_atexit_flushes_internal_queue():
+    """
+    Test that PostHog's atexit handler flushes its internal log_queue.
+
+    This works in conjunction with GLOBAL_LOGGING_WORKER:
+    1. GLOBAL_LOGGING_WORKER atexit invokes pending callbacks
+    2. Callbacks add events to PostHog's internal log_queue
+    3. PostHog's atexit flushes log_queue via HTTP POST
+    """
+    from unittest.mock import Mock, patch
+    import httpx
+
+    posthog_logger = PostHogLogger()
+
+    # Add mock events to internal queue (simulating what callbacks do)
+    standard_payload = create_standard_logging_payload()
+    kwargs = {"standard_logging_object": standard_payload}
+    event_payload = posthog_logger.create_posthog_event_payload(kwargs)
+
+    posthog_logger.log_queue.append({
+        "event": event_payload,
+        "api_key": "test_key",
+        "api_url": "https://app.posthog.com"
+    })
+
+    assert len(posthog_logger.log_queue) == 1, "Queue should have 1 event"
+
+    # Mock the sync HTTP client to avoid real API calls
+    with patch.object(posthog_logger.sync_client, 'post') as mock_post:
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.raise_for_status = Mock()
+        mock_post.return_value = mock_response
+
+        # Trigger atexit flush
+        posthog_logger._flush_on_exit()
+
+        # Verify HTTP POST was called
+        assert mock_post.called, "HTTP POST should be called during flush"
+        assert len(posthog_logger.log_queue) == 0, "Queue should be empty after flush"
+
+        # Verify correct endpoint was called
+        call_args = mock_post.call_args
+        assert "/batch/" in call_args.kwargs['url'], "Should POST to /batch/ endpoint"
+
+
+@pytest.mark.asyncio
+async def test_sync_callback_not_affected_by_atexit():
+    """
+    Regression test: ensure sync completions still work immediately.
+
+    Sync callbacks should be invoked immediately during completion(),
+    not deferred to atexit. This test verifies atexit handlers don't
+    interfere with the sync path.
+    """
+    from unittest.mock import Mock, patch
+
+    # Track when callback is invoked
+    callback_invoked_immediately = False
+
+    def mock_log_success(self, kwargs, response_obj, start_time, end_time):
+        nonlocal callback_invoked_immediately
+        callback_invoked_immediately = True
+
+    with patch.object(PostHogLogger, 'log_success_event', mock_log_success):
+        with patch('httpx.Client.post') as mock_post:
+            mock_response = Mock()
+            mock_response.status_code = 200
+            mock_response.raise_for_status = Mock()
+            mock_post.return_value = mock_response
+
+            posthog_logger = PostHogLogger()
+            standard_payload = create_standard_logging_payload()
+            kwargs = {"standard_logging_object": standard_payload}
+
+            # Call sync method directly (simulates what completion() does)
+            posthog_logger.log_success_event(kwargs, None, 0.0, 0.0)
+
+            # Callback should be invoked immediately, not queued for atexit
+            assert callback_invoked_immediately, "Sync callback should be invoked immediately"