fix: add thread safety to SSE connection for realtime flags

- Import threading module at top of client.py
- Add threading.Lock (_sse_lock) for sse_connected flag
- Wrap all sse_connected reads/writes with lock
- Store SSE response and call close() for graceful shutdown
- Ensure thread-safe cleanup in _close_sse_connection()

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

Author: gustavohstrassburger

posthog/client.py

@@ -2,6 +2,7 @@
 import logging
 import os
 import sys
+import threading
 from datetime import datetime, timedelta
 from typing import Any, Dict, Optional, Union
 from typing_extensions import Unpack
@@ -255,6 +256,7 @@ def __init__(
         self.sse_connection = None  # type: Optional[Any]
         self.sse_response = None
         self.sse_connected = False
+        self._sse_lock = threading.Lock()
 
         self.capture_exception_code_variables = capture_exception_code_variables
         self.code_variables_mask_patterns = (
@@ -2245,9 +2247,11 @@ def _setup_sse_connection(self):
             )
             return
 
-        if self.sse_connected:
-            self.log.debug("[FEATURE FLAGS] SSE connection already established")
-            return
+        with self._sse_lock:
+            if self.sse_connected:
+                self.log.debug("[FEATURE FLAGS] SSE connection already established")
+                return
+            self.sse_connected = True
 
         try:
             import threading
@@ -2274,16 +2278,17 @@ def sse_listener():
                         self.log.warning(
                             f"[FEATURE FLAGS] SSE connection failed with status {response.status_code}"
                         )
-                        self.sse_connected = False
+                        with self._sse_lock:
+                            self.sse_connected = False
                         return
 
-                    self.sse_connected = True
                     self.log.debug("[FEATURE FLAGS] SSE connection established")
 
                     # Process the stream line by line, checking sse_connected periodically
                     for line in response.iter_lines():
-                        if not self.sse_connected:
-                            break
+                        with self._sse_lock:
+                            if not self.sse_connected:
+                                break
 
                         if not line:
                             continue
@@ -2304,7 +2309,8 @@ def sse_listener():
                     self.log.warning(
                         f"[FEATURE FLAGS] SSE connection error: {e}. Reconnecting in 5 seconds..."
                     )
-                    self.sse_connected = False
+                    with self._sse_lock:
+                        self.sse_connected = False
 
                     # Attempt to reconnect after 5 seconds if realtime_flags is still enabled
                     if self.realtime_flags:
@@ -2338,7 +2344,8 @@ def _close_sse_connection(self):
         """
         if self.sse_connection:
             self.log.debug("[FEATURE FLAGS] Closing SSE connection")
-            self.sse_connected = False
+            with self._sse_lock:
+                self.sse_connected = False
 
             # Close the response to interrupt iter_lines()
             if self.sse_response:

posthog/test/test_realtime_feature_flags.py

@@ -132,7 +132,9 @@ def test_process_flag_update(self, mock_get):
 
         # Verify flag was updated
         self.assertIn("test-flag", client.feature_flags_by_key)
-        self.assertEqual(client.feature_flags_by_key["test-flag"]["name"], "Updated Test Flag")
+        self.assertEqual(
+            client.feature_flags_by_key["test-flag"]["name"], "Updated Test Flag"
+        )
         self.assertFalse(client.feature_flags_by_key["test-flag"]["active"])
 
         # Cleanup
@@ -289,7 +291,6 @@ def test_sse_cleanup_on_shutdown(self, mock_get):
         self.assertFalse(client.sse_connected)
         self.assertIsNone(client.sse_connection)
 
-
     @mock.patch("posthog.client.get")
     def test_on_feature_flags_update_callback(self, mock_get):
         """Test that the callback is called when flags are updated"""
@@ -306,9 +307,7 @@ def test_on_feature_flags_update_callback(self, mock_get):
         callback_calls = []
 
         def flag_update_callback(flag_key, flag_data):
-            callback_calls.append(
-                {"flag_key": flag_key, "flag_data": flag_data}
-            )
+            callback_calls.append({"flag_key": flag_key, "flag_data": flag_data})
 
         client = Client(
             FAKE_TEST_API_KEY,