diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index c609d30ccff..468b1a441fb 100644 --- a/litellm/integrations/posthog.py +++ b/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)}") diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 3c475f133a8..20f0d70160a 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/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() diff --git a/tests/logging_callback_tests/test_posthog.py b/tests/logging_callback_tests/test_posthog.py index addd1c4917b..ddc60885f3b 100644 --- a/tests/logging_callback_tests/test_posthog.py +++ b/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"