fix: return the uuid the wire event carries

Generated-By: PostHog Code
Task-Id: 4b5e5cda-81f1-41a2-8fcd-4f335d19217e

Author: carlos-marchal-ph

posthog/client.py

@@ -2147,6 +2147,21 @@ def _reinit_after_fork(self):
         else:
             self.poller = None
 
+    def _normalize_event_uuid(self, msg):
+        # type: (...) -> None
+        """Ensure `msg["uuid"]` is a valid uuid string, generating one if missing or invalid."""
+        if "uuid" in msg:
+            uuid = msg.pop("uuid")
+            if uuid is not None:
+                try:
+                    msg["uuid"] = _stringify_event_uuid(uuid)
+                except ValueError as e:
+                    self.log.error("%s Falling back to a generated UUID.", e)
+
+        if "uuid" not in msg:
+            # Always send a uuid, so we can always return one
+            msg["uuid"] = stringify_id(uuid4())
+
     def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None):
         # type: (...) -> Optional[str]
         """Push a new `msg` onto a lane's queue (analytics when unspecified), return the event uuid or None."""
@@ -2165,19 +2180,7 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None):
         timestamp = guess_timezone(timestamp)
         msg["timestamp"] = timestamp.isoformat()
 
-        if "uuid" in msg:
-            uuid = msg.pop("uuid")
-            if uuid is not None:
-                try:
-                    msg["uuid"] = _stringify_event_uuid(uuid)
-                except ValueError as e:
-                    self.log.error("%s Falling back to a generated UUID.", e)
-
-        if "uuid" not in msg:
-            # Always send a uuid, so we can always return one
-            msg["uuid"] = stringify_id(uuid4())
-
-        sent_uuid = msg["uuid"]
+        self._normalize_event_uuid(msg)
 
         if not msg.get("properties"):
             msg["properties"] = {}
@@ -2223,6 +2226,11 @@ def _enqueue(self, msg, disable_geoip, lane=None, property_allowlist=None):
                 self.log.exception(f"Error in before_send callback: {e}")
                 # Continue with the original message if callback fails
 
+        # Re-normalized after before_send, which may have replaced or removed
+        # msg["uuid"], so the returned uuid always matches the wire event.
+        self._normalize_event_uuid(msg)
+        sent_uuid = msg["uuid"]
+
         self.log.debug("queueing: %s", msg)
 
         # if send is False, return msg as if it was successfully queued, unless

posthog/test/test_ai_capture_lane.py

@@ -558,6 +558,36 @@ def test_supplied_uuid_is_preserved_end_to_end(self):
         batch = mock_post.call_args.kwargs["batch"]
         self.assertEqual(batch[0]["uuid"], supplied_uuid)
 
+    def test_returned_uuid_reflects_before_send_replacement(self):
+        replacement_uuid = str(uuid.uuid4())
+
+        def replace_uuid(event):
+            event["uuid"] = replacement_uuid
+            return event
+
+        client = self._client(before_send=replace_uuid)
+        with mock.patch("posthog.consumer.batch_post") as mock_post:
+            returned_uuid = client.capture_ai("$ai_generation", distinct_id="d")
+            client.flush()
+
+        self.assertEqual(returned_uuid, replacement_uuid)
+        batch = mock_post.call_args.kwargs["batch"]
+        self.assertEqual(batch[0]["uuid"], replacement_uuid)
+
+    def test_returned_uuid_is_regenerated_when_before_send_removes_it(self):
+        def drop_uuid(event):
+            del event["uuid"]
+            return event
+
+        client = self._client(before_send=drop_uuid)
+        with mock.patch("posthog.consumer.batch_post") as mock_post:
+            returned_uuid = client.capture_ai("$ai_generation", distinct_id="d")
+            client.flush()
+
+        self.assertIsNotNone(returned_uuid)
+        batch = mock_post.call_args.kwargs["batch"]
+        self.assertEqual(batch[0]["uuid"], returned_uuid)
+
 
 class TestCaptureAiPrivacyMode(unittest.TestCase):
     """Privacy mode always wins over `enable_full_ai_capture`."""

posthog/test/test_before_send.py

@@ -67,6 +67,53 @@ def my_before_send(event):
             self.assertEqual(len(processed_events), 1)
             self.assertEqual(processed_events[0]["event"], "test_event")
 
+    def test_before_send_callback_replacing_uuid_changes_the_returned_uuid(self):
+        """capture()'s return value must match the uuid on the wire event."""
+        replacement_uuid = "12345678-1234-5678-1234-567812345678"
+
+        def replace_uuid(event):
+            event["uuid"] = replacement_uuid
+            return event
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                before_send=replace_uuid,
+                sync_mode=True,
+            )
+            msg_uuid = client.capture("test_event", distinct_id="user1")
+
+            self.assertEqual(msg_uuid, replacement_uuid)
+
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            enqueued_msg = batch_data[0]
+            self.assertEqual(enqueued_msg["uuid"], replacement_uuid)
+
+    def test_before_send_callback_removing_uuid_regenerates_it(self):
+        """If before_send drops the uuid, a fresh one is generated and returned."""
+
+        def remove_uuid(event):
+            del event["uuid"]
+            return event
+
+        with mock.patch("posthog.client.batch_post") as mock_post:
+            client = Client(
+                FAKE_TEST_API_KEY,
+                on_error=self.set_fail,
+                before_send=remove_uuid,
+                sync_mode=True,
+            )
+            msg_uuid = client.capture("test_event", distinct_id="user1")
+
+            self.assertIsNotNone(msg_uuid)
+
+            mock_post.assert_called_once()
+            batch_data = mock_post.call_args[1]["batch"]
+            enqueued_msg = batch_data[0]
+            self.assertEqual(enqueued_msg["uuid"], msg_uuid)
+
     def test_before_send_callback_drops_event(self):
         """Test that before_send callback can drop events by returning None."""