chore: import upstream snapshot with attribution

sdks/python/tests/unit/message_processing/batching/test_batch_manager.py

@@ -0,0 +1,116 @@
+from unittest import mock
+import time
+
+from opik.message_processing import messages
+from opik.message_processing.batching import batch_manager
+from opik.message_processing.batching import batchers
+
+from ....testlib import fake_message_factory
+
+NOT_USED = None
+
+
+def test_batch_manager__messages_processing_methods():
+    integers_batcher = mock.Mock()
+    strings_batcher = mock.Mock()
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+
+    assert tested.message_supports_batching("a-string")
+    assert tested.message_supports_batching(42)
+    assert not tested.message_supports_batching(float(42.5))
+
+    tested.process_message(42)
+    integers_batcher.add.assert_called_once_with(42)
+
+    tested.process_message("a-string")
+    strings_batcher.add.assert_called_once_with("a-string")
+
+
+def test_batch_manager__all_batchers_are_empty__batch_manager_is_empty():
+    integers_batcher = mock.Mock()
+    integers_batcher.is_empty.return_value = True
+    strings_batcher = mock.Mock()
+    strings_batcher.is_empty.return_value = True
+
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+
+    assert tested.is_empty()
+    strings_batcher.is_empty.assert_called_once()
+    integers_batcher.is_empty.assert_called_once()
+
+
+def test_batch_manager__at_least_one_batcher_is_not_empty__batch_manager_is_not_empty():
+    integers_batcher = mock.Mock()
+    integers_batcher.is_empty.return_value = True
+    strings_batcher = mock.Mock()
+    strings_batcher.is_empty.return_value = False
+
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+
+    assert not tested.is_empty()
+    strings_batcher.is_empty.assert_called_once()
+    integers_batcher.is_empty.assert_called_once()
+
+
+def test_batch_manager__flush_is_called__all_batchers_are_flushed():
+    integers_batcher = mock.Mock()
+    strings_batcher = mock.Mock()
+
+    MESSAGE_BATCHERS = {
+        int: integers_batcher,
+        str: strings_batcher,
+    }
+
+    tested = batch_manager.BatchManager(MESSAGE_BATCHERS)
+    tested.flush()
+    integers_batcher.flush.assert_called_once()
+    strings_batcher.flush.assert_called_once()
+
+
+def test_batch_manager__flush_then_stop__accumulated_data_is_flushed_and_thread_is_stopped():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    span_messages_batch = fake_message_factory.fake_span_create_message_batch(
+        count=10, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    example_span_batcher = batchers.CreateSpanMessageBatcher(
+        flush_callback=flush_callback,
+        max_batch_size=42,
+        flush_interval_seconds=0.1,
+        batch_memory_limit_mb=5,
+    )
+    tested = batch_manager.BatchManager(
+        {messages.CreateSpanMessage: example_span_batcher}
+    )
+
+    tested.start()
+    time.sleep(0.1)
+
+    for span_message in span_messages_batch:
+        tested.process_message(span_message)
+
+    # BatchManager.stop() is a pure signal now — callers must flush first if
+    # they want pending batches pushed downstream.
+    tested.flush()
+    tested.stop()
+
+    assert len(collected_messages) >= 2

sdks/python/tests/unit/message_processing/batching/test_experiment_items_batcher.py

@@ -0,0 +1,194 @@
+from unittest import mock
+import time
+
+
+from opik.message_processing.batching import batchers
+from opik.message_processing import messages
+
+NOT_USED = None
+
+
+def test_create_experiment_items_batch_message_batcher__exactly_max_batch_size_reached__batch_is_flushed():
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    experiment_item_batch_messages = [
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_1",
+                    experiment_id="exp1",
+                    trace_id="trace1",
+                    dataset_item_id="item1",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_2",
+                    experiment_id="exp1",
+                    trace_id="trace2",
+                    dataset_item_id="item2",
+                ),
+            ]
+        ),
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_3",
+                    experiment_id="exp1",
+                    trace_id="trace3",
+                    dataset_item_id="item3",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_4",
+                    experiment_id="exp1",
+                    trace_id="trace4",
+                    dataset_item_id="item4",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_5",
+                    experiment_id="exp1",
+                    trace_id="trace5",
+                    dataset_item_id="item5",
+                ),
+            ]
+        ),
+    ]
+
+    for experiment_items_batch in experiment_item_batch_messages:
+        batcher.add(experiment_items_batch)
+    assert batcher.is_empty()
+
+    # Verify that flush callback was called with all items combined
+    flush_callback.assert_called_once()
+    call_args = flush_callback.call_args[0][0]
+    assert isinstance(call_args, messages.CreateExperimentItemsBatchMessage)
+    assert len(call_args.batch) == 5
+    assert call_args.supports_batching is False
+
+
+def test_create_experiment_items_batch_message_batcher__more_than_max_batch_size__multiple_flushes():
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 3
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    # Add 2 messages
+    batcher.add(
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_1",
+                    experiment_id="exp1",
+                    trace_id="trace1",
+                    dataset_item_id="item1",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_2",
+                    experiment_id="exp1",
+                    trace_id="trace2",
+                    dataset_item_id="item2",
+                ),
+            ]
+        )
+    )
+    assert flush_callback.call_count == 0
+
+    # Add 2 more messages (total 4, exceeds max of 3)
+    batcher.add(
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_3",
+                    experiment_id="exp1",
+                    trace_id="trace3",
+                    dataset_item_id="item3",
+                ),
+                messages.ExperimentItemMessage(
+                    id="exp_item_4",
+                    experiment_id="exp1",
+                    trace_id="trace4",
+                    dataset_item_id="item4",
+                ),
+            ]
+        )
+    )
+
+    # Should have flushed the first batch of 3
+    assert flush_callback.call_count == 1
+    first_call_args = flush_callback.call_args_list[0][0][0]
+    assert len(first_call_args.batch) == 3
+
+    # The batcher should have 1 item left
+    assert not batcher.is_empty()
+    assert batcher.size() == 1
+
+
+def test_create_experiment_items_batch_message_batcher__flush_interval_reached__batch_is_flushed():
+    flush_callback = mock.Mock()
+
+    FLUSH_INTERVAL = 0.1  # 100ms
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=1000,  # High enough to not trigger size-based flush
+        flush_callback=flush_callback,
+        flush_interval_seconds=FLUSH_INTERVAL,
+    )
+
+    # Add a message
+    batcher.add(
+        messages.CreateExperimentItemsBatchMessage(
+            batch=[
+                messages.ExperimentItemMessage(
+                    id="exp_item_1",
+                    experiment_id="exp1",
+                    trace_id="trace1",
+                    dataset_item_id="item1",
+                ),
+            ]
+        )
+    )
+
+    # Batcher should not be empty yet
+    assert not batcher.is_empty()
+    assert not batcher.is_ready_to_flush()
+
+    # Wait for flush interval
+    time.sleep(FLUSH_INTERVAL + 0.05)
+
+    # Now it should be ready to flush
+    assert batcher.is_ready_to_flush()
+
+    # Manual flush
+    batcher.flush()
+
+    # Should have flushed
+    assert flush_callback.call_count == 1
+    assert batcher.is_empty()
+
+
+def test_create_experiment_items_batch_message_batcher__empty_batch__nothing_flushed():
+    flush_callback = mock.Mock()
+
+    batcher = batchers.CreateExperimentItemsBatchMessageBatcher(
+        max_batch_size=5,
+        flush_callback=flush_callback,
+        flush_interval_seconds=1.0,
+    )
+
+    # Flush empty batcher
+    batcher.flush()
+
+    # Should not have called the callback
+    flush_callback.assert_not_called()

sdks/python/tests/unit/message_processing/batching/test_feedback_scores_batcher.py

@@ -0,0 +1,165 @@
+from unittest import mock
+import time
+
+import pytest
+
+from opik.message_processing.batching import batchers
+from opik.message_processing import messages
+
+NOT_USED = None
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class, batch_message_class",
+    [
+        (
+            batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+            messages.AddSpanFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+            messages.AddTraceFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+            messages.AddThreadsFeedbackScoresBatchMessage,
+        ),
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__exactly_max_batch_size_reached__batch_is_flushed(
+    message_batcher_class,
+    batch_message_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    add_feedback_score_batch_messages = [
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[1, 2]),
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[3, 4, 5]),
+    ]  # batcher doesn't care about the content
+
+    for feedback_scores_batch in add_feedback_score_batch_messages:
+        batcher.add(feedback_scores_batch)
+    assert batcher.is_empty()
+
+    flush_callback.assert_called_once_with(
+        batch_message_class(batch=[1, 2, 3, 4, 5], supports_batching=False)
+    )
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class,batch_message_class",
+    [
+        (
+            batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+            messages.AddSpanFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+            messages.AddTraceFeedbackScoresBatchMessage,
+        ),
+        (
+            batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+            messages.AddThreadsFeedbackScoresBatchMessage,
+        ),
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__more_than_max_batch_size_items_added__one_batch_flushed__some_data_remains_in_batcher(
+    message_batcher_class,
+    batch_message_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    add_feedback_score_batch_messages = [
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[1, 2]),
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[3, 4, 5, 6]),
+        messages.AddSpanFeedbackScoresBatchMessage(batch=[7, 8]),
+    ]  # batcher doesn't care about the content
+
+    for feedback_scores_batch in add_feedback_score_batch_messages:
+        batcher.add(feedback_scores_batch)
+
+    assert not batcher.is_empty()
+    flush_callback.assert_called_once_with(
+        batch_message_class(batch=[1, 2, 3, 4, 5], supports_batching=False)
+    )
+    flush_callback.reset_mock()
+
+    batcher.flush()
+    flush_callback.assert_called_once_with(
+        batch_message_class(batch=[6, 7, 8], supports_batching=False)
+    )
+    assert batcher.is_empty()
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+        batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+        batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__batcher_doesnt_have_items__flush_is_called__flush_callback_NOT_called(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+    batcher.flush()
+    flush_callback.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.AddSpanFeedbackScoresBatchMessageBatcher,
+        batchers.AddTraceFeedbackScoresBatchMessageBatcher,
+        batchers.AddThreadsFeedbackScoresBatchMessageBatcher,
+    ],
+)
+def test_add_feedback_scores_batch_message_batcher__ready_to_flush_returns_True__is_flush_interval_passed(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+    FLUSH_INTERVAL = 0.1
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=FLUSH_INTERVAL,
+    )
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()
+    batcher.flush()
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()

sdks/python/tests/unit/message_processing/batching/test_flushing_thread.py

@@ -0,0 +1,174 @@
+import threading
+import time
+from unittest import mock
+from opik.message_processing import messages
+from opik.message_processing.batching import (
+    batch_manager,
+    batchers,
+    flushing_thread,
+)
+from ....testlib import fake_message_factory
+
+
+def test_flushing_thread__batcher_is_flushed__every_time_flush_interval_time_passes():
+    flush_callback = mock.Mock()
+    FLUSH_INTERVAL = 0.2
+    very_big_batch_size = float("inf")
+    batcher = batchers.CreateSpanMessageBatcher(
+        flush_callback=flush_callback,
+        max_batch_size=very_big_batch_size,
+        flush_interval_seconds=FLUSH_INTERVAL,
+    )
+    manager = batch_manager.BatchManager(
+        message_to_batcher_mapping={messages.CreateSpanMessage: batcher},
+    )
+
+    spans_messages = fake_message_factory.fake_span_create_message_batch(
+        count=2, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    manager.start()
+    try:
+        manager.process_message(spans_messages[0])
+        flush_callback.assert_not_called()
+
+        time.sleep(FLUSH_INTERVAL + 0.1)
+        # flush interval has passed after batcher was created, batcher is ready to be flushed
+        # (0.1 is added because a thread probation interval is 0.1, and it's already made its first check)
+        flush_callback.assert_called_once()
+
+        flush_callback.reset_mock()
+
+        manager.process_message(spans_messages[1])
+        time.sleep(FLUSH_INTERVAL)
+        # flush interval has passed after a previous flush, batcher is ready to be flushed again
+        flush_callback.assert_called_once()
+    finally:
+        manager.stop(flush=False)
+
+
+def test_flushing_thread__concurrent_adds_during_flush__no_messages_dropped():
+    """Regression test for OPIK-6444.
+
+    The FlushingThread used to call ``batcher.flush()`` without holding the
+    BatchManager lock, while ``BatchManager.process_message`` holds it for
+    ``add()``. Messages appended during the post-iteration tail of
+    ``_create_batches_from_accumulated_messages()`` (notably during
+    ``sequence_splitter.split_into_batches()``) were not in the returned
+    batches, but were wiped by ``self._accumulated_messages = []``.
+    """
+    delivered_trace_ids = set()
+    delivered_lock = threading.Lock()
+
+    def flush_callback(batch_message):
+        with delivered_lock:
+            for trace in batch_message.batch:
+                delivered_trace_ids.add(trace.id)
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        flush_callback=flush_callback,
+        max_batch_size=10_000,
+        flush_interval_seconds=0.005,
+    )
+    manager = batch_manager.BatchManager(
+        message_to_batcher_mapping={messages.CreateTraceMessage: batcher},
+    )
+    # Probe aggressively to maximise the race window in the buggy code path.
+    manager._flushing_thread._probe_interval_seconds = 0.001
+    manager.start()
+
+    WRITERS = 4
+    ADDS_PER_WRITER = 100
+    PAYLOAD_SIZE = 50 * fake_message_factory.ONE_KILOBYTE
+    added_trace_ids = []
+    added_lock = threading.Lock()
+
+    def writer():
+        trace_messages = fake_message_factory.fake_create_trace_message_batch(
+            count=ADDS_PER_WRITER,
+            approximate_trace_size=PAYLOAD_SIZE,
+            has_ended=True,
+        )
+        for message in trace_messages:
+            manager.process_message(message)
+            with added_lock:
+                added_trace_ids.append(message.trace_id)
+
+    writer_threads = [threading.Thread(target=writer) for _ in range(WRITERS)]
+    for writer_thread in writer_threads:
+        writer_thread.start()
+    for writer_thread in writer_threads:
+        writer_thread.join()
+
+    manager.stop(flush=True)
+
+    assert set(added_trace_ids) == delivered_trace_ids, (
+        f"dropped {len(set(added_trace_ids) - delivered_trace_ids)} of "
+        f"{len(added_trace_ids)} trace messages"
+    )
+
+
+def test_flushing_thread__exception_in_flush_callable__thread_keeps_running():
+    """Regression test for OPIK-6444.
+
+    ``FlushingThread.run`` had no try/except, so any exception in the
+    periodic callable silently killed the daemon thread, after which
+    periodic flushing stopped.
+    """
+    call_counter = {"n": 0}
+
+    def flaky_callable():
+        call_counter["n"] += 1
+        if call_counter["n"] == 1:
+            raise RuntimeError("simulated transient failure")
+
+    tested = flushing_thread.FlushingThread(
+        flush_callable=flaky_callable, probe_interval_seconds=0.01
+    )
+    tested.start()
+    try:
+        time.sleep(0.2)
+    finally:
+        tested.close()
+        tested.join(timeout=2)
+
+    assert call_counter["n"] >= 2, "FlushingThread stopped after the first failure"
+
+
+def test_batch_manager_flush_ready__failing_batcher__remaining_batchers_still_flushed():
+    """A failing flush on one batcher must not skip the others in the same tick."""
+    good_callback = mock.Mock()
+    bad_callback = mock.Mock(side_effect=RuntimeError("boom"))
+
+    bad_batcher = batchers.CreateSpanMessageBatcher(
+        flush_callback=bad_callback,
+        max_batch_size=float("inf"),
+        flush_interval_seconds=0.0,
+    )
+    good_batcher = batchers.CreateTraceMessageBatcher(
+        flush_callback=good_callback,
+        max_batch_size=float("inf"),
+        flush_interval_seconds=0.0,
+    )
+    manager = batch_manager.BatchManager(
+        message_to_batcher_mapping={
+            messages.CreateSpanMessage: bad_batcher,
+            messages.CreateTraceMessage: good_batcher,
+        },
+    )
+
+    span = fake_message_factory.fake_span_create_message_batch(
+        count=1, approximate_span_size=fake_message_factory.ONE_KILOBYTE
+    )[0]
+    trace = fake_message_factory.fake_create_trace_message_batch(
+        count=1,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=True,
+    )[0]
+    manager.process_message(span)
+    manager.process_message(trace)
+
+    manager.flush_ready()
+
+    bad_callback.assert_called_once()
+    good_callback.assert_called_once()

sdks/python/tests/unit/message_processing/batching/test_sequence_splitter.py

@@ -0,0 +1,134 @@
+import json
+
+import pytest
+
+from opik.message_processing.batching import sequence_splitter
+
+from ....testlib import fake_message_factory
+
+
+ONE_MEGABYTE_OBJECT_A = fake_message_factory.LongStr(
+    "a" * fake_message_factory.ONE_MEGABYTE
+)
+ONE_MEGABYTE_OBJECT_B = fake_message_factory.LongStr(
+    "b" * fake_message_factory.ONE_MEGABYTE
+)
+ONE_MEGABYTE_OBJECT_C = fake_message_factory.LongStr(
+    "c" * fake_message_factory.ONE_MEGABYTE
+)
+
+
+def test_split_list_into_batches__by_size_only():
+    items = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    batches = sequence_splitter.split_into_batches(items, max_length=4)
+
+    assert batches == [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]]
+
+
+def test_split_list_into_batches__by_memory_only():
+    items = [ONE_MEGABYTE_OBJECT_A] * 2 + [ONE_MEGABYTE_OBJECT_B] * 2
+    batches = sequence_splitter.split_into_batches(items, max_payload_size_MB=3.5)
+
+    assert batches == [
+        [ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_B],
+        [ONE_MEGABYTE_OBJECT_B],
+    ]
+
+
+def test_split_list_into_batches__by_memory_and_by_size():
+    FOUR_MEGABYTE_OBJECT_C = [ONE_MEGABYTE_OBJECT_C] * 4
+    items = (
+        [ONE_MEGABYTE_OBJECT_A] * 2
+        + [FOUR_MEGABYTE_OBJECT_C]
+        + [ONE_MEGABYTE_OBJECT_B] * 2
+    )
+    batches = sequence_splitter.split_into_batches(
+        items, max_length=2, max_payload_size_MB=3.5
+    )
+
+    # Object C comes before object A because if item is bigger than the max payload size
+    # it is immediately added to the result batches list before batch which is currently accumulating
+    assert batches == [
+        [FOUR_MEGABYTE_OBJECT_C],
+        [ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_A],
+        [ONE_MEGABYTE_OBJECT_B, ONE_MEGABYTE_OBJECT_B],
+    ]
+
+
+def test_split_list_into_batches__empty_list():
+    items = []
+    batches = sequence_splitter.split_into_batches(
+        items, max_length=3, max_payload_size_MB=3.5
+    )
+
+    assert batches == []
+
+
+def test_split_list_into_batches__multiple_large_objects():
+    items = [ONE_MEGABYTE_OBJECT_A, ONE_MEGABYTE_OBJECT_B, ONE_MEGABYTE_OBJECT_C]
+    batches = sequence_splitter.split_into_batches(
+        items, max_length=2, max_payload_size_MB=0.5
+    )
+
+    assert batches == [
+        [ONE_MEGABYTE_OBJECT_A],
+        [ONE_MEGABYTE_OBJECT_B],
+        [ONE_MEGABYTE_OBJECT_C],
+    ]
+
+
+@pytest.mark.parametrize(
+    "test_input",
+    [
+        ("string"),  # Basic string
+        (""),  # Empty string
+        ("Unicode👋🏼"),  # String with non-UTF8 characters
+        (123.76),  # Float
+        (123456789),  # Integer
+        (True),  # Boolean
+        (False),  # Boolean
+        ([1, 2, 3]),  # List
+        ({"key": "value"}),  # Dictionary
+        (None),  # null
+    ],
+)
+def test_get_json_size_basic(test_input):
+    expected = len(
+        json.dumps(test_input, separators=(",", ":"), ensure_ascii=False).encode(
+            "utf-8"
+        )
+    )
+    assert sequence_splitter._get_json_size(test_input) == expected
+
+
+def test_get_json_size_complex_nested():
+    test_input = {
+        "string": "simple value",
+        "number": 42,
+        "boolean": True,
+        "null": None,
+        "nested_dict": {
+            "level1": {
+                "level2": {
+                    "level3": "deep value",
+                    "numbers": [1, 2, 3],
+                    "mixed_list": [True, "string", 42.0, None],
+                },
+                "sibling2": {
+                    "data": [{"id": 1, "name": "item1"}, {"id": 2, "name": "item2"}]
+                },
+            }
+        },
+        "array_of_arrays": [[1, 2, 3], ["a", "b", "c"], [{"x": 1}, {"y": 2}]],
+        "mixed_nesting": {
+            "lists": [[1, 2], [3, 4]],
+            "dicts": {"a": {"b": {"c": "value"}}, "x": {"y": {"z": 100}}},
+        },
+    }
+
+    expected = len(
+        json.dumps(test_input, separators=(",", ":"), ensure_ascii=False).encode(
+            "utf-8"
+        )
+    )
+    assert sequence_splitter._get_json_size(test_input) == expected

sdks/python/tests/unit/message_processing/batching/test_span_batcher.py

@@ -0,0 +1,236 @@
+from unittest import mock
+import time
+
+import pytest
+
+from opik.message_processing.batching import batchers
+from opik.message_processing import messages
+from ....testlib import fake_message_factory
+
+NOT_USED = mock.sentinel.NOT_USED
+
+
+def test_create_span_message_batcher__exactly_max_batch_size_reached__batch_is_flushed():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=10,
+    )
+
+    assert batcher.is_empty()
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=MAX_BATCH_SIZE, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    for span_message in span_messages:
+        batcher.add(span_message)
+    assert batcher.is_empty()
+
+    assert len(collected_messages) == 1
+
+
+def test_create_span_message_batcher__more_than_max_batch_size_items_added__one_batch_flushed__some_data_remains_in_batcher():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=10,
+    )
+
+    assert batcher.is_empty()
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=MAX_BATCH_SIZE + 2,
+        approximate_span_size=fake_message_factory.ONE_MEGABYTE,
+    )
+
+    for span_message in span_messages:
+        batcher.add(span_message)
+
+    assert not batcher.is_empty()
+    assert len(collected_messages) == 1
+
+    batcher.flush()
+    assert len(collected_messages) == 2
+
+
+def test_create_span_message_batcher__split_message_into_batches__size_limit_reached():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=3,
+    )
+
+    assert batcher.is_empty()
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=2 * 2, approximate_span_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    for span_message in span_messages:
+        batcher.add(span_message)
+
+    assert not batcher.is_empty()
+    batcher.flush()
+
+    # we created 4 messages about 1 MB each, and batcher is limited to 3 MB - thus we expect 2 batches
+    assert len(collected_messages) == 2
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.CreateSpanMessageBatcher,
+        batchers.CreateTraceMessageBatcher,
+    ],
+)
+def test_create_message_batcher__batcher_doesnt_have_items__flush_is_called__flush_callback_NOT_called(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=10,
+    )
+
+    assert batcher.is_empty()
+    batcher.flush()
+    flush_callback.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "message_batcher_class",
+    [
+        batchers.CreateSpanMessageBatcher,
+        batchers.CreateTraceMessageBatcher,
+    ],
+)
+def test_create_message_batcher__ready_to_flush_returns_True__is_flush_interval_passed(
+    message_batcher_class,
+):
+    flush_callback = mock.Mock()
+
+    MAX_BATCH_SIZE = 5
+    FLUSH_INTERVAL = 0.1
+
+    batcher = message_batcher_class(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=FLUSH_INTERVAL,
+        batch_memory_limit_mb=10,
+    )
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()
+    batcher.flush()
+    assert not batcher.is_ready_to_flush()
+    time.sleep(0.1)
+    assert batcher.is_ready_to_flush()
+
+
+def test_create_span_message_batcher__add_duplicated_span__previous_span_is_removed_from_batcher():
+    """
+    Tests the handling of duplicate spans in a span message batcher, specifically focusing on ensuring
+    that newer ones replace older duplicated spans.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+    assert batcher.is_empty()
+
+    span_messages = []
+    span_messages += fake_message_factory.fake_span_create_message_batch(
+        count=1,
+        approximate_span_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+    span_messages += fake_message_factory.fake_span_create_message_batch(
+        count=1, approximate_span_size=fake_message_factory.ONE_KILOBYTE, has_ended=True
+    )
+
+    batcher.add(span_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second span to be considered the same as the first one
+    span_messages[1].span_id = span_messages[0].span_id
+    batcher.add(span_messages[1])
+
+    # assert that the first span is removed from the batcher
+    assert batcher.size() == 1
+
+    # check that the second message is in the batch
+    batcher.flush()
+    assert len(batches) == 1
+    assert len(batches[0].batch) == 1
+    assert batches[0].batch[0].id == span_messages[1].span_id
+
+
+def test_create_span_message_batcher__add_duplicated_span__previous_span_is_not_removed_from_batcher__not_ended_span():
+    """
+    Tests the behavior of adding a duplicate span to a span message batcher and ensures
+    that the previous span is not removed when adding a duplicate span if both are not ended.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateSpanMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+    assert batcher.is_empty()
+
+    span_messages = fake_message_factory.fake_span_create_message_batch(
+        count=2,
+        approximate_span_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+
+    batcher.add(span_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second span to be considered the same as the first one
+    span_messages[1].span_id = span_messages[0].span_id
+    batcher.add(span_messages[1])
+
+    # assert that both spans are in the batcher
+    assert batcher.size() == 2

sdks/python/tests/unit/message_processing/batching/test_trace_batcher.py

@@ -0,0 +1,123 @@
+from unittest import mock
+
+from opik.message_processing import messages
+from opik.message_processing.batching import batchers
+
+from ....testlib import fake_message_factory
+
+NOT_USED = mock.sentinel.NOT_USED
+
+
+def test_create_trace_message_batcher__split_message_into_batches__size_limit_reached():
+    collected_messages = []
+
+    def flush_callback(message: messages.BaseMessage):
+        collected_messages.append(message)
+
+    MAX_BATCH_SIZE = 5
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+        batch_memory_limit_mb=3,
+    )
+
+    assert batcher.is_empty()
+    trace_messages = fake_message_factory.fake_create_trace_message_batch(
+        count=2 * 2, approximate_trace_size=fake_message_factory.ONE_MEGABYTE
+    )
+
+    for trace_message in trace_messages:
+        batcher.add(trace_message)
+
+    assert not batcher.is_empty()
+    batcher.flush()
+
+    # we created 4 messages about 1 MB each, and batcher is limited to 3 MB - thus we expect 2 batches
+    assert len(collected_messages) == 2
+
+
+def test_create_trace_message_batcher__add_duplicated_trace__previous_trace_is_removed_from_batcher():
+    """
+    Tests the handling of duplicate traces in a trace message batcher, specifically focusing on ensuring
+    that newer ones replace older duplicated traces.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+
+    trace_messages = []
+    trace_messages += fake_message_factory.fake_create_trace_message_batch(
+        count=1,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+    trace_messages += fake_message_factory.fake_create_trace_message_batch(
+        count=1,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=True,
+    )
+
+    batcher.add(trace_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second trace to be considered the same as the first one
+    trace_messages[1].trace_id = trace_messages[0].trace_id
+
+    # add the second trace to the batcher and check that the first one is removed from the batcher
+    batcher.add(trace_messages[1])
+    assert batcher.size() == 1
+
+    batcher.flush()
+    assert len(batches) == 1
+    assert len(batches[0].batch) == 1
+    assert batches[0].batch[0].id == trace_messages[1].trace_id
+
+
+def test_create_trace_message_batcher__add_duplicated_trace__previous_trace_is_not_removed_from_batcher__not_ended_trace():
+    """
+    Tests the behavior of adding a duplicate trace to a trace message batcher and ensures
+    that the previous trace is not removed when adding a duplicate trace if both are not ended.
+    """
+    MAX_BATCH_SIZE = 5
+
+    batches = []
+
+    def flush_callback(batch):
+        batches.append(batch)
+
+    batcher = batchers.CreateTraceMessageBatcher(
+        max_batch_size=MAX_BATCH_SIZE,
+        flush_callback=flush_callback,
+        flush_interval_seconds=NOT_USED,
+    )
+
+    assert batcher.is_empty()
+
+    trace_messages = fake_message_factory.fake_create_trace_message_batch(
+        count=2,
+        approximate_trace_size=fake_message_factory.ONE_KILOBYTE,
+        has_ended=False,
+    )
+
+    batcher.add(trace_messages[0])
+    assert batcher.size() == 1
+
+    # modify the second trace to be considered the same as the first one
+    trace_messages[1].trace_id = trace_messages[0].trace_id
+
+    # add the second trace to the batcher and check that the first one is removed from the batcher
+    batcher.add(trace_messages[1])
+    assert batcher.size() == 2

sdks/python/tests/unit/message_processing/emulation/test_emulator_activation.py

@@ -0,0 +1,115 @@
+"""Ref-counted activation + single recording-context slot on the emulator.
+
+These guard the behavior that lets several concurrent users of a *shared*
+processing chain (e.g. evaluate() runs sharing one connection) coordinate
+instead of deactivating the emulator out from under each other.
+"""
+
+import datetime
+
+from opik.message_processing import messages
+from opik.message_processing.emulation import local_emulator_message_processor
+
+
+def _emulator():
+    return local_emulator_message_processor.LocalEmulatorMessageProcessor(active=False)
+
+
+def _trace_message(trace_id: str = "t1"):
+    now = datetime.datetime(2026, 1, 1, tzinfo=datetime.timezone.utc)
+    return messages.CreateTraceMessage(
+        trace_id=trace_id,
+        project_name="p",
+        name="trace",
+        start_time=now,
+        end_time=now,
+        input={},
+        output={},
+        metadata={},
+        tags=[],
+        error_info=None,
+        thread_id=None,
+        last_updated_at=now,
+        source="sdk",
+    )
+
+
+def test_acquire_release__concurrent_holders__active_until_last_release():
+    emulator = _emulator()
+    assert not emulator.is_active()
+
+    emulator.acquire()
+    emulator.acquire()
+    assert emulator.is_active()
+
+    emulator.release()
+    assert emulator.is_active()  # one holder remains
+
+    emulator.release()
+    assert not emulator.is_active()
+
+
+def test_release__when_not_acquired__is_noop():
+    emulator = _emulator()
+    emulator.release()
+    assert not emulator.is_active()
+
+
+def test_release__last_holder_with_reset__clears_recorded_state():
+    emulator = _emulator()
+    emulator.acquire(reset=True)
+    emulator.process(_trace_message())
+    assert len(emulator.trace_trees) == 1
+
+    emulator.release(reset=True)
+
+    assert emulator.trace_trees == []
+
+
+def test_release__reset_deferred_until_last_holder__cleared_only_at_zero():
+    emulator = _emulator()
+    emulator.acquire(reset=True)
+    emulator.process(_trace_message())
+    assert len(emulator.trace_trees) == 1
+
+    # A second holder acquires and releases; reset must not fire while a holder
+    # remains, so the recorded trace stays intact.
+    emulator.acquire(reset=True)
+    emulator.release(reset=True)
+    assert len(emulator.trace_trees) == 1
+
+    # Only the final release (refcount -> 0) resets the recorded state.
+    emulator.release(reset=True)
+    assert emulator.trace_trees == []
+
+
+def test_acquire__first_holder_with_reset__clears_prior_state():
+    emulator = _emulator()
+    emulator.acquire(reset=True)
+    emulator.process(_trace_message())
+    emulator.release(reset=False)  # leave the recorded trace in place
+    assert len(emulator.trace_trees) == 1
+
+    emulator.acquire(reset=True)  # fresh activation must start clean
+
+    assert emulator.trace_trees == []
+
+
+def test_begin_recording_context__nested__second_call_refused():
+    emulator = _emulator()
+    assert emulator.begin_recording_context() is True
+    assert emulator.begin_recording_context() is False
+
+    emulator.end_recording_context()
+    assert emulator.begin_recording_context() is True
+
+
+def test_begin_recording_context__while_activated_for_evaluation__still_allowed():
+    emulator = _emulator()
+    emulator.acquire()  # simulates evaluate() activating the shared emulator
+
+    # A running evaluation must not block record_traces_locally().
+    assert emulator.begin_recording_context() is True
+
+    emulator.end_recording_context()
+    emulator.release()

sdks/python/tests/unit/message_processing/permissions/test_unauthorized_message_type_registry.py

@@ -0,0 +1,182 @@
+from unittest import mock
+
+from opik.message_processing.permissions import unauthorized_message_type_registry
+
+MESSAGE_TYPE_A = "message_type_a"
+MESSAGE_TYPE_B = "message_type_b"
+RETRY_INTERVAL = 60.0
+NOW = 1000.0
+
+
+class TestIsAuthorized:
+    def test_is_authorized__message_type_not_in_registry__returns_true(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        assert registry.is_authorized(MESSAGE_TYPE_A) is True
+
+    def test_is_authorized__message_added_within_retry_interval__returns_false(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL * 0.5):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is False
+
+    def test_is_authorized__retry_interval_elapsed__returns_true(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL + 1.0):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is True
+
+    def test_is_authorized__elapsed_time_equals_retry_interval__returns_true(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is True
+
+    def test_is_authorized__different_message_type_not_in_registry__returns_true(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL * 0.5):
+            assert registry.is_authorized(MESSAGE_TYPE_B) is True
+
+    def test_is_authorized__message_re_added_after_interval_elapsed__returns_false(
+        self,
+    ):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        later = NOW + RETRY_INTERVAL + 1.0
+        registry.add(MESSAGE_TYPE_A, attempt_time=later)
+
+        with mock.patch("time.time", return_value=later + RETRY_INTERVAL * 0.5):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is False
+
+
+class TestIsAuthorizedWithMaxRetryCount:
+    def test_is_authorized__max_retry_count_reached__returns_false(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL, max_retry_count=2
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)  # retry_count=0
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)  # retry_count=1
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)  # retry_count=2
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL + 1.0):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is False
+
+    def test_is_authorized__retry_count_below_max_and_interval_elapsed__returns_true(
+        self,
+    ):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL, max_retry_count=3
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)  # retry_count=0
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)  # retry_count=1
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL + 1.0):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is True
+
+    def test_is_authorized__retry_count_below_max_within_interval__returns_false(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL, max_retry_count=3
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)  # retry_count=0
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL * 0.5):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is False
+
+    def test_is_authorized__max_retry_count_not_set_and_interval_elapsed__returns_true(
+        self,
+    ):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL, max_retry_count=None
+        )
+        for _ in range(100):
+            registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        with mock.patch("time.time", return_value=NOW + RETRY_INTERVAL + 1.0):
+            assert registry.is_authorized(MESSAGE_TYPE_A) is True
+
+
+class TestAdd:
+    def test_add__message_type_not_in_registry__retry_count_is_zero(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        assert registry.registry[MESSAGE_TYPE_A].retry_count == 0
+
+    def test_add__message_type_not_in_registry__last_attempt_at_is_stored(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        assert registry.registry[MESSAGE_TYPE_A].last_attempt_at == NOW
+
+    def test_add__message_type_already_in_registry__retry_count_incremented(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        assert registry.registry[MESSAGE_TYPE_A].retry_count == 1
+
+    def test_add__message_type_already_in_registry__last_attempt_at_updated(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        later = NOW + RETRY_INTERVAL + 1.0
+        registry.add(MESSAGE_TYPE_A, attempt_time=later)
+
+        assert registry.registry[MESSAGE_TYPE_A].last_attempt_at == later
+
+    def test_add__attempt_time_not_provided__last_attempt_at_set_to_current_time(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+
+        with mock.patch("time.time", return_value=NOW):
+            registry.add(MESSAGE_TYPE_A)
+
+        assert registry.registry[MESSAGE_TYPE_A].last_attempt_at == NOW
+
+    def test_add__message_type_added_multiple_times__retry_count_incremented_cumulatively(
+        self,
+    ):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        for _ in range(5):
+            registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+
+        assert registry.registry[MESSAGE_TYPE_A].retry_count == 4
+
+    def test_add__different_message_types__tracked_independently(self):
+        registry = unauthorized_message_type_registry.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=RETRY_INTERVAL
+        )
+        registry.add(MESSAGE_TYPE_A, attempt_time=NOW)
+        registry.add(MESSAGE_TYPE_B, attempt_time=NOW)
+        registry.add(MESSAGE_TYPE_B, attempt_time=NOW)
+
+        assert registry.registry[MESSAGE_TYPE_A].retry_count == 0
+        assert registry.registry[MESSAGE_TYPE_B].retry_count == 1

sdks/python/tests/unit/message_processing/preprocessing/test_attachments_preprocessor.py

@@ -0,0 +1,278 @@
+from unittest import mock
+
+from opik.message_processing import messages
+from opik.message_processing.preprocessing import constants, attachments_preprocessor
+
+
+def test_preprocess_message_does_not_wrap_create_span_without_end_time():
+    """Test preprocess_message does not wrap CreateSpanMessage without end_time."""
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"input_key": "input_value"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(span_message)
+
+    # Should return the original message unchanged since end_time is None
+    assert result is span_message
+    assert not isinstance(result, messages.AttachmentSupportingMessage)
+
+
+def test_preprocess_message_does_not_wrap_create_trace_without_end_time():
+    """Test preprocess_message does not wrap CreateTraceMessage without end_time."""
+    trace_message = messages.CreateTraceMessage(
+        trace_id="trace-789",
+        project_name="test-project",
+        name="test-trace",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output={"output_key": "output_value"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(trace_message)
+
+    # Should return the original message unchanged since end_time is None
+    assert result is trace_message
+    assert not isinstance(result, messages.AttachmentSupportingMessage)
+
+
+def test_preprocess_message_does_not_wrap_other_messages():
+    """Test preprocess_message doesn't wrap unsupported message types."""
+    base_message = messages.BaseMessage()
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(base_message)
+
+    # Should return the original message unchanged
+    assert result is base_message
+    assert not isinstance(result, messages.AttachmentSupportingMessage)
+
+
+def test_preprocess_message_avoids_double_wrapping():
+    """Test preprocess_message doesn't re-wrap already marked messages."""
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output={"output_key": "output_value"},
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    # Mark the message as already preprocessed
+    setattr(span_message, constants.MARKER_ATTRIBUTE_NAME, True)
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(span_message)
+
+    # Should return an original message without wrapping
+    assert result is span_message
+    assert not isinstance(result, messages.AttachmentSupportingMessage)
+
+
+def test_preprocess_message_avoids_wrapping_message_with_empty_candidate_fields():
+    """Test preprocess_message doesn't wrap messages with empty candidate fields."""
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(span_message)
+
+    # Should return an original message without wrapping
+    assert result is span_message
+    assert not isinstance(result, messages.AttachmentSupportingMessage)
+
+
+def test_preprocess_message_wraps_create_span_with_end_time():
+    """Test preprocess_message wraps CreateSpanMessage with end_time set."""
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test",
+        start_time=mock.Mock(),
+        end_time=mock.Mock(),  # end_time is set
+        input={"input_key": "input_value"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(span_message)
+
+    assert isinstance(result, messages.AttachmentSupportingMessage)
+    assert result.original_message is span_message
+
+
+def test_preprocess_message_wraps_create_trace_with_end_time():
+    """Test preprocess_message wraps CreateTraceMessage with end_time set."""
+    trace_message = messages.CreateTraceMessage(
+        trace_id="trace-789",
+        project_name="test-project",
+        name="test-trace",
+        start_time=mock.Mock(),
+        end_time=mock.Mock(),  # end_time is set
+        input=None,
+        output={"output_key": "output_value"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(trace_message)
+
+    assert isinstance(result, messages.AttachmentSupportingMessage)
+    assert result.original_message is trace_message
+
+
+def test_preprocess_message_wraps_update_span():
+    """Test preprocess_message wraps UpdateSpanMessage."""
+    span_message = messages.UpdateSpanMessage(
+        span_id="span-123",
+        parent_span_id=None,
+        trace_id="trace-456",
+        project_name="test-project",
+        end_time=None,
+        input={"input_key": "input_value"},
+        output=None,
+        metadata=None,
+        tags=None,
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(span_message)
+
+    assert isinstance(result, messages.AttachmentSupportingMessage)
+    assert result.original_message is span_message
+
+
+def test_preprocess_message_wraps_update_trace():
+    """Test preprocess_message wraps UpdateTraceMessage."""
+    trace_message = messages.UpdateTraceMessage(
+        trace_id="trace-789",
+        project_name="test-project",
+        end_time=None,
+        input=None,
+        output={"output_key": "output_value"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(trace_message)
+
+    assert isinstance(result, messages.AttachmentSupportingMessage)
+    assert result.original_message is trace_message
+
+
+def test_preprocess_message_does_not_wrap_create_span_with_end_time_but_empty_fields():
+    """Test preprocess_message doesn't wrap CreateSpanMessage with end_time if all candidate fields are empty."""
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test",
+        start_time=mock.Mock(),
+        end_time=mock.Mock(),  # end_time is set
+        input=None,  # but all candidate fields are None
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    processor = attachments_preprocessor.AttachmentsPreprocessor()
+    result = processor.preprocess(span_message)
+
+    # Should return the original message unchanged since all candidate fields are None
+    assert result is span_message
+    assert not isinstance(result, messages.AttachmentSupportingMessage)

sdks/python/tests/unit/message_processing/processors/test_assertion_results_processor.py

@@ -0,0 +1,185 @@
+"""
+Unit tests for AssertionResultsMessageProcessor — the migration-period write
+path for OPIK-6054 / OPIK-6048.
+
+Targets the small dedicated class directly so the tests don't depend on the
+larger OpikMessageProcessor / replay / permissions plumbing.
+"""
+
+from unittest import mock
+
+import pytest
+
+from opik.message_processing import messages
+from opik.message_processing.processors.assertion_results_processor import (
+    AssertionResultsMessageProcessor,
+)
+from opik.rest_api import core as rest_api_core
+from opik.rest_api.types import assertion_result_batch_item, feedback_score_batch_item
+
+
+@pytest.fixture
+def mock_rest_client() -> mock.MagicMock:
+    client = mock.MagicMock()
+    # MagicMock blocks attributes starting with "assert" as protection against
+    # typoed assertion methods; configure assertion_results explicitly.
+    client.assertion_results = mock.MagicMock(unsafe=True)
+    return client
+
+
+@pytest.fixture
+def processor(
+    mock_rest_client: mock.MagicMock,
+) -> AssertionResultsMessageProcessor:
+    return AssertionResultsMessageProcessor(rest_client=mock_rest_client)
+
+
+def _build_message(
+    items=None,
+    entity_type="TRACE",
+) -> messages.AddAssertionResultsBatchMessage:
+    items = items or [
+        messages.AssertionResultMessage(
+            entity_id="trace-1",
+            project_name="proj-A",
+            name="assertion 1",
+            status="passed",
+            source="sdk",
+            reason="looked good",
+        ),
+        messages.AssertionResultMessage(
+            entity_id="trace-1",
+            project_name="proj-A",
+            name="assertion 2",
+            status="failed",
+            source="sdk",
+            reason=None,
+        ),
+    ]
+    msg = messages.AddAssertionResultsBatchMessage(batch=items, entity_type=entity_type)
+    msg.message_id = 1
+    return msg
+
+
+class TestAssertionResultsProcessorHappyPath:
+    def test_process__happy_path__forwards_to_assertion_results_endpoint(
+        self,
+        processor: AssertionResultsMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        processor.process(_build_message())
+
+        mock_rest_client.assertion_results.store_assertions_batch.assert_called_once()
+        call_kwargs = (
+            mock_rest_client.assertion_results.store_assertions_batch.call_args.kwargs
+        )
+        assert call_kwargs["entity_type"] == "TRACE"
+        sent_items = call_kwargs["assertion_results"]
+        assert len(sent_items) == 2
+        assert all(
+            isinstance(item, assertion_result_batch_item.AssertionResultBatchItem)
+            for item in sent_items
+        )
+        assert sent_items[0].entity_id == "trace-1"
+        assert sent_items[0].name == "assertion 1"
+        assert sent_items[0].status == "passed"
+        assert sent_items[0].source == "sdk"
+        assert sent_items[0].reason == "looked good"
+        assert sent_items[1].status == "failed"
+        assert sent_items[1].reason is None
+
+    def test_process__happy_path__does_not_use_feedback_scores_endpoint(
+        self,
+        processor: AssertionResultsMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        processor.process(_build_message())
+
+        mock_rest_client.traces.score_batch_of_traces.assert_not_called()
+        mock_rest_client.spans.score_batch_of_spans.assert_not_called()
+
+
+class TestAssertionResultsProcessorFallback:
+    """
+    Forward-compat fallback: if the dedicated assertion-results endpoint is
+    missing on an older self-hosted backend, write through the legacy
+    feedback-scores piggyback and remember the choice for the rest of the
+    session.
+    """
+
+    def test_process__bad_endpoint_404__falls_back_to_feedback_scores(
+        self,
+        processor: AssertionResultsMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        mock_rest_client.assertion_results.store_assertions_batch.side_effect = (
+            rest_api_core.ApiError(status_code=404, body="Not Found")
+        )
+
+        processor.process(_build_message())
+
+        mock_rest_client.traces.score_batch_of_traces.assert_called_once()
+        sent = mock_rest_client.traces.score_batch_of_traces.call_args.kwargs["scores"]
+        assert len(sent) == 2
+        assert all(
+            isinstance(item, feedback_score_batch_item.FeedbackScoreBatchItem)
+            for item in sent
+        )
+        assert all(item.category_name == "suite_assertion" for item in sent)
+        assert sent[0].id == "trace-1"
+        assert sent[0].name == "assertion 1"
+        assert sent[0].value == 1.0
+        assert sent[0].source == "sdk"
+        assert sent[1].value == 0.0  # status="failed"
+        assert processor._use_assertion_results_endpoint is False
+
+    def test_process__bad_endpoint_405__falls_back_to_feedback_scores(
+        self,
+        processor: AssertionResultsMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        mock_rest_client.assertion_results.store_assertions_batch.side_effect = (
+            rest_api_core.ApiError(status_code=405, body="Method Not Allowed")
+        )
+
+        processor.process(_build_message())
+
+        mock_rest_client.traces.score_batch_of_traces.assert_called_once()
+        assert processor._use_assertion_results_endpoint is False
+
+    def test_process__non_404_api_error__propagates_and_does_not_set_fallback(
+        self,
+        processor: AssertionResultsMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        mock_rest_client.assertion_results.store_assertions_batch.side_effect = (
+            rest_api_core.ApiError(status_code=500, body="Internal Server Error")
+        )
+
+        with pytest.raises(rest_api_core.ApiError):
+            processor.process(_build_message())
+
+        mock_rest_client.traces.score_batch_of_traces.assert_not_called()
+        assert processor._use_assertion_results_endpoint is True
+
+    def test_process__sticky_fallback__second_call_skips_new_endpoint(
+        self,
+        processor: AssertionResultsMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        mock_rest_client.assertion_results.store_assertions_batch.side_effect = (
+            rest_api_core.ApiError(status_code=404, body="Not Found")
+        )
+
+        processor.process(_build_message())
+        # First call: tried new endpoint, hit 404, flipped the flag.
+        assert mock_rest_client.assertion_results.store_assertions_batch.call_count == 1
+        assert processor._use_assertion_results_endpoint is False
+
+        mock_rest_client.traces.score_batch_of_traces.reset_mock()
+        processor.process(_build_message())
+
+        # Second call must NOT retry the new endpoint.
+        assert mock_rest_client.assertion_results.store_assertions_batch.call_count == 1
+        # It went straight to the fallback.
+        mock_rest_client.traces.score_batch_of_traces.assert_called_once()

sdks/python/tests/unit/message_processing/processors/test_attachments_extraction_processor.py

@@ -0,0 +1,817 @@
+from unittest import mock
+
+import pytest
+
+from opik.api_objects.attachment import attachment, attachment_context
+from opik.message_processing import messages
+from opik.message_processing.preprocessing import constants
+from opik.message_processing.processors import attachments_extraction_processor
+
+
+@pytest.fixture
+def mock_streamer():
+    """Create a mock streamer."""
+    return mock.Mock()
+
+
+@pytest.fixture
+def processor(mock_streamer):
+    """Create an AttachmentsExtractionProcessor with default settings."""
+    return attachments_extraction_processor.AttachmentsExtractionProcessor(
+        min_attachment_size=20,
+        messages_streamer=mock_streamer,
+        url_override="https://example.com",
+        is_active=True,
+    )
+
+
+@pytest.fixture
+def inactive_processor(mock_streamer):
+    """Create an inactive AttachmentsExtractionProcessor."""
+    return attachments_extraction_processor.AttachmentsExtractionProcessor(
+        min_attachment_size=20,
+        messages_streamer=mock_streamer,
+        url_override="https://example.com",
+        is_active=False,
+    )
+
+
+def test_process_non_attachment_supporting_message_skips_processing(processor):
+    """Test that non-AttachmentSupportingMessage messages are skipped."""
+    # Create a regular message (not wrapped)
+    regular_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"key": "value"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    # Process should return immediately without doing anything
+    processor.process(regular_message)
+
+    # Streamer should not be called
+    processor.messages_streamer.put.assert_not_called()
+
+
+def test_process_attachment_support_message_with_no_attachments(
+    processor, mock_streamer
+):
+    """Test processing a message with no extractable attachments."""
+    # Create an original message with plain text
+    original_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"text": "just plain text"},
+        output={"result": "no attachments here"},
+        metadata={"info": "metadata"},
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    # Wrap in AttachmentSupportingMessage
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Mock extractor to return no attachments
+    with mock.patch.object(processor.extractor, "extract_and_replace", return_value=[]):
+        processor.process(wrapped_message)
+
+    # The original message should be re-queued with marker
+    assert mock_streamer.put.call_count == 1
+    requeued_message = mock_streamer.put.call_args[0][0]
+    assert requeued_message is original_message
+    assert hasattr(requeued_message, constants.MARKER_ATTRIBUTE_NAME)
+    assert getattr(requeued_message, constants.MARKER_ATTRIBUTE_NAME) is True
+
+
+def test_process_span_message_extracts_from_input(
+    processor, mock_streamer, temp_file_15kb
+):
+    """Test extraction of attachments from span input."""
+    # Create a span message with input containing base64
+    original_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"image": "fake-base64-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Mock extracted attachment - use actual temp file
+    mock_attachment = attachment.Attachment(
+        data=temp_file_15kb.name,
+        file_name="input-attachment-123.png",
+        content_type="image/png",
+        create_temp_copy=False,
+    )
+
+    mock_attachment_with_context = attachment_context.AttachmentWithContext(
+        attachment_data=mock_attachment,
+        entity_type="span",
+        entity_id="span-123",
+        project_name="test-project",
+        context="input",
+    )
+
+    # Mock extractor to return one attachment from input
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        if context == "input":
+            return [mock_attachment_with_context]
+        return []
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    # Should have created one CreateAttachmentMessage and re-queued original
+    assert mock_streamer.put.call_count == 2
+
+    # First call should be the attachment message
+    first_call = mock_streamer.put.call_args_list[0][0][0]
+    assert isinstance(first_call, messages.CreateAttachmentMessage)
+    assert first_call.entity_type == "span"
+    assert first_call.entity_id == "span-123"
+    assert first_call.project_name == "test-project"
+    assert first_call.delete_after_upload is True
+
+    # Second call should be the original message with marker
+    second_call = mock_streamer.put.call_args_list[1][0][0]
+    assert second_call is original_message
+    assert hasattr(second_call, constants.MARKER_ATTRIBUTE_NAME)
+
+
+def test_process_span_message_extracts_from_output(
+    processor, mock_streamer, temp_file_15kb
+):
+    """Test extraction of attachments from the span output."""
+    original_message = messages.CreateSpanMessage(
+        span_id="span-789",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output={"image": "fake-base64-data"},
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    mock_attachment = attachment.Attachment(
+        data=temp_file_15kb.name,
+        file_name="output-attachment-456.png",
+        content_type="image/png",
+        create_temp_copy=False,
+    )
+
+    mock_attachment_with_context = attachment_context.AttachmentWithContext(
+        attachment_data=mock_attachment,
+        entity_type="span",
+        entity_id="span-789",
+        project_name="test-project",
+        context="output",
+    )
+
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        if context == "output":
+            return [mock_attachment_with_context]
+        return []
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    assert mock_streamer.put.call_count == 2
+    attachment_msg = mock_streamer.put.call_args_list[0][0][0]
+    assert isinstance(attachment_msg, messages.CreateAttachmentMessage)
+
+
+def test_process_span_message_extracts_from_metadata(
+    processor, mock_streamer, temp_file_15kb
+):
+    """Test extraction of attachments from span metadata."""
+    original_message = messages.CreateSpanMessage(
+        span_id="span-999",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output=None,
+        metadata={"doc": "fake-base64-pdf"},
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    mock_attachment = attachment.Attachment(
+        data=temp_file_15kb.name,
+        file_name="metadata-attachment-789.pdf",
+        content_type="application/pdf",
+        create_temp_copy=False,
+    )
+
+    mock_attachment_with_context = attachment_context.AttachmentWithContext(
+        attachment_data=mock_attachment,
+        entity_type="span",
+        entity_id="span-999",
+        project_name="test-project",
+        context="metadata",
+    )
+
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        if context == "metadata":
+            return [mock_attachment_with_context]
+        return []
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    assert mock_streamer.put.call_count == 2
+    attachment_msg = mock_streamer.put.call_args_list[0][0][0]
+    assert isinstance(attachment_msg, messages.CreateAttachmentMessage)
+
+
+def test_process_trace_message_extracts_attachments(
+    processor, mock_streamer, temp_file_15kb
+):
+    """Test extraction from trace messages."""
+    original_message = messages.CreateTraceMessage(
+        trace_id="trace-123",
+        project_name="test-project",
+        name="test-trace",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"data": "fake-base64"},
+        output=None,
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    mock_attachment = attachment.Attachment(
+        data=temp_file_15kb.name,
+        file_name="input-attachment-trace.png",
+        content_type="image/png",
+        create_temp_copy=False,
+    )
+
+    mock_attachment_with_context = attachment_context.AttachmentWithContext(
+        attachment_data=mock_attachment,
+        entity_type="trace",
+        entity_id="trace-123",
+        project_name="test-project",
+        context="input",
+    )
+
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        if context == "input":
+            return [mock_attachment_with_context]
+        return []
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    # Should create an attachment message and re-queue the original message
+    assert mock_streamer.put.call_count == 2
+
+    # Verify it's a trace entity type
+    attachment_msg = mock_streamer.put.call_args_list[0][0][0]
+    assert isinstance(attachment_msg, messages.CreateAttachmentMessage)
+    assert attachment_msg.entity_type == "trace"
+    assert attachment_msg.entity_id == "trace-123"
+    assert attachment_msg.delete_after_upload is True
+
+    # Last call should be the original message with marker
+    second_call = mock_streamer.put.call_args_list[1][0][0]
+    assert second_call is original_message
+    assert hasattr(second_call, constants.MARKER_ATTRIBUTE_NAME)
+
+
+def test_process_update_span_message(processor, mock_streamer, temp_file_15kb):
+    """Test extraction from UpdateSpanMessage."""
+    original_message = messages.UpdateSpanMessage(
+        span_id="span-update-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        end_time=mock.Mock(),
+        input={"image": "base64-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    mock_attachment = attachment.Attachment(
+        data=temp_file_15kb.name,
+        file_name="input-attachment.png",
+        content_type="image/png",
+        create_temp_copy=False,
+    )
+
+    mock_attachment_with_context = attachment_context.AttachmentWithContext(
+        attachment_data=mock_attachment,
+        entity_type="span",
+        entity_id="span-update-123",
+        project_name="test-project",
+        context="input",
+    )
+
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        if context == "input":
+            return [mock_attachment_with_context]
+        return []
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    assert mock_streamer.put.call_count == 2
+
+    attachment_msg = mock_streamer.put.call_args_list[0][0][0]
+    assert isinstance(attachment_msg, messages.CreateAttachmentMessage)
+
+
+def test_process_update_trace_message(processor, mock_streamer, temp_file_15kb):
+    """Test extraction from UpdateTraceMessage."""
+    original_message = messages.UpdateTraceMessage(
+        trace_id="trace-update-123",
+        project_name="test-project",
+        end_time=mock.Mock(),
+        input=None,
+        output={"result": "base64-data"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    mock_attachment = attachment.Attachment(
+        data=temp_file_15kb.name,
+        file_name="output-attachment.png",
+        content_type="image/png",
+        create_temp_copy=False,
+    )
+
+    mock_attachment_with_context = attachment_context.AttachmentWithContext(
+        attachment_data=mock_attachment,
+        entity_type="trace",
+        entity_id="trace-update-123",
+        project_name="test-project",
+        context="output",
+    )
+
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        if context == "output":
+            return [mock_attachment_with_context]
+        return []
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    assert mock_streamer.put.call_count == 2
+
+    attachment_msg = mock_streamer.put.call_args_list[0][0][0]
+    assert isinstance(attachment_msg, messages.CreateAttachmentMessage)
+
+
+def test_process_multiple_attachments_from_different_contexts(
+    processor, mock_streamer, temp_file_15kb
+):
+    """Test extraction of multiple attachments from input, output, and metadata."""
+    original_message = messages.CreateSpanMessage(
+        span_id="span-multi",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"image1": "base64"},
+        output={"image2": "base64"},
+        metadata={"doc": "base64"},
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Create three attachments for different contexts
+    attachments_map = {
+        "input": [
+            attachment_context.AttachmentWithContext(
+                attachment_data=attachment.Attachment(
+                    data=temp_file_15kb.name,
+                    file_name="input-att.png",
+                    content_type="image/png",
+                    create_temp_copy=False,
+                ),
+                entity_type="span",
+                entity_id="span-multi",
+                project_name="test-project",
+                context="input",
+            )
+        ],
+        "output": [
+            attachment_context.AttachmentWithContext(
+                attachment_data=attachment.Attachment(
+                    data=temp_file_15kb.name,
+                    file_name="output-att.png",
+                    content_type="image/png",
+                    create_temp_copy=False,
+                ),
+                entity_type="span",
+                entity_id="span-multi",
+                project_name="test-project",
+                context="output",
+            )
+        ],
+        "metadata": [
+            attachment_context.AttachmentWithContext(
+                attachment_data=attachment.Attachment(
+                    data=temp_file_15kb.name,
+                    file_name="meta-att.pdf",
+                    content_type="application/pdf",
+                    create_temp_copy=False,
+                ),
+                entity_type="span",
+                entity_id="span-multi",
+                project_name="test-project",
+                context="metadata",
+            )
+        ],
+    }
+
+    def mock_extract(data, entity_type, entity_id, project_name, context):
+        return attachments_map.get(context, [])
+
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", side_effect=mock_extract
+    ):
+        processor.process(wrapped_message)
+
+    # Should create 3 attachment messages + 1 re-queued original = 4 total
+    assert mock_streamer.put.call_count == 4
+
+    # First 3 should be attachment messages
+    for i in range(3):
+        msg = mock_streamer.put.call_args_list[i][0][0]
+        assert isinstance(msg, messages.CreateAttachmentMessage)
+
+    # Last should be the original message
+    last_msg = mock_streamer.put.call_args_list[3][0][0]
+    assert last_msg is original_message
+
+
+def test_process_with_null_input_output_metadata(processor, mock_streamer):
+    """Test processing message with null input/output/metadata fields."""
+    original_message = messages.CreateSpanMessage(
+        span_id="span-nulls",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Extractor should never be called since all fields are None
+    with mock.patch.object(processor.extractor, "extract_and_replace") as mock_extract:
+        processor.process(wrapped_message)
+
+        # Extractor should not be called
+        mock_extract.assert_not_called()
+
+    # Only the original message should be re-queued
+    assert mock_streamer.put.call_count == 1
+
+
+def test_process_inactive_processor_skips_extraction(inactive_processor, mock_streamer):
+    """Test that the inactive processor skips extraction but still re-queues a message."""
+    original_message = messages.CreateSpanMessage(
+        span_id="span-inactive",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"image": "base64-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Should not call extractor since the processor is inactive
+    with mock.patch.object(
+        inactive_processor.extractor, "extract_and_replace"
+    ) as mock_extract:
+        inactive_processor.process(wrapped_message)
+
+        # Extractor should not be called
+        mock_extract.assert_not_called()
+
+    # Original message should still be re-queued with marker
+    assert mock_streamer.put.call_count == 1
+    requeued_message = mock_streamer.put.call_args[0][0]
+    assert requeued_message is original_message
+    assert hasattr(requeued_message, constants.MARKER_ATTRIBUTE_NAME)
+
+
+def test_process_handles_extraction_exception_gracefully(processor, mock_streamer):
+    """Test that extraction exceptions are caught and logged, but processing continues."""
+    original_message = messages.CreateSpanMessage(
+        span_id="span-error",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"data": "causes-error"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Mock extractor to raise an exception
+    with mock.patch.object(
+        processor.extractor,
+        "extract_and_replace",
+        side_effect=Exception("Extraction failed"),
+    ):
+        # Should not raise exception
+        processor.process(wrapped_message)
+
+    # Original message should still be re-queued despite the error
+    assert mock_streamer.put.call_count == 1
+    requeued_message = mock_streamer.put.call_args[0][0]
+    assert requeued_message is original_message
+    assert hasattr(requeued_message, constants.MARKER_ATTRIBUTE_NAME)
+
+
+def test_process_message_already_marked_returns_early(processor):
+    """Test that messages already marked with a preprocessing attribute are handled."""
+    # This test verifies behavior when a message with marker is wrapped
+    # In practice, preprocess_message() should prevent this, but we test defense
+    original_message = messages.CreateSpanMessage(
+        span_id="span-marked",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=mock.Mock(),
+        end_time=None,
+        input={"data": "value"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    # Mark the message as already processed
+    setattr(original_message, constants.MARKER_ATTRIBUTE_NAME, True)
+
+    wrapped_message = messages.AttachmentSupportingMessage(
+        original_message=original_message
+    )
+
+    # Mock extractor to verify it's still called (marker is on the original, not wrapper)
+    with mock.patch.object(
+        processor.extractor, "extract_and_replace", return_value=[]
+    ) as mock_extract:
+        processor.process(wrapped_message)
+
+        # Extractor should still be called since the wrapper doesn't have the marker
+        mock_extract.assert_called()
+
+
+def test_entity_type_from_attachment_message_span():
+    """Test entity_type_from_attachment_message extracts correct details for span."""
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    result = attachments_extraction_processor.entity_type_from_attachment_message(
+        span_message
+    )
+
+    assert result is not None
+    assert result.entity_type == "span"
+    assert result.entity_id == "span-123"
+    assert result.project_name == "test-project"
+
+
+def test_entity_type_from_attachment_message_trace():
+    """Test entity_type_from_attachment_message extracts correct details for trace."""
+    trace_message = messages.CreateTraceMessage(
+        trace_id="trace-789",
+        project_name="my-project",
+        name="test-trace",
+        start_time=mock.Mock(),
+        end_time=None,
+        input=None,
+        output=None,
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=mock.Mock(),
+        source="sdk",
+    )
+
+    result = attachments_extraction_processor.entity_type_from_attachment_message(
+        trace_message
+    )
+
+    assert result is not None
+    assert result.entity_type == "trace"
+    assert result.entity_id == "trace-789"
+    assert result.project_name == "my-project"
+
+
+def test_entity_type_from_attachment_message_unsupported():
+    """Test entity_type_from_attachment_message returns None for unsupported message types."""
+    unsupported_message = messages.BaseMessage()
+
+    result = attachments_extraction_processor.entity_type_from_attachment_message(
+        unsupported_message
+    )
+
+    assert result is None

sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py

@@ -0,0 +1,905 @@
+"""
+Unit tests for OpikMessageProcessor interaction with ReplayManager.
+
+Tests the three replay-related code paths in OpikMessageProcessor.process():
+- register_message: called before every handler invocation
+- unregister_message: called after successful handler execution
+- message_sent_failed: called on httpx.ConnectError / httpx.TimeoutException
+"""
+
+import datetime
+from unittest import mock
+
+import httpx
+import pydantic
+import pytest
+import tenacity
+
+from opik.message_processing import messages
+from opik.message_processing.processors import online_message_processor
+from opik.message_processing.processors.online_message_processor import (
+    OpikMessageProcessor,
+)
+from opik.file_upload.s3_multipart_upload import s3_upload_error
+from opik.message_processing.replay import replay_manager, db_manager
+from opik.rest_api import core as rest_api_core
+from opik import exceptions
+
+
+def _create_trace_message(message_id: int = 1) -> messages.CreateTraceMessage:
+    msg = messages.CreateTraceMessage(
+        trace_id="trace-1",
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
+        input={"query": "test"},
+        output={"answer": "response"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _create_attachment_message(message_id: int = 1) -> messages.CreateAttachmentMessage:
+    msg = messages.CreateAttachmentMessage(
+        file_path="/tmp/test-file.bin",
+        file_name="test-file.bin",
+        mime_type="application/octet-stream",
+        entity_type="trace",
+        entity_id="trace-1",
+        project_name="test-project",
+        encoded_url_override="aHR0cDovL2xvY2FsaG9zdA==",
+        delete_after_upload=False,
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _create_span_message(message_id: int = 1) -> messages.CreateSpanMessage:
+    msg = messages.CreateSpanMessage(
+        span_id="span-1",
+        trace_id="trace-1",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
+        input={"prompt": "test"},
+        output={"response": "result"},
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+@pytest.fixture
+def mock_rest_client() -> mock.MagicMock:
+    # No spec= so nested attributes like .traces.create_trace are auto-created
+    return mock.MagicMock()
+
+
+@pytest.fixture
+def mock_file_uploader() -> mock.MagicMock:
+    return mock.MagicMock()
+
+
+@pytest.fixture
+def mock_replay() -> mock.MagicMock:
+    return mock.MagicMock(spec=replay_manager.ReplayManager)
+
+
+@pytest.fixture
+def processor(
+    mock_rest_client: mock.MagicMock,
+    mock_file_uploader: mock.MagicMock,
+    mock_replay: mock.MagicMock,
+) -> OpikMessageProcessor:
+    return OpikMessageProcessor(
+        rest_client=mock_rest_client,
+        file_upload_manager=mock_file_uploader,
+        fallback_replay_manager=mock_replay,
+        unauthorized_message_types_registry=mock.Mock(),
+    )
+
+
+class TestRegisterMessage:
+    def test_process__known_message__registers_before_handler(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """register_message must be called before the REST API handler fires."""
+        call_order = []
+        mock_replay.register_message.side_effect = lambda m: call_order.append(
+            "register"
+        )
+        mock_rest_client.traces.create_trace.side_effect = (
+            lambda **kw: call_order.append("handler")
+        )
+
+        msg = _create_trace_message()
+        processor.process(msg)
+
+        assert "register" in call_order
+        assert "handler" in call_order
+        assert call_order.index("register") < call_order.index("handler")
+
+    def test_process__trace_message__registers_with_correct_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """register_message is called with the exact message object."""
+        msg = _create_trace_message(message_id=42)
+        processor.process(msg)
+
+        mock_replay.register_message.assert_called_once_with(msg)
+
+    def test_process__unknown_message_type__does_not_register(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """Unknown message types should be skipped without registration."""
+        msg = mock.MagicMock(spec=messages.BaseMessage)
+        # Use a type not in the handlers dict
+        msg.__class__ = type("UnknownMessage", (messages.BaseMessage,), {})
+
+        processor.process(msg)
+
+        mock_replay.register_message.assert_not_called()
+
+    def test_process__inactive_processor__does_not_register(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        mock_replay: mock.MagicMock,
+    ):
+        """When the processor is inactive, no registration should occur."""
+        inactive_processor = OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=mock_replay,
+            active=False,
+            unauthorized_message_types_registry=mock.Mock(),
+        )
+
+        msg = _create_trace_message()
+        inactive_processor.process(msg)
+
+        mock_replay.register_message.assert_not_called()
+
+
+class TestNoServerConnection:
+    """Tests for when has_server_connection is False — messages should be
+    registered as failed and the handler should NOT execute."""
+
+    @pytest.fixture
+    def offline_replay(self) -> mock.MagicMock:
+        m = mock.MagicMock(spec=replay_manager.ReplayManager)
+        m.has_server_connection = False
+        return m
+
+    @pytest.fixture
+    def offline_processor(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        offline_replay: mock.MagicMock,
+    ) -> OpikMessageProcessor:
+        return OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=offline_replay,
+            unauthorized_message_types_registry=mock.Mock(),
+        )
+
+    def test_process__no_connection__registers_message_as_failed(
+        self,
+        offline_processor: OpikMessageProcessor,
+        offline_replay: mock.MagicMock,
+    ):
+        """When there is no server connection, register_message must be called
+        with status=MessageStatus.failed."""
+        msg = _create_trace_message(message_id=10)
+        offline_processor.process(msg)
+
+        offline_replay.register_message.assert_called_once_with(
+            msg, status=db_manager.MessageStatus.failed
+        )
+        offline_replay.unregister_message.assert_not_called()
+
+    def test_process__no_connection__handler_not_called(
+        self,
+        offline_processor: OpikMessageProcessor,
+        offline_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """When there is no server connection, the REST API handler must NOT
+        be invoked."""
+        msg = _create_trace_message()
+        offline_processor.process(msg)
+
+        mock_rest_client.traces.create_trace.assert_not_called()
+        offline_replay.unregister_message.assert_not_called()
+
+    def test_process__no_connection__unregister_not_called(
+        self,
+        offline_processor: OpikMessageProcessor,
+        offline_replay: mock.MagicMock,
+    ):
+        """When there is no server connection, unregister_message must NOT be
+        called since the handler is skipped."""
+        msg = _create_trace_message()
+        offline_processor.process(msg)
+
+        offline_replay.unregister_message.assert_not_called()
+
+    def test_process__no_connection__message_sent_failed_not_called(
+        self,
+        offline_processor: OpikMessageProcessor,
+        offline_replay: mock.MagicMock,
+    ):
+        """message_sent_failed should NOT be called — it is only for connection
+        errors that occur during handler execution."""
+        msg = _create_trace_message()
+        offline_processor.process(msg)
+
+        offline_replay.message_sent_failed.assert_not_called()
+
+
+class TestUnregisterMessage:
+    def test_process__successful_handler__unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """After successful handler execution, unregister_message is called."""
+        msg = _create_trace_message(message_id=10)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(10)
+
+    def test_process__successful_handler__unregisters_after_handler(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """unregister_message must be called after the handler completes."""
+        call_order = []
+        mock_rest_client.traces.create_trace.side_effect = (
+            lambda **kw: call_order.append("handler")
+        )
+        mock_replay.unregister_message.side_effect = lambda mid: call_order.append(
+            "unregister"
+        )
+
+        msg = _create_trace_message()
+        processor.process(msg)
+
+        assert call_order.index("handler") < call_order.index("unregister")
+
+    def test_process__api_error_409__unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """409 Conflict triggers an explicit unregister_message call and early
+        return (the message is considered delivered)."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=409, body="conflict"
+        )
+
+        msg = _create_trace_message(message_id=11)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(11)
+
+    def test_process__api_error_500__unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """Non-connection API errors (e.g. 500) still unregister the message."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=500, body="internal error"
+        )
+
+        msg = _create_trace_message(message_id=12)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(12)
+
+    def test_process__retry_error__unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """tenacity.RetryError unregisters the message (not a connection error)."""
+        future = tenacity.Future(attempt_number=1)
+        future.set_exception(RuntimeError("underlying cause"))
+        mock_rest_client.traces.create_trace.side_effect = tenacity.RetryError(
+            last_attempt=future,
+        )
+
+        msg = _create_trace_message(message_id=13)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(13)
+
+    def test_process__validation_error__unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """pydantic.ValidationError unregisters the message."""
+        # Generate a real ValidationError via pydantic
+        validation_error: pydantic.ValidationError
+        try:
+            pydantic.TypeAdapter(int).validate_python("not-an-int")
+        except pydantic.ValidationError as e:
+            validation_error = e
+
+        mock_rest_client.traces.create_trace.side_effect = validation_error
+
+        msg = _create_trace_message(message_id=14)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(14)
+
+    def test_process__generic_exception__unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """A generic Exception still unregisters the message."""
+        mock_rest_client.traces.create_trace.side_effect = RuntimeError("unexpected")
+
+        msg = _create_trace_message(message_id=15)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(15)
+
+    @pytest.mark.parametrize(
+        "connection_error",
+        [httpx.ConnectError("refused"), httpx.ReadTimeout("timed out")],
+    )
+    def test_process__connection_error__does_not_unregister(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+        connection_error: Exception,
+    ):
+        """httpx connection/timeout errors must NOT unregister — the message
+        stays registered for replay."""
+        mock_rest_client.traces.create_trace.side_effect = connection_error
+
+        msg = _create_trace_message(message_id=16)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_not_called()
+
+    def test_process__api_error_429_with_rate_limit__does_not_unregister(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """429 with valid rate-limit headers raises and leaves the message
+        registered for retry — unregister must NOT be called."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=429, body="rate limited", headers={"RateLimit-Reset": "10"}
+        )
+
+        msg = _create_trace_message(message_id=17)
+
+        with pytest.raises(exceptions.OpikCloudRequestsRateLimited):
+            processor.process(msg)
+
+        mock_replay.unregister_message.assert_not_called()
+
+
+class TestMessageSentFailed:
+    @pytest.mark.parametrize(
+        "connection_error",
+        [httpx.ConnectError("connection refused"), httpx.ReadTimeout("read timed out")],
+    )
+    def test_process__connection_error__marks_message_as_failed(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+        connection_error: Exception,
+    ):
+        """httpx.TimeoutException should trigger message_sent_failed."""
+        mock_rest_client.traces.create_trace.side_effect = connection_error
+
+        msg = _create_trace_message(message_id=8)
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_called_once_with(
+            8, failure_reason=str(connection_error)
+        )
+
+        # assert that unregister_message was not called
+        mock_replay.unregister_message.assert_not_called()
+
+
+class TestErrorPathsNoFailedMark:
+    """Verify that non-connection errors do NOT call message_sent_failed.
+    The message stays registered (not unregistered, not marked failed)."""
+
+    def test_process__api_error_409__no_failed_mark(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """409 Conflict is silently ignored and a message unregistered."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=409, body="conflict"
+        )
+
+        msg = _create_trace_message(message_id=1)
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+        mock_replay.unregister_message.assert_called()
+
+    def test_process__api_error_500__no_failed_mark(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """500 Server Error logs but does not mark as failed."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=500, body="internal error"
+        )
+
+        msg = _create_trace_message(message_id=2)
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+        mock_replay.unregister_message.assert_called()
+
+    def test_process__generic_exception__no_failed_mark(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """A generic Exception does not mark as failed."""
+        mock_rest_client.traces.create_trace.side_effect = RuntimeError("unexpected")
+
+        msg = _create_trace_message(message_id=3)
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+        mock_replay.unregister_message.assert_called()
+
+
+class TestApiError429WithRateLimiter:
+    """Verify that a 429 ApiError with valid rate-limit headers raises
+    OpikCloudRequestsRateLimited, leaving the message registered for retry."""
+
+    def test_process__api_error_429_with_rate_limit_headers__no_unregister_no_failed(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """Neither unregister nor message_sent_failed is called — a message stays
+        registered so the queue consumer can retry it after the rate-limit window."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=429, body="rate limited", headers={"RateLimit-Reset": "10"}
+        )
+
+        msg = _create_trace_message(message_id=3)
+
+        with pytest.raises(exceptions.OpikCloudRequestsRateLimited):
+            processor.process(msg)
+
+        mock_replay.register_message.assert_called_once_with(msg)
+        mock_replay.unregister_message.assert_not_called()
+        mock_replay.message_sent_failed.assert_not_called()
+
+    def test_process__api_error_429_without_rate_limit_headers__no_raise(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """429 without parseable rate-limit headers falls through to generic
+        ApiError logging, unregisters the message, and does not raise."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=429, body="rate limited", headers={"Other-Header": "value"}
+        )
+
+        msg = _create_trace_message(message_id=4)
+
+        # Should NOT raise — falls through to the generic ApiError branch
+        processor.process(msg)
+
+        mock_replay.register_message.assert_called_once_with(msg)
+
+        # Not a connection error, so a message is unregistered
+        mock_replay.unregister_message.assert_called_once_with(4)
+        mock_replay.message_sent_failed.assert_not_called()
+
+    def test_process__api_error_429_with_none_headers__no_raise(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """429 with headers=None skips rate-limit parsing, unregisters the
+        message, and does not raise."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=429, body="rate limited", headers=None
+        )
+
+        msg = _create_trace_message(message_id=5)
+
+        processor.process(msg)
+
+        mock_replay.register_message.assert_called_once_with(msg)
+
+        # Not a connection error, so a message is unregistered
+        mock_replay.unregister_message.assert_called_once_with(5)
+        mock_replay.message_sent_failed.assert_not_called()
+
+
+class TestMultipleMessagesSequence:
+    def test_process__success_then_connection_error__correct_replay_calls(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """The first message succeeds (register+unregister), the second fails with
+        ConnectError (register+message_sent_failed)."""
+        # The first call succeeds, the second raises ConnectError
+        mock_rest_client.traces.create_trace.side_effect = [
+            None,
+            httpx.ConnectError("connection refused"),
+        ]
+
+        msg1 = _create_trace_message(message_id=1)
+        msg2 = _create_trace_message(message_id=2)
+
+        processor.process(msg1)
+        processor.process(msg2)
+
+        # Both registered
+        assert mock_replay.register_message.call_count == 2
+
+        # Only first unregistered
+        mock_replay.unregister_message.assert_called_once_with(1)
+
+        # Only second marked as failed
+        mock_replay.message_sent_failed.assert_called_once_with(
+            2, failure_reason="connection refused"
+        )
+
+    def test_process__connection_error_then_success__correct_replay_calls(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """The first message fails with a connection error, the second succeeds."""
+        mock_rest_client.traces.create_trace.side_effect = [
+            httpx.ConnectError("connection refused"),
+            None,
+        ]
+
+        msg1 = _create_trace_message(message_id=1)
+        msg2 = _create_trace_message(message_id=2)
+
+        processor.process(msg1)
+        processor.process(msg2)
+
+        assert mock_replay.register_message.call_count == 2
+        mock_replay.message_sent_failed.assert_called_once_with(
+            1, failure_reason="connection refused"
+        )
+        mock_replay.unregister_message.assert_called_once_with(2)
+
+
+class TestAttachmentUploadCallbacks:
+    """Tests for CreateAttachmentMessage upload callback interactions with
+    the replay manager.
+
+    The attachment handler passes on_upload_success / on_upload_failed callbacks
+    to the file uploader. These callbacks are responsible for calling
+    unregister_message or message_sent_failed on the replay manager."""
+
+    def test_process__attachment__registers_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """CreateAttachmentMessage should be registered before the upload handler."""
+        msg = _create_attachment_message(message_id=50)
+        processor.process(msg)
+
+        mock_replay.register_message.assert_called_once_with(msg)
+
+    def test_process__attachment__passes_callbacks_to_uploader(
+        self,
+        processor: OpikMessageProcessor,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """The file uploader must receive on_upload_success and on_upload_failed
+        callbacks."""
+        msg = _create_attachment_message(message_id=50)
+        processor.process(msg)
+
+        mock_file_uploader.upload.assert_called_once()
+        call_kwargs = mock_file_uploader.upload.call_args
+        assert call_kwargs.kwargs["on_upload_success"] is not None
+        assert call_kwargs.kwargs["on_upload_failed"] is not None
+
+    def test_process__attachment__success_callback_unregisters_message(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When the on_upload_success callback fires, it should call
+        unregister_message with the correct message_id."""
+        msg = _create_attachment_message(message_id=60)
+
+        # Capture the callback by invoking it inside the mock upload
+        def fake_upload(message, on_upload_success=None, on_upload_failed=None):
+            on_upload_success()
+
+        mock_file_uploader.upload.side_effect = fake_upload
+        mock_replay.reset_mock()
+
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_any_call(60)
+
+    def test_process__attachment__failed_callback_connect_error__marks_failed(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When on_upload_failed fires with httpx.ConnectError, it should call
+        message_sent_failed."""
+        msg = _create_attachment_message(message_id=70)
+        error = httpx.ConnectError("connection refused")
+
+        def fake_upload(message, on_upload_success=None, on_upload_failed=None):
+            on_upload_failed(error)
+
+        mock_file_uploader.upload.side_effect = fake_upload
+        mock_replay.reset_mock()
+
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_called_once_with(
+            70, failure_reason=str(error)
+        )
+
+    def test_process__attachment__failed_callback_timeout__marks_failed(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When on_upload_failed fires with httpx.TimeoutException, it should call
+        message_sent_failed."""
+        msg = _create_attachment_message(message_id=71)
+        error = httpx.ReadTimeout("read timed out")
+
+        def fake_upload(message, on_upload_success=None, on_upload_failed=None):
+            on_upload_failed(error)
+
+        mock_file_uploader.upload.side_effect = fake_upload
+        mock_replay.reset_mock()
+
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_called_once_with(
+            71, failure_reason=str(error)
+        )
+
+    def test_process__attachment__failed_callback_s3_connection_error__marks_failed(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When on_upload_failed fires with S3UploadError(connection_error=True),
+        it should call message_sent_failed."""
+        msg = _create_attachment_message(message_id=72)
+        error = s3_upload_error.S3UploadError(
+            reason="connection lost", connection_error=True
+        )
+
+        def fake_upload(message, on_upload_success=None, on_upload_failed=None):
+            on_upload_failed(error)
+
+        mock_file_uploader.upload.side_effect = fake_upload
+        mock_replay.reset_mock()
+
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_called_once_with(
+            72, failure_reason=str(error)
+        )
+
+    def test_process__attachment__failed_callback_s3_non_connection_error__no_failed_mark(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When on_upload_failed fires with S3UploadError(connection_error=False),
+        message_sent_failed should NOT be called."""
+        msg = _create_attachment_message(message_id=73)
+        error = s3_upload_error.S3UploadError(
+            reason="file too large", connection_error=False
+        )
+
+        def fake_upload(message, on_upload_success=None, on_upload_failed=None):
+            on_upload_failed(error)
+
+        mock_file_uploader.upload.side_effect = fake_upload
+        mock_replay.reset_mock()
+
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+        mock_replay.unregister_message.assert_called_once_with(73)
+
+    def test_process__attachment__failed_callback_generic_error__no_failed_mark(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When on_upload_failed fires with a non-connection error (e.g. ValueError),
+        message_sent_failed should NOT be called."""
+        msg = _create_attachment_message(message_id=74)
+        error = ValueError("bad data")
+
+        def fake_upload(message, on_upload_success=None, on_upload_failed=None):
+            on_upload_failed(error)
+
+        mock_file_uploader.upload.side_effect = fake_upload
+        mock_replay.reset_mock()
+
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+        mock_replay.unregister_message.assert_called_once_with(74)
+
+    def test_process__attachment_no_connection__registers_as_failed_and_skips_upload(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """When has_server_connection is False, the attachment message should be
+        registered as failed and the file uploader should NOT be called."""
+        offline_replay = mock.MagicMock(spec=replay_manager.ReplayManager)
+        offline_replay.has_server_connection = False
+        offline_processor = OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=offline_replay,
+            unauthorized_message_types_registry=mock.Mock(),
+        )
+
+        msg = _create_attachment_message(message_id=80)
+        offline_processor.process(msg)
+
+        offline_replay.register_message.assert_called_once_with(
+            msg, status=db_manager.MessageStatus.failed
+        )
+        mock_file_uploader.upload.assert_not_called()
+        offline_replay.unregister_message.assert_not_called()
+
+
+def _create_attachment_supporting_message(
+    message_id: int = 1,
+) -> messages.AttachmentSupportingMessage:
+    inner = _create_trace_message(message_id)
+    msg = messages.AttachmentSupportingMessage(original_message=inner)
+    msg.message_id = message_id
+    return msg
+
+
+class TestIgnoredMessageTypes:
+    """Tests for messages listed in _ignored_message_types_for_replay.
+
+    AttachmentSupportingMessage is the sole type on this list. Such messages
+    bypass replay registration entirely — the handler runs directly regardless
+    of connection state, and replay manager methods are never invoked.
+    """
+
+    def test_process__ignored_type__no_replay_on_online_connection__handler_still_called(
+        self,
+        processor: OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """All three replay manager methods are untouched for ignored types when online."""
+        noop_handler = mock.MagicMock(
+            spec=online_message_processor.MessageProcessingHandler
+        )
+        processor.register_message_handler(
+            handler=noop_handler, message_type=messages.AttachmentSupportingMessage
+        )
+
+        msg = _create_attachment_supporting_message(message_id=106)
+        processor.process(msg)
+
+        # check that the noop handler was called
+        noop_handler.assert_called_once_with(msg)
+
+        # No replay interaction at all
+        mock_replay.register_message.assert_not_called()
+        mock_replay.unregister_message.assert_not_called()
+        mock_replay.message_sent_failed.assert_not_called()
+
+    def test_process__ignored_type__no_replay_called_on_offline_connection__handler_still_called(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+    ):
+        """The handler (noop) must be invoked for ignored types even when in offline mode —
+        replay state does not block execution for these message types."""
+        offline_replay = mock.MagicMock(spec=replay_manager.ReplayManager)
+        offline_replay.has_server_connection = False
+        offline_processor = OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=offline_replay,
+            unauthorized_message_types_registry=mock.Mock(),
+        )
+
+        noop_handler = mock.MagicMock(
+            spec=online_message_processor.MessageProcessingHandler
+        )
+        offline_processor.register_message_handler(
+            handler=noop_handler, message_type=messages.AttachmentSupportingMessage
+        )
+
+        # Process completes without error — the noop handler runs
+        msg = _create_attachment_supporting_message(message_id=105)
+        offline_processor.process(msg)
+
+        # check that the noop handler was called
+        noop_handler.assert_called_once_with(msg)
+
+        # No replay interaction at all
+        offline_replay.register_message.assert_not_called()
+        offline_replay.unregister_message.assert_not_called()
+        offline_replay.message_sent_failed.assert_not_called()

sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_unauthorized_message_type.py

@@ -0,0 +1,409 @@
+"""
+Unit tests for OpikMessageProcessor interaction with UnauthorizedMessageTypeRegistry.
+
+Tests cover:
+- Unauthorized message types (is_authorized returns False) are skipped:
+  handler not called, replay manager not touched.
+- 401 ApiError triggers registry.add() with the message type.
+- 401 ApiError still causes message unregistration (not a connection error).
+- Non-401 errors do NOT add to the registry.
+- End-to-end: after a 401 causes a type to be registered, later messages
+  of that type are blocked; messages of other types are unaffected.
+"""
+
+import datetime
+from unittest import mock
+
+import pytest
+
+from opik.message_processing import messages
+from opik.message_processing import permissions
+from opik.message_processing.processors import online_message_processor
+from opik.message_processing.replay import replay_manager
+from opik.rest_api import core as rest_api_core
+
+
+def _create_trace_message(message_id: int = 1) -> messages.CreateTraceMessage:
+    msg = messages.CreateTraceMessage(
+        trace_id="trace-1",
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.datetime(
+            2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc
+        ),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1, tzinfo=datetime.timezone.utc),
+        input={"query": "test"},
+        output={"answer": "response"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _create_span_message(message_id: int = 1) -> messages.CreateSpanMessage:
+    msg = messages.CreateSpanMessage(
+        span_id="span-1",
+        trace_id="trace-1",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.datetime(
+            2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc
+        ),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1, tzinfo=datetime.timezone.utc),
+        input={"prompt": "test"},
+        output={"response": "result"},
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+@pytest.fixture
+def mock_rest_client() -> mock.MagicMock:
+    return mock.MagicMock()
+
+
+@pytest.fixture
+def mock_file_uploader() -> mock.MagicMock:
+    return mock.MagicMock()
+
+
+@pytest.fixture
+def mock_replay() -> mock.MagicMock:
+    return mock.MagicMock(spec=replay_manager.ReplayManager)
+
+
+@pytest.fixture
+def mock_registry() -> mock.MagicMock:
+    registry = mock.MagicMock(spec=permissions.UnauthorizedMessageTypeRegistry)
+    registry.is_authorized.return_value = True
+    return registry
+
+
+@pytest.fixture
+def processor(
+    mock_rest_client: mock.MagicMock,
+    mock_file_uploader: mock.MagicMock,
+    mock_replay: mock.MagicMock,
+    mock_registry: mock.MagicMock,
+) -> online_message_processor.OpikMessageProcessor:
+    return online_message_processor.OpikMessageProcessor(
+        rest_client=mock_rest_client,
+        file_upload_manager=mock_file_uploader,
+        fallback_replay_manager=mock_replay,
+        unauthorized_message_types_registry=mock_registry,
+    )
+
+
+def test_process__authorized_message_type__handler_called(
+    processor: online_message_processor.OpikMessageProcessor,
+    mock_rest_client: mock.MagicMock,
+):
+    """Control: when is_authorized returns True, the handler must be invoked."""
+    msg = _create_trace_message()
+    processor.process(msg)
+
+    mock_rest_client.traces.create_trace.assert_called_once()
+
+
+class TestUnauthorizedMessageTypeBlocking:
+    """Tests for the early-exit path when is_authorized() returns False.
+
+    When is_authorized returns False, process() should return immediately
+    without invoking the handler or any replay manager methods.
+    """
+
+    @pytest.fixture
+    def unauthorized_processor(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        mock_replay: mock.MagicMock,
+    ) -> online_message_processor.OpikMessageProcessor:
+        registry = mock.MagicMock(spec=permissions.UnauthorizedMessageTypeRegistry)
+        registry.is_authorized.return_value = False
+        return online_message_processor.OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=mock_replay,
+            unauthorized_message_types_registry=registry,
+        )
+
+    def test_process__unauthorized_message_type__handler_not_called(
+        self,
+        unauthorized_processor: online_message_processor.OpikMessageProcessor,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """When the message type is unauthorized, no REST handler should fire."""
+        msg = _create_trace_message()
+        unauthorized_processor.process(msg)
+
+        mock_rest_client.traces.create_trace.assert_not_called()
+
+    def test_process__unauthorized_message_type__register_not_called(
+        self,
+        unauthorized_processor: online_message_processor.OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """When the message type is unauthorized, register_message must not be called."""
+        msg = _create_trace_message()
+        unauthorized_processor.process(msg)
+
+        mock_replay.register_message.assert_not_called()
+
+    def test_process__unauthorized_message_type__unregister_not_called(
+        self,
+        unauthorized_processor: online_message_processor.OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """When the message type is unauthorized, unregister_message must not be called."""
+        msg = _create_trace_message()
+        unauthorized_processor.process(msg)
+
+        mock_replay.unregister_message.assert_not_called()
+
+    def test_process__unauthorized_message_type__message_sent_failed_not_called(
+        self,
+        unauthorized_processor: online_message_processor.OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+    ):
+        """When the message type is unauthorized, message_sent_failed must not be called."""
+        msg = _create_trace_message()
+        unauthorized_processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+
+    def test_process__unauthorized_message_type__is_authorized_called_with_message_type(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        mock_replay: mock.MagicMock,
+    ):
+        """is_authorized must be invoked with message.message_type."""
+        registry = mock.MagicMock(spec=permissions.UnauthorizedMessageTypeRegistry)
+        registry.is_authorized.return_value = False
+
+        proc = online_message_processor.OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=mock_replay,
+            unauthorized_message_types_registry=registry,
+        )
+
+        msg = _create_trace_message()
+        proc.process(msg)
+
+        registry.is_authorized.assert_called_once_with(msg.message_type)
+
+
+class TestUnauthorizedOnApiError401:
+    """Tests for the 401 ApiError path.
+
+    A 401 response must cause the message type to be added to the unauthorized
+    registry. It should also unregister the message (not a connection error) but
+    must not call message_sent_failed.
+    """
+
+    def test_process__api_error_401__adds_message_type_to_registry(
+        self,
+        processor: online_message_processor.OpikMessageProcessor,
+        mock_registry: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """On 401, registry.add must be called with message.message_type."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=401, body="Unauthorized"
+        )
+
+        msg = _create_trace_message()
+        processor.process(msg)
+
+        mock_registry.add.assert_called_once_with(msg.message_type)
+
+    def test_process__api_error_401__unregisters_message(
+        self,
+        processor: online_message_processor.OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """On 401, unregister_message must be called (not a connection error)."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=401, body="Unauthorized"
+        )
+
+        msg = _create_trace_message(message_id=20)
+        processor.process(msg)
+
+        mock_replay.unregister_message.assert_called_once_with(20)
+
+    def test_process__api_error_401__message_sent_failed_not_called(
+        self,
+        processor: online_message_processor.OpikMessageProcessor,
+        mock_replay: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """On 401, message_sent_failed must NOT be called."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=401, body="Unauthorized"
+        )
+
+        msg = _create_trace_message()
+        processor.process(msg)
+
+        mock_replay.message_sent_failed.assert_not_called()
+
+    @pytest.mark.parametrize(
+        "status_code,body",
+        [
+            (500, "Internal Server Error"),
+            (409, "Conflict"),
+        ],
+    )
+    def test_process__api_error_not_401__does_not_add_to_registry(
+        self,
+        status_code,
+        body,
+        processor: online_message_processor.OpikMessageProcessor,
+        mock_registry: mock.MagicMock,
+        mock_rest_client: mock.MagicMock,
+    ):
+        """A non-401 error must NOT call registry.add — only 401 triggers this."""
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=status_code, body=body
+        )
+
+        msg = _create_trace_message()
+        processor.process(msg)
+
+        mock_registry.add.assert_not_called()
+
+
+class TestUnauthorizedMessageTypeEndToEnd:
+    """Integration tests using a real UnauthorizedMessageTypeRegistry.
+
+    These tests verify that after a 401 causes a type to be registered,
+    later messages of that type are blocked, while other types continue
+    to be processed normally.
+    """
+
+    def test_process__after_401_error__same_type_blocked(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        mock_replay: mock.MagicMock,
+        capture_log,
+    ):
+        """After a 401 error registers a type, the next message of that type must
+        be skipped — handler not called, replay not touched for the second message."""
+        real_registry = permissions.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=3600  # long interval — type stays blocked
+        )
+        proc = online_message_processor.OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=mock_replay,
+            unauthorized_message_types_registry=real_registry,
+        )
+
+        # The first call gets a 401 → type is registered as unauthorized
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=401, body="Unauthorized"
+        )
+        proc.process(_create_trace_message(message_id=1))
+
+        # Reset call tracking for the second call
+        mock_rest_client.reset_mock()
+        mock_replay.reset_mock()
+
+        # The second call — type is now blocked, handler should not fire
+        proc.process(_create_trace_message(message_id=2))
+
+        mock_rest_client.traces.create_trace.assert_not_called()
+        mock_replay.register_message.assert_not_called()
+
+        assert (
+            capture_log.records[0].message
+            == "Unauthorized message type 'CreateTraceMessage' processing request: Unauthorized"
+        )
+        assert capture_log.records[0].levelname == "ERROR"
+
+    def test_process__after_401_error__different_type_not_blocked(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        mock_replay: mock.MagicMock,
+    ):
+        """After 401 blocks CreateTraceMessage, CreateSpanMessage must still
+        be processed normally."""
+        real_registry = permissions.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=3600
+        )
+        proc = online_message_processor.OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=mock_replay,
+            unauthorized_message_types_registry=real_registry,
+        )
+
+        # Cause CreateTraceMessage to be blocked via a 401
+        mock_rest_client.traces.create_trace.side_effect = rest_api_core.ApiError(
+            status_code=401, body="Unauthorized"
+        )
+        proc.process(_create_trace_message(message_id=1))
+
+        mock_rest_client.reset_mock()
+        mock_replay.reset_mock()
+
+        # CreateSpanMessage is a different type — should still go through
+        proc.process(_create_span_message(message_id=2))
+
+        mock_rest_client.spans.create_span.assert_called_once()
+
+    def test_process__after_retry_interval__type_authorized_again(
+        self,
+        mock_rest_client: mock.MagicMock,
+        mock_file_uploader: mock.MagicMock,
+        mock_replay: mock.MagicMock,
+    ):
+        """After the retry interval has elapsed, an unauthorized type should be
+        re-authorized and its handler invoked."""
+        real_registry = permissions.UnauthorizedMessageTypeRegistry(
+            retry_interval_seconds=60
+        )
+        proc = online_message_processor.OpikMessageProcessor(
+            rest_client=mock_rest_client,
+            file_upload_manager=mock_file_uploader,
+            fallback_replay_manager=mock_replay,
+            unauthorized_message_types_registry=real_registry,
+        )
+
+        # Manually add the type with a timestamp far in the past so the
+        # retry interval (60 s) has already elapsed
+        past_time = 0.0
+        real_registry.add(
+            messages.CreateTraceMessage.message_type, attempt_time=past_time
+        )
+
+        mock_rest_client.reset_mock()
+        mock_replay.reset_mock()
+
+        # The handler should be called because the interval has elapsed
+        mock_rest_client.traces.create_trace.side_effect = None  # success
+        proc.process(_create_trace_message(message_id=3))
+
+        mock_rest_client.traces.create_trace.assert_called_once()

sdks/python/tests/unit/message_processing/replay/test_db_manager.py

@@ -0,0 +1,891 @@
+import datetime
+import sqlite3
+import threading
+import time
+from typing import Generator
+from unittest import mock
+
+import pytest
+
+from opik.message_processing import messages
+from opik.message_processing.replay import message_serialization, db_manager
+
+
+@pytest.fixture
+def manager() -> Generator[db_manager.DBManager, None, None]:
+    """Fixture that creates a DBManager and ensures cleanup after test."""
+    mgr = db_manager.DBManager(batch_size=10, batch_replay_delay=0.5)
+    yield mgr
+    mgr.close()
+
+
+@pytest.fixture
+def small_batch_manager() -> Generator[db_manager.DBManager, None, None]:
+    """Fixture that creates a DBManager with a small batch size for testing batching."""
+    mgr = db_manager.DBManager(batch_size=10, batch_replay_delay=0.1)
+    yield mgr
+    mgr.close()
+
+
+def _create_trace_message(
+    message_id: int, trace_id: str = "trace-1"
+) -> messages.CreateTraceMessage:
+    """Helper to create a CreateTraceMessage with required fields."""
+    msg = messages.CreateTraceMessage(
+        trace_id=trace_id,
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.datetime(
+            2024, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc
+        ),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1, tzinfo=datetime.timezone.utc),
+        input={"query": "test"},
+        output={"answer": "response"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _create_span_message(
+    message_id: int, span_id: str = "span-1"
+) -> messages.CreateSpanMessage:
+    """Helper to create a CreateSpanMessage with required fields."""
+    msg = messages.CreateSpanMessage(
+        span_id=span_id,
+        trace_id="trace-1",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
+        input={"prompt": "test"},
+        output={"response": "result"},
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _verify_message_was_inserted(
+    message_id: int,
+    message: messages.BaseMessage,
+    conn: sqlite3.Connection,
+    status: db_manager.MessageStatus = db_manager.MessageStatus.registered,
+) -> None:
+    """Verify a message row matches the provided message."""
+    # Verify a message was inserted
+    cursor = conn.execute(
+        "SELECT message_id, status, message_type, message_json FROM messages WHERE message_id = ?",
+        (message_id,),
+    )
+    row = cursor.fetchone()
+    assert row is not None
+    assert row[0] == message.message_id
+    assert row[1] == status
+    assert row[2] == message.message_type
+    assert row[3] == message_serialization.serialize_message(message)
+
+
+class TestReplayManagerInitialization:
+    def test_init__default_parameters__creates_db_and_schema(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that ReplayManager initializes with default parameters."""
+        assert manager.initialized
+        assert not manager.closed
+        assert not manager.failed
+        assert manager.conn is not None
+
+        # Verify schema was created
+        cursor = manager.conn.execute(
+            "SELECT name FROM sqlite_master WHERE type='table' AND name='messages'"
+        )
+        assert cursor.fetchone() is not None
+
+    def test_init__custom_connection__uses_provided_connection(self):
+        """Test that ReplayManager uses a provided connection."""
+        conn = sqlite3.connect(":memory:", check_same_thread=False)
+        mgr = db_manager.DBManager(conn=conn, batch_size=10, batch_replay_delay=0.1)
+
+        try:
+            assert mgr.conn is conn
+            assert mgr.initialized
+        finally:
+            mgr.close()
+
+    def test_init__connection_creation_failure__marks_as_error(self):
+        """Test that schema creation failure marks the manager as an error."""
+        # Create a read-only connection that will fail on CREATE TABLE
+        conn = sqlite3.connect(":memory:", check_same_thread=False)
+        # Close the connection to make it unusable
+        conn.close()
+
+        mgr = db_manager.DBManager(conn=conn, batch_size=10, batch_replay_delay=0.1)
+
+        assert mgr.failed
+        assert not mgr.initialized
+
+
+class TestReplayManagerClose:
+    def test_close__normal_operation__closes_connection(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that close properly closes the connection."""
+        assert not manager.closed
+
+        manager.close()
+
+        assert manager.closed
+
+    def test_close__called_twice__no_error(self, manager: db_manager.DBManager):
+        """Test that calling close twice doesn't raise an error."""
+        manager.close()
+        manager.close()  # Should not raise
+
+        assert manager.closed
+
+
+class TestRegisterMessage:
+    def test_register_message__single_message__inserts_into_db(
+        self, manager: db_manager.DBManager
+    ):
+        """Test registering a single message inserts it into the database."""
+        message = _create_trace_message(message_id=1)
+
+        manager.register_message(message)
+
+        # Verify a message was inserted
+        _verify_message_was_inserted(1, message, manager.conn)
+
+    def test_register_message__single_message__upsert_into_db(
+        self, manager: db_manager.DBManager
+    ):
+        """Test registering a single message inserts it into the database and upserts if exists."""
+        message = _create_trace_message(message_id=1)
+
+        manager.register_message(message, status=db_manager.MessageStatus.registered)
+
+        # Verify a message was inserted
+        _verify_message_was_inserted(1, message, manager.conn)
+
+        # register the same message again, it should be updated
+        manager.register_message(message, status=db_manager.MessageStatus.failed)
+
+        # Verify a message was upserted
+        cursor = manager.conn.execute(
+            "SELECT status FROM messages WHERE message_id = ?", (1,)
+        )
+        row = cursor.fetchone()
+        assert row[0] == db_manager.MessageStatus.failed
+
+    def test_register_message__custom_status__uses_provided_status(
+        self, manager: db_manager.DBManager
+    ):
+        """Test registering a message with a custom status."""
+        message = _create_trace_message(message_id=1)
+
+        manager.register_message(message, status=db_manager.MessageStatus.failed)
+
+        cursor = manager.conn.execute(
+            "SELECT status FROM messages WHERE message_id = ?", (1,)
+        )
+        row = cursor.fetchone()
+        assert row[0] == db_manager.MessageStatus.failed
+
+    def test_register_message__manager_failed__ignores_message(self):
+        """Test that registering a message on an uninitialized manager is ignored."""
+        conn = sqlite3.connect(":memory:", check_same_thread=False)
+        conn.execute(
+            "CREATE TABLE messages (id INTEGER)"
+        )  # Force failure on a register message operation
+        mgr = db_manager.DBManager(conn=conn, batch_size=10, batch_replay_delay=0.1)
+
+        try:
+            message = _create_trace_message(message_id=1)
+            with pytest.raises(sqlite3.OperationalError):
+                mgr.register_message(message)  # the first operation should fail
+
+            mgr.register_message(
+                message
+            )  # Should not raise as DBManager marked as failed
+        finally:
+            mgr.close()
+
+    def test_register_message__manager_closed__ignores_message(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that registering a message on a closed manager is ignored."""
+        manager.close()
+
+        message = _create_trace_message(message_id=1)
+        manager.register_message(message)  # Should not raise
+
+    def test_register_message__message_id_none__raises_value_error(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that registering a message without message_id raises an error."""
+        message = _create_trace_message(message_id=1)
+        message.message_id = None
+
+        with pytest.raises(ValueError, match="Message ID expected"):
+            manager.register_message(message)
+
+
+class TestRegisterMessages:
+    def test_register_messages__batch_of_messages__inserts_all(
+        self, manager: db_manager.DBManager
+    ):
+        """Test registering a batch of messages inserts all of them."""
+        messages_list = [
+            _create_trace_message(message_id=1, trace_id="trace-1"),
+            _create_trace_message(message_id=2, trace_id="trace-2"),
+            _create_trace_message(message_id=3, trace_id="trace-3"),
+        ]
+
+        manager.register_messages(messages_list)
+
+        cursor = manager.conn.execute("SELECT COUNT(*) FROM messages")
+        assert cursor.fetchone()[0] == 3
+
+        for message in messages_list:
+            _verify_message_was_inserted(message.message_id, message, manager.conn)
+
+    def test_register_messages__batch_of_messages__upserts_all(
+        self, manager: db_manager.DBManager
+    ):
+        """Test registering a batch of messages inserts all of them and  upserts if exists."""
+        messages_list = [
+            _create_trace_message(message_id=1, trace_id="trace-1"),
+            _create_trace_message(message_id=2, trace_id="trace-2"),
+            _create_trace_message(message_id=3, trace_id="trace-3"),
+        ]
+
+        manager.register_messages(messages_list)
+
+        # register the same messages again, it should be updated
+        manager.register_messages(messages_list, status=db_manager.MessageStatus.failed)
+
+        cursor = manager.conn.execute("SELECT COUNT(*) FROM messages")
+        assert cursor.fetchone()[0] == 3
+
+        for message in messages_list:
+            _verify_message_was_inserted(
+                message.message_id,
+                message,
+                manager.conn,
+                status=db_manager.MessageStatus.failed,
+            )
+
+    def test_register_messages__manager_closed__ignores_messages(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that registering messages on a closed manager is ignored."""
+        manager.close()
+
+        messages_list = [_create_trace_message(message_id=1)]
+        manager.register_messages(messages_list)  # Should not raise
+
+
+class TestUpdateMessage:
+    def test_update_message__to_failed_status__updates_status(
+        self, manager: db_manager.DBManager
+    ):
+        """Test updating a message status to failed."""
+        message = _create_trace_message(message_id=1)
+        manager.register_message(message)
+
+        manager.update_message(1, db_manager.MessageStatus.failed)
+
+        cursor = manager.conn.execute(
+            "SELECT status FROM messages WHERE message_id = ?", (1,)
+        )
+        assert cursor.fetchone()[0] == db_manager.MessageStatus.failed
+
+    def test_update_message__to_delivered_status__deletes_message(
+        self, manager: db_manager.DBManager
+    ):
+        """Test updating a message status to delivered deletes the record."""
+        message = _create_trace_message(message_id=1)
+        manager.register_message(message)
+
+        manager.update_message(1, db_manager.MessageStatus.delivered)
+
+        cursor = manager.conn.execute(
+            "SELECT COUNT(*) FROM messages WHERE message_id = ?", (1,)
+        )
+        assert cursor.fetchone()[0] == 0
+
+    def test_update_message__manager_closed__ignores_update(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that an updating message on closed manager is ignored."""
+        message = _create_trace_message(message_id=1)
+        manager.register_message(message)
+        manager.close()
+
+        manager.update_message(1, db_manager.MessageStatus.failed)  # Should not raise
+
+
+class TestUpdateMessagesBatch:
+    def test_update_messages_batch__to_failed_status__updates_all(
+        self, manager: db_manager.DBManager
+    ):
+        """Test updating a batch of messages to failed status."""
+        messages_list = [
+            _create_trace_message(message_id=i, trace_id=f"trace-{i}")
+            for i in range(1, 4)
+        ]
+        manager.register_messages(messages_list)
+
+        cursor = manager.conn.execute(
+            "SELECT COUNT(*) FROM messages WHERE status = ?",
+            (db_manager.MessageStatus.failed,),
+        )
+        assert cursor.fetchone()[0] == 0  # no messages updated yet
+
+        manager.update_messages_batch([1, 2, 3], db_manager.MessageStatus.failed)
+
+        cursor = manager.conn.execute(
+            "SELECT COUNT(*) FROM messages WHERE status = ?",
+            (db_manager.MessageStatus.failed,),
+        )
+        assert cursor.fetchone()[0] == 3  # all three messages updated to failed
+
+    def test_update_messages_batch__to_delivered_status__deletes_all(
+        self, manager: db_manager.DBManager
+    ):
+        """Test updating a batch of messages to delivered deletes them."""
+        messages_list = [
+            _create_trace_message(message_id=i, trace_id=f"trace-{i}")
+            for i in range(1, 4)
+        ]
+        manager.register_messages(messages_list)
+
+        manager.update_messages_batch([1, 2, 3], db_manager.MessageStatus.delivered)
+
+        cursor = manager.conn.execute("SELECT COUNT(*) FROM messages")
+        assert cursor.fetchone()[0] == 0
+
+
+class TestGetMessage:
+    def test_get_message__existing_message__returns_deserialized_message(
+        self, manager: db_manager.DBManager
+    ):
+        """Test getting an existing message returns the deserialized BaseMessage."""
+        original = _create_trace_message(message_id=1)
+        manager.register_message(original)
+
+        result = manager.get_message(message_id=1)
+
+        assert result is not None
+        assert isinstance(result, messages.CreateTraceMessage)
+        assert result.trace_id == "trace-1"
+        assert result.project_name == "test-project"
+
+    def test_get_message__nonexistent_message__returns_none(
+        self, manager: db_manager.DBManager
+    ):
+        """Test getting a nonexistent message returns None."""
+        result = manager.get_message(message_id=999)
+
+        assert result is None
+
+
+class TestFetchFailedMessagesBatched:
+    def test_fetch_failed_messages_batched__no_failed_messages__yields_nothing(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that no batches are yielded when there are no failed messages."""
+        message = _create_trace_message(message_id=1)
+        manager.register_message(message)  # Status is 'registered', not 'failed'
+
+        batches = list(manager.fetch_failed_messages_batched(batch_size=10))
+
+        assert len(batches) == 0
+
+    def test_fetch_failed_messages_batched__fewer_than_batch_size__yields_single_batch(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that a single batch is yielded when messages < batch_size."""
+        for i in range(5):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        batches = list(manager.fetch_failed_messages_batched(batch_size=10))
+
+        assert len(batches) == 1
+        assert len(batches[0]) == 5
+
+    def test_fetch_failed_messages_batched__more_than_batch_size__yields_multiple_batches(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that multiple batches are yielded when messages > batch_size."""
+        for i in range(25):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        batches = list(manager.fetch_failed_messages_batched(batch_size=10))
+
+        assert len(batches) == 3
+        assert len(batches[0]) == 10
+        assert len(batches[1]) == 10
+        assert len(batches[2]) == 5
+
+    def test_fetch_failed_messages_batched__exact_batch_size__yields_correct_batches(
+        self, manager: db_manager.DBManager
+    ):
+        """Test batching when message count is exact multiple of batch_size."""
+        for i in range(20):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        batches = list(manager.fetch_failed_messages_batched(batch_size=10))
+
+        assert len(batches) == 2
+        assert len(batches[0]) == 10
+        assert len(batches[1]) == 10
+
+    def test_fetch_failed_messages_batched__reverse_insert_order__returns_ordered_by_id(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that batches return messages ordered by message_id."""
+        # Insert in reverse order
+        for i in range(10, 0, -1):
+            msg = _create_trace_message(message_id=i, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        batches = list(manager.fetch_failed_messages_batched(batch_size=5))
+
+        # Should be ordered by message_id ascending
+        assert [m.id for m in batches[0]] == [1, 2, 3, 4, 5]
+        assert [m.id for m in batches[1]] == [6, 7, 8, 9, 10]
+
+
+class TestReplayFailedMessages:
+    def test_replay_failed_messages__no_failed_messages__returns_zero(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that replay returns 0 when there are no failed messages."""
+        callback = mock.Mock()
+
+        result = manager.replay_failed_messages(callback)
+
+        assert result == 0
+        callback.assert_not_called()
+
+    def test_replay_failed_messages__with_failed_messages__replays_all(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that all failed messages are replayed."""
+        for i in range(3):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        callback = mock.Mock()
+        result = manager.replay_failed_messages(callback)
+
+        assert result == 3
+        assert callback.call_count == 3
+
+    def test_replay_failed_messages__after_replay__status_updated_to_registered(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that replayed messages have their status updated to registered."""
+        message_id = 1
+        msg = _create_trace_message(message_id=message_id)
+        manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        callback = mock.Mock()
+        manager.replay_failed_messages(callback)
+
+        # Check status was updated to registered
+        cursor = manager.conn.execute(
+            "SELECT status FROM messages WHERE message_id = ?", (message_id,)
+        )
+        assert cursor.fetchone()[0] == db_manager.MessageStatus.registered
+
+    def test_replay_failed_messages__callback_invocation__receives_deserialized_messages(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that callback receives properly deserialized BaseMessage objects."""
+        original = _create_trace_message(message_id=1)
+        manager.register_message(original, status=db_manager.MessageStatus.failed)
+
+        received_messages = []
+
+        def callback(msg: messages.BaseMessage) -> None:
+            received_messages.append(msg)
+
+        manager.replay_failed_messages(callback)
+
+        assert len(received_messages) == 1
+        assert isinstance(received_messages[0], messages.CreateTraceMessage)
+        assert received_messages[0].trace_id == "trace-1"
+
+    def test_replay_failed_messages__manager_closed__returns_zero(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that replay returns 0 when the manager is closed."""
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg, status=db_manager.MessageStatus.failed)
+        manager.close()
+
+        callback = mock.Mock()
+        result = manager.replay_failed_messages(callback)
+
+        assert result == 0
+        callback.assert_not_called()
+
+    def test_replay_failed_messages__large_batch__processes_all_batches(
+        self, small_batch_manager: db_manager.DBManager
+    ):
+        """Test that batched processing handles all messages across batches."""
+        count = 150
+        for i in range(count):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            small_batch_manager.register_message(
+                msg, status=db_manager.MessageStatus.failed
+            )
+
+        callback = mock.Mock()
+        small_batch_manager.batch_replay_delay = (
+            0.01  # Short delay to process all batches faster for testing
+        )
+        result = small_batch_manager.replay_failed_messages(callback)
+
+        assert result == count
+        assert callback.call_count == count
+
+    def test_replay_failed_messages__callback_raises_exception__continues_processing(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that exceptions in callback don't stop processing other messages."""
+        for i in range(3):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        call_count = 0
+
+        def callback(msg: messages.BaseMessage) -> None:
+            nonlocal call_count
+            call_count += 1
+            if call_count == 2:
+                raise ValueError("Test error")
+
+        result = manager.replay_failed_messages(callback)
+
+        # Should return 2 because one callback failed
+        assert result == 2
+        assert call_count == 3
+
+    def test_replay_failed_messages__lock_not_held_during_sleep__producer_not_blocked(
+        self,
+    ):
+        """Lock must be released before time.sleep so producers are never blocked.
+
+        With batch_size=1 and two failed messages there is exactly one inter-batch
+        sleep.  A producer thread that calls register_message while that sleep is
+        in progress must complete in a fraction of the sleep duration — proving the
+        lock is not held during time.sleep().
+        """
+        long_delay = 0.5
+        mgr = db_manager.DBManager(batch_size=1, batch_replay_delay=long_delay)
+
+        # Two failed messages → two batches → one sleep after the first batch.
+        msg1 = _create_trace_message(message_id=1, trace_id="trace-1")
+        msg2 = _create_trace_message(message_id=2, trace_id="trace-2")
+        mgr.register_message(msg1, status=db_manager.MessageStatus.failed)
+        mgr.register_message(msg2, status=db_manager.MessageStatus.failed)
+
+        sleep_entered = threading.Event()
+        producer_elapsed: list = []
+
+        real_sleep = time.sleep
+
+        def patched_sleep(seconds: float) -> None:
+            sleep_entered.set()
+            real_sleep(seconds)
+
+        def producer() -> None:
+            # Wait until the replay thread is inside time.sleep, then race it.
+            sleep_entered.wait(timeout=5.0)
+            t0 = time.monotonic()
+            msg3 = _create_trace_message(message_id=3, trace_id="trace-3")
+            mgr.register_message(msg3)
+            producer_elapsed.append(time.monotonic() - t0)
+
+        producer_thread = threading.Thread(target=producer, daemon=True)
+        producer_thread.start()
+
+        callback = mock.Mock()
+
+        with mock.patch(
+            "opik.message_processing.replay.db_manager.time.sleep",
+            side_effect=patched_sleep,
+        ):
+            mgr.replay_failed_messages(callback)
+
+        producer_thread.join(timeout=long_delay * 3)
+        mgr.close()
+
+        assert len(producer_elapsed) == 1, "Producer thread did not complete"
+        # If the lock were held during sleep, elapsed ≈ long_delay (0.5 s).
+        # With the lock released, elapsed should be well under 50 ms.
+        assert producer_elapsed[0] < long_delay * 0.5, (
+            f"register_message blocked for {producer_elapsed[0]:.3f}s — "
+            f"lock was likely held during sleep (delay={long_delay}s)"
+        )
+
+        # the callback should be called twice, once for each message
+        expected_calls = [mock.call(msg1), mock.call(msg2)]
+        callback.assert_has_calls(expected_calls, any_order=False)
+
+    def test_replay_failed_messages__concurrent_calls__second_returns_zero_no_duplicates(
+        self,
+    ):
+        """Concurrent replay_failed_messages calls must not re-fetch the same messages.
+
+        The _replay_mutex ensures the second caller returns 0 immediately while
+        the first is still running, so the same failed messages are never replayed
+        twice simultaneously.
+        """
+        mgr = db_manager.DBManager(batch_size=10, batch_replay_delay=0.0)
+
+        for i in range(3):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            mgr.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        first_started = threading.Event()
+        first_may_finish = threading.Event()
+        second_result: list = []
+
+        def slow_callback(m: messages.BaseMessage) -> None:
+            # Signal that the first replay is inside the callback, then wait.
+            first_started.set()
+            first_may_finish.wait(timeout=5.0)
+
+        def run_second_replay() -> None:
+            second_result.append(mgr.replay_failed_messages(mock.Mock()))
+
+        second_thread = threading.Thread(target=run_second_replay, daemon=True)
+
+        # Start the first replay (will block inside slow_callback).
+        first_thread = threading.Thread(
+            target=mgr.replay_failed_messages, args=(slow_callback,), daemon=True
+        )
+        first_thread.start()
+
+        # Wait until the first replay is inside the callback, then launch the second.
+        first_started.wait(timeout=5.0)
+        second_thread.start()
+        second_thread.join(timeout=2.0)
+
+        # The second call must have returned 0 (mutex not acquired).
+        assert second_result == [0], (
+            f"Expected second replay to return 0, got {second_result}"
+        )
+
+        # Let the first replay finish.
+        first_may_finish.set()
+        first_thread.join(timeout=5.0)
+        mgr.close()
+
+    def test_replay_failed_messages__lock_not_held_during_callback__producer_not_blocked(
+        self,
+    ):
+        """Lock must be released before the replay callback so producers are not blocked.
+
+        A producer thread that calls register_message while the replay callback is
+        executing must complete quickly — proving the lock is not held during
+        _replay_messages / replay_callback.
+        """
+        long_delay = 0.0  # no sleep; focus purely on the callback window
+        mgr = db_manager.DBManager(batch_size=10, batch_replay_delay=long_delay)
+
+        msg = _create_trace_message(message_id=1, trace_id="trace-1")
+        mgr.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        callback_entered = threading.Event()
+        callback_may_exit = threading.Event()
+        producer_elapsed: list = []
+
+        def slow_callback(m: messages.BaseMessage) -> None:
+            callback_entered.set()
+            # Hold the callback open so the producer has time to race it.
+            callback_may_exit.wait(timeout=5.0)
+
+        def producer() -> None:
+            callback_entered.wait(timeout=5.0)
+            t0 = time.monotonic()
+            msg2 = _create_trace_message(message_id=2, trace_id="trace-2")
+            mgr.register_message(msg2)
+            producer_elapsed.append(time.monotonic() - t0)
+            callback_may_exit.set()
+
+        producer_thread = threading.Thread(target=producer, daemon=True)
+        producer_thread.start()
+
+        mgr.replay_failed_messages(slow_callback)
+        producer_thread.join(timeout=2.0)
+        mgr.close()
+
+        assert len(producer_elapsed) == 1, "Producer thread did not complete"
+        # register_message should return almost immediately (well under 100 ms).
+        assert producer_elapsed[0] < 0.1, (
+            f"register_message blocked for {producer_elapsed[0]:.3f}s — "
+            f"lock was likely held during callback"
+        )
+
+
+class TestReplayManagerStatusProperties:
+    def test_closed__not_closed__returns_false(self, manager: db_manager.DBManager):
+        """Test closed property returns False when not closed."""
+        assert not manager.closed
+
+    def test_closed__after_close__returns_true(self, manager: db_manager.DBManager):
+        """Test closed property returns True after close."""
+        manager.close()
+        assert manager.closed
+
+    def test_initialized__after_init__returns_true(self, manager: db_manager.DBManager):
+        """Test initialized property returns True after successful init."""
+        assert manager.initialized
+
+    def test_failed__normal_operation__returns_false(
+        self, manager: db_manager.DBManager
+    ):
+        """Test failed property returns False during normal operation."""
+        assert not manager.failed
+
+    def test_failed__after_db_error__returns_true(self):
+        """Test failed property returns True after a DB error."""
+        # Use a closed connection to force init failure
+        conn = sqlite3.connect(":memory:", check_same_thread=False)
+        conn.close()
+        mgr = db_manager.DBManager(conn=conn, batch_size=10, batch_replay_delay=0.1)
+
+        try:
+            assert mgr.failed
+        finally:
+            mgr.close()
+
+
+class TestFailedMessagesCount:
+    def test_failed_messages_count__no_failed_messages__returns_zero(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that count returns 0 when there are no failed messages."""
+        message = _create_trace_message(message_id=1)
+        manager.register_message(message)  # Status is 'registered', not 'failed'
+
+        result = manager.failed_messages_count()
+
+        assert result == 0
+
+    def test_failed_messages_count__with_failed_messages__returns_correct_count(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that count returns the correct number of failed messages."""
+        for i in range(5):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        result = manager.failed_messages_count()
+
+        assert result == 5
+
+    def test_failed_messages_count__mixed_statuses__counts_only_failed(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that only messages with 'failed' status are counted."""
+        # Register some as 'registered' and some as 'failed'
+        for i in range(3):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg)  # registered status
+
+        for i in range(3, 5):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        result = manager.failed_messages_count()
+
+        assert result == 2
+
+    def test_failed_messages_count__manager_not_initialized__returns_negative_one(self):
+        """Test that count returns -1 when the manager is not initialized."""
+        conn = sqlite3.connect(":memory:", check_same_thread=False)
+        conn.close()
+        mgr = db_manager.DBManager(conn=conn, batch_size=10, batch_replay_delay=0.1)
+
+        result = mgr.failed_messages_count()
+
+        assert result == -1
+
+    def test_failed_messages_count__manager_closed__returns_negative_one(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that count returns -1 when the manager is closed."""
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg, status=db_manager.MessageStatus.failed)
+        manager.close()
+
+        result = manager.failed_messages_count()
+
+        assert result == -1
+
+    def test_failed_messages_count__db_error__returns_negative_one_and_marks_failed(
+        self, manager: db_manager.DBManager
+    ):
+        """Test that a DB error returns -1 and marks the manager as failed."""
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg, status=db_manager.MessageStatus.failed)
+
+        # Force a DB error by dropping the messages table
+        manager.conn.execute("DROP TABLE messages")
+
+        result = manager.failed_messages_count()
+
+        assert result == -1
+        assert manager.failed
+
+
+class TestDbMessageToMessage:
+    def test_db_message_to_message__supported_type__returns_correct_message(
+        self, manager: db_manager.DBManager
+    ):
+        """Test conversion of DBMessage to BaseMessage for supported types."""
+        original = _create_trace_message(message_id=1)
+        manager.register_message(original)
+
+        db_message = manager.get_db_message(1)
+        result = db_manager.db_message_to_message(db_message)
+
+        assert isinstance(result, messages.CreateTraceMessage)
+        assert result.trace_id == "trace-1"
+
+    def test_db_message_to_message__unsupported_type__raises_value_error(self):
+        """Test that an unsupported message type raises ValueError."""
+        db_message = db_manager.DBMessage(
+            id=1,
+            type="UnsupportedMessageType",
+            json="{}",
+            status=db_manager.MessageStatus.registered,
+        )
+
+        with pytest.raises(ValueError, match="Unsupported message type"):
+            db_manager.db_message_to_message(db_message)

sdks/python/tests/unit/message_processing/replay/test_replay_manager.py

@@ -0,0 +1,571 @@
+import datetime
+import threading
+import time
+from typing import Generator, List, Tuple, Optional
+from unittest import mock
+
+import pytest
+
+from opik.healthcheck import connection_monitor
+from opik.message_processing import messages
+from opik.message_processing.replay import db_manager, replay_manager
+
+
+def _create_trace_message(
+    message_id: Optional[int], trace_id: str = "trace-1"
+) -> messages.CreateTraceMessage:
+    msg = messages.CreateTraceMessage(
+        trace_id=trace_id,
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
+        input={"query": "test"},
+        output={"answer": "response"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _make_monitor(
+    tick_return: connection_monitor.ConnectionStatus = connection_monitor.ConnectionStatus.connection_ok,
+) -> mock.MagicMock:
+    """Create a mock OpikConnectionMonitor with configurable tick() return."""
+    monitor = mock.MagicMock(spec=connection_monitor.OpikConnectionMonitor)
+    monitor.tick.return_value = tick_return
+    monitor.has_server_connection = True
+    return monitor
+
+
+def _make_manager(
+    monitor: mock.MagicMock,
+    tick_interval: float = 0.05,
+) -> replay_manager.ReplayManager:
+    """Create a ReplayManager with a fast tick for testing."""
+    return replay_manager.ReplayManager(
+        monitor=monitor,
+        batch_size=10,
+        batch_replay_delay=0.01,
+        tick_interval_seconds=tick_interval,
+    )
+
+
+@pytest.fixture
+def monitor() -> mock.MagicMock:
+    return _make_monitor()
+
+
+@pytest.fixture
+def manager(
+    monitor: mock.MagicMock,
+) -> Generator[replay_manager.ReplayManager, None, None]:
+    rm = _make_manager(monitor)
+    yield rm
+    rm.close()
+    if rm.is_alive():
+        rm.join(timeout=2)
+
+
+@pytest.fixture
+def manager_monitor(
+    monitor: mock.MagicMock,
+) -> Generator[
+    Tuple[replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor],
+    None,
+    None,
+]:
+    rm = _make_manager(monitor, tick_interval=0.05)
+    yield rm, monitor
+    rm.close()
+    if rm.is_alive():
+        rm.join(timeout=2)
+
+
+class TestReplayManagerInitialization:
+    def test_init__default_parameters__creates_thread(
+        self, manager: replay_manager.ReplayManager
+    ):
+        assert manager.daemon is True
+        assert manager.name == "ReplayManager"
+        assert not manager.is_alive()
+
+    def test_init_db_manager_initialized(self, manager: replay_manager.ReplayManager):
+        assert manager.database_manager.initialized
+        assert not manager.database_manager.closed
+
+
+class TestStart:
+    def test_start__without_callback__raises_value_error(
+        self, manager: replay_manager.ReplayManager
+    ):
+        with pytest.raises(ValueError, match="Replay callback must be set"):
+            manager.start()
+
+    def test_start__with_callback__thread_starts(
+        self, manager: replay_manager.ReplayManager
+    ):
+        manager.set_replay_callback(mock.Mock())
+        manager.start()
+
+        assert manager.is_alive()
+
+        manager.close()
+        manager.join(timeout=2)
+        assert not manager.is_alive()
+
+
+class TestClose:
+    def test_close__running_thread__stops_loop(
+        self, manager: replay_manager.ReplayManager
+    ):
+        manager.set_replay_callback(mock.Mock())
+        manager.start()
+        assert manager.is_alive()
+
+        manager.close()
+        manager.join(timeout=2)
+
+        assert not manager.is_alive()
+
+    def test_close_db_manager_closed_after_thread_exits(
+        self, manager: replay_manager.ReplayManager
+    ):
+        manager.set_replay_callback(mock.Mock())
+        manager.start()
+
+        manager.close()
+        manager.join(timeout=2)
+
+        assert manager.database_manager.closed
+
+    def test_close__interruptible_sleep(self, monitor: mock.MagicMock):
+        """close() should interrupt the sleep and stop quickly, not wait full tick."""
+        rm = _make_manager(monitor, tick_interval=5.0)
+        rm.set_replay_callback(mock.Mock())
+        rm.start()
+
+        start_time = time.time()
+        rm.close()
+        rm.join(timeout=2)
+        elapsed = time.time() - start_time
+
+        assert elapsed < 2.0, f"close() took {elapsed:.2f}s, should be near-instant"
+
+
+class TestHasServerConnection:
+    def test_has_server_connection__delegates_to_monitor(
+        self, manager: replay_manager.ReplayManager, monitor: mock.MagicMock
+    ):
+        monitor.has_server_connection = True
+        assert manager.has_server_connection is True
+
+        monitor.has_server_connection = False
+        assert manager.has_server_connection is False
+
+
+class TestRegisterMessage:
+    def test_register_message__stores_in_db(
+        self, manager: replay_manager.ReplayManager
+    ):
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg)
+
+        db_msg = manager.database_manager.get_db_message(1)
+        assert db_msg is not None
+        assert db_msg.id == 1
+        assert db_msg.status == db_manager.MessageStatus.registered
+
+    def test_register_message__automatically_set_message_id(
+        self, manager: replay_manager.ReplayManager
+    ):
+        msg = _create_trace_message(message_id=None)
+        manager.register_message(
+            msg
+        )  # this will set msg.message_id to a particular value
+
+        assert msg.message_id is not None
+
+        db_msg = manager.database_manager.get_db_message(msg.message_id)
+        assert db_msg is not None
+        assert db_msg.id == msg.message_id
+        assert db_msg.status == db_manager.MessageStatus.registered
+
+
+class TestUnregisterMessage:
+    def test_unregister_message__deletes_from_db(
+        self, manager: replay_manager.ReplayManager
+    ):
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg)
+
+        manager.unregister_message(1)
+
+        db_msg = manager.database_manager.get_db_message(1)
+        assert db_msg is None
+
+
+class TestMessageSentFailed:
+    def test_message_sent_failed__marks_message_as_failed(
+        self, manager: replay_manager.ReplayManager
+    ):
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg)
+
+        manager.message_sent_failed(1, failure_reason="timeout")
+
+        db_msg = manager.database_manager.get_db_message(1)
+        assert db_msg is not None
+        assert db_msg.status == db_manager.MessageStatus.failed
+
+    def test_message_sent_failed__notifies_monitor(
+        self, manager: replay_manager.ReplayManager, monitor: mock.MagicMock
+    ):
+        msg = _create_trace_message(message_id=1)
+        manager.register_message(msg)
+
+        manager.message_sent_failed(1, failure_reason="timeout")
+
+        monitor.connection_failed.assert_called_once_with(failure_reason="timeout")
+
+
+class TestFlush:
+    def test_flush__without_callback__raises_value_error(
+        self, manager: replay_manager.ReplayManager
+    ):
+        with pytest.raises(ValueError, match="Replay callback must be set"):
+            manager.flush()
+
+    def test_flush__no_failed_messages__callback_not_called(
+        self, manager: replay_manager.ReplayManager
+    ):
+        callback = mock.Mock()
+        manager.set_replay_callback(callback)
+
+        manager.flush()
+
+        callback.assert_not_called()
+
+    def test_flush__with_failed_messages__replays_all(
+        self, manager: replay_manager.ReplayManager
+    ):
+        for i in range(3):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            manager.database_manager.register_message(
+                msg, status=db_manager.MessageStatus.failed
+            )
+
+        callback = mock.Mock()
+        manager.set_replay_callback(callback)
+
+        manager.flush()
+
+        assert callback.call_count == 3
+        for call_args in callback.call_args_list:
+            replayed_msg = call_args[0][0]
+            assert isinstance(replayed_msg, messages.CreateTraceMessage)
+
+    def test_flush__callback_failure__re_marks_message_as_failed(
+        self, manager: replay_manager.ReplayManager
+    ):
+        msg = _create_trace_message(message_id=1)
+        manager.database_manager.register_message(
+            msg, status=db_manager.MessageStatus.failed
+        )
+
+        callback = mock.Mock(side_effect=RuntimeError("send error"))
+        manager.set_replay_callback(callback)
+
+        manager.flush()
+
+        db_msg = manager.database_manager.get_db_message(1)
+        assert db_msg is not None
+        assert db_msg.status == db_manager.MessageStatus.failed
+
+
+class TestLoopConnectionRestored:
+    def test_loop__connection_restored__replays_failed_messages(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        rm, monitor = manager_monitor
+        monitor.tick.return_value = (
+            connection_monitor.ConnectionStatus.connection_restored
+        )
+
+        # Register a failed message
+        msg = _create_trace_message(message_id=1)
+        rm.database_manager.register_message(
+            msg, status=db_manager.MessageStatus.failed
+        )
+
+        replayed: List[messages.BaseMessage] = []
+
+        def callback(m: messages.BaseMessage) -> None:
+            replayed.append(m)
+
+        rm.set_replay_callback(callback)
+        rm.start()
+
+        # Wait for at least one tick cycle
+        deadline = time.time() + 0.5
+        while not replayed and time.time() < deadline:
+            time.sleep(0.05)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        assert len(replayed) == 1
+        assert isinstance(replayed[0], messages.CreateTraceMessage)
+        assert replayed[0].trace_id == "trace-1"
+
+    def test_loop__connection_restored__resets_monitor(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        rm, monitor = manager_monitor
+        monitor.tick.return_value = (
+            connection_monitor.ConnectionStatus.connection_restored
+        )
+
+        rm.set_replay_callback(mock.Mock())
+        rm.start()
+
+        # Wait for at least one tick
+        deadline = time.time() + 2.0
+        while monitor.reset.call_count == 0 and time.time() < deadline:
+            time.sleep(0.05)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        monitor.reset.assert_called()
+
+    def test_loop__connection_ok__no_replay(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        rm, monitor = manager_monitor
+        monitor.tick.return_value = connection_monitor.ConnectionStatus.connection_ok
+
+        msg = _create_trace_message(message_id=1)
+        rm.database_manager.register_message(
+            msg, status=db_manager.MessageStatus.failed
+        )
+
+        callback = mock.Mock()
+        rm.set_replay_callback(callback)
+        rm.start()
+
+        # Let several ticks pass
+        time.sleep(0.1)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        callback.assert_not_called()
+
+    def test_loop__connection_failed__no_replay(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        rm, monitor = manager_monitor
+        monitor.tick.return_value = (
+            connection_monitor.ConnectionStatus.connection_failed
+        )
+
+        msg = _create_trace_message(message_id=1)
+        rm.database_manager.register_message(
+            msg, status=db_manager.MessageStatus.failed
+        )
+
+        callback = mock.Mock()
+        rm.set_replay_callback(callback)
+        rm.start()
+
+        time.sleep(0.1)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        callback.assert_not_called()
+
+
+class TestLoopExceptionHandling:
+    def test_loop__tick_raises__thread_continues(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        """The loop should survive exceptions from monitor.tick() and keep running."""
+        call_count = 0
+
+        def tick_side_effect():
+            nonlocal call_count
+            call_count += 1
+            if call_count <= 2:
+                raise ConnectionError("probe failed")
+            return connection_monitor.ConnectionStatus.connection_ok
+
+        rm, monitor = manager_monitor
+        monitor.tick.side_effect = tick_side_effect
+
+        rm.set_replay_callback(mock.Mock())
+        rm.start()
+
+        # Wait for at least 3 ticks
+        deadline = time.time() + 1.0
+        while call_count < 3 and time.time() < deadline:
+            time.sleep(0.05)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        assert call_count >= 3, "Thread should have survived the first two exceptions"
+
+    def test_loop__replay_callback_raises__thread_continues(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        """If replay_failed_messages raises, the loop should continue."""
+        rm, monitor = manager_monitor
+
+        msg = _create_trace_message(message_id=1)
+        rm.database_manager.register_message(
+            msg, status=db_manager.MessageStatus.failed
+        )
+        # mark connection as restored to enable replay callback
+        monitor.tick.return_value = (
+            connection_monitor.ConnectionStatus.connection_restored
+        )
+
+        call_count = 0
+
+        def failing_callback(m: messages.BaseMessage) -> None:
+            nonlocal call_count
+            call_count += 1
+            raise RuntimeError("replay error")
+
+        rm.set_replay_callback(failing_callback)
+        rm.start()
+
+        # Wait for a few ticks
+        deadline = time.time() + 2.0
+        while call_count < 2 and time.time() < deadline:
+            time.sleep(0.05)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        assert call_count >= 2, "Loop should have continued despite replay errors"
+
+        assert rm.is_alive() is False, "Thread should have exited cleanly"
+
+
+class TestFlushConcurrency:
+    def test_flush_and_loop__serialized_by_replay_lock(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        """flush() and the loop both acquire _replay_lock, so only one replay
+        can happen at a time."""
+        rm, monitor = manager_monitor
+
+        for i in range(5):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            rm.database_manager.register_message(
+                msg, status=db_manager.MessageStatus.failed
+            )
+
+        replayed_ids: List[str] = []
+        lock = threading.Lock()
+
+        def callback(m: messages.BaseMessage) -> None:
+            with lock:
+                assert isinstance(m, messages.CreateTraceMessage)
+                replayed_ids.append(m.trace_id)
+
+        rm.set_replay_callback(callback)
+        rm.start()
+
+        # Also call flush from the main thread
+        rm.flush()
+
+        # Let the loop run too
+        time.sleep(0.3)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        # Each message should appear at least once (no corruption)
+        assert len(replayed_ids) == 5
+
+
+class TestEndToEnd:
+    def test_full_lifecycle__register_fail_restore_replay(
+        self,
+        manager_monitor: Tuple[
+            replay_manager.ReplayManager, connection_monitor.OpikConnectionMonitor
+        ],
+    ):
+        """Test the complete message lifecycle:
+        register → fail → connection restored → replay → delivered."""
+        rm, monitor = manager_monitor
+
+        replayed: List[messages.BaseMessage] = []
+
+        def callback(m: messages.BaseMessage) -> None:
+            replayed.append(m)
+
+        rm.set_replay_callback(callback)
+
+        # 1. Register a message
+        msg = _create_trace_message(message_id=1)
+        rm.register_message(msg)
+
+        # 2. Mark it as failed (simulating a connection error during sending)
+        rm.message_sent_failed(1, failure_reason="connection timeout")
+        db_msg = rm.database_manager.get_db_message(1)
+        assert db_msg is not None
+        assert db_msg.status == db_manager.MessageStatus.failed
+
+        # 3. Start the thread — connection is still failed, no replay
+        rm.start()
+        time.sleep(0.2)
+        assert len(replayed) == 0
+
+        # 4. Simulate connection restored
+        monitor.tick.return_value = (
+            connection_monitor.ConnectionStatus.connection_restored
+        )
+
+        # 5. Wait for replay
+        deadline = time.time() + 2.0
+        while not replayed and time.time() < deadline:
+            time.sleep(0.05)
+
+        rm.close()
+        rm.join(timeout=2)
+
+        assert len(replayed) == 1
+        assert isinstance(replayed[0], messages.CreateTraceMessage)
+        assert replayed[0].trace_id == "trace-1"
+        monitor.reset.assert_called()

sdks/python/tests/unit/message_processing/test_dynamic_rate_limiting.py

@@ -0,0 +1,129 @@
+import time
+from unittest import mock
+
+import pytest
+
+from opik import exceptions, synchronization
+from opik.message_processing import streamer_constructors, messages, queue_consumer
+
+MAX_QUEUE_SIZE = 10
+
+
+@pytest.fixture
+def streamer_with_mock_message_processor(fake_file_upload_manager):
+    mock_message_processor = mock.Mock()
+    streamer = streamer_constructors.construct_streamer(
+        message_processor=mock_message_processor,
+        n_consumers=1,
+        use_batching=False,
+        use_attachment_extraction=False,
+        file_uploader=fake_file_upload_manager,
+        max_queue_size=MAX_QUEUE_SIZE,
+        fallback_replay_manager=mock.Mock(),
+    )
+
+    yield streamer, mock_message_processor
+
+    streamer.close(flush=False)
+
+
+def test_dynamic_rate_limiting__rate_limited__check_queue_messages_are_put_back(
+    streamer_with_mock_message_processor,
+):
+    streamer, mock_message_processor = streamer_with_mock_message_processor
+
+    retry_after = queue_consumer.SLEEP_BETWEEN_LOOP_ITERATIONS * 3
+    mock_message_processor.process.side_effect = (
+        exceptions.OpikCloudRequestsRateLimited(
+            headers={},
+            retry_after=retry_after,
+        )
+    )
+
+    now = time.monotonic()
+    messages_number = 5
+    for i in range(messages_number):
+        streamer.put(messages.BaseMessage())
+
+    # Don't pin the assertion to a fixed sleep matching `retry_after` — that
+    # races with the consumer waking from its rate-limit wait at the same
+    # instant. Poll the actual signals instead: the rate-limited message has
+    # been pushed back (queue_size is at the original count) and the consumer
+    # has set a future `next_message_time` (so it's parked in the wait branch
+    # and won't pop another message).
+    synchronization.until(
+        lambda: (
+            streamer.queue_size() == messages_number
+            and streamer._queue_consumers[0].next_message_time > now
+        ),
+        max_try_seconds=1.0,
+    )
+    assert streamer.queue_size() == messages_number
+    assert streamer._queue_consumers[0].next_message_time > now
+
+    # Same idea for the drain: poll for queue_size == 0 instead of a fixed
+    # sleep tied to retry_after.
+    mock_message_processor.process = lambda message: None
+    synchronization.until(
+        lambda: streamer.queue_size() == 0,
+        max_try_seconds=retry_after * 4,
+    )
+    assert streamer.queue_size() == 0
+
+
+def test_dynamic_rate_limiting__rate_limited__check_queue_size_is_bounded_by_max_queue_size(
+    streamer_with_mock_message_processor,
+):
+    streamer, mock_message_processor = streamer_with_mock_message_processor
+
+    # to allow a few skipped loop iterations
+    retry_after = queue_consumer.SLEEP_BETWEEN_LOOP_ITERATIONS * 3
+    mock_message_processor.process.side_effect = (
+        exceptions.OpikCloudRequestsRateLimited(
+            headers={},
+            retry_after=retry_after,
+        )
+    )
+
+    messages_number = 2 * MAX_QUEUE_SIZE
+    for i in range(messages_number):
+        streamer.put(messages.BaseMessage())
+        time.sleep(queue_consumer.SLEEP_BETWEEN_LOOP_ITERATIONS)
+
+    # sleep for a while to allow queue_consumer to execute a few loop iterations
+    time.sleep(retry_after)
+    # check that queue size doesn't exceed the maximal allowed size
+    assert streamer.queue_size() == MAX_QUEUE_SIZE
+
+    # check that all queue messages are processed when no exception is raised
+    mock_message_processor.process = lambda message: None
+    time.sleep(retry_after * 2)
+    assert streamer.queue_size() == 0
+
+
+def test_dynamic_rate_limiting__rate_limited__check_streamer_flush_is_blocking_until_all_messages_are_processed(
+    streamer_with_mock_message_processor,
+):
+    streamer, mock_message_processor = streamer_with_mock_message_processor
+    retry_after = queue_consumer.SLEEP_BETWEEN_LOOP_ITERATIONS * 3
+    max_limit = 5
+
+    counter = {"count": 0}
+
+    def process(message: messages.BaseMessage):
+        counter["count"] += 1
+        if counter["count"] == max_limit:
+            raise exceptions.OpikCloudRequestsRateLimited(
+                headers={},
+                retry_after=retry_after,
+            )
+
+    mock_message_processor.process = process
+
+    for i in range(max_limit):
+        streamer.put(messages.BaseMessage())
+
+    streamer.flush(timeout=None)
+
+    # we expect max_limit calls before retry plus one call after waiting for the retry
+    assert counter["count"] == max_limit + 1

sdks/python/tests/unit/message_processing/test_encoder_helpers.py

@@ -0,0 +1,410 @@
+from typing import Dict, Optional, Any
+
+from opik.anonymizer import anonymizer, recursive_anonymizer
+
+from opik.message_processing import encoder_helpers
+
+
+class MockAnonymizer(anonymizer.Anonymizer):
+    """Mock anonymizer for testing purposes."""
+
+    def anonymize(self, data, **kwargs):
+        """Mock anonymization that replaces strings with '[ANONYMIZED]'."""
+        if isinstance(data, str):
+            return "[ANONYMIZED]"
+        return data
+
+
+class TestEncodeAndAnonymize:
+    """Test suite for anonymize_encoded_obj functionality."""
+
+    def test_anonymize_encoded_obj__no_anonymizers__returns_encoded_only(self):
+        """Test that with an empty anonymizers list, only encoding is performed."""
+        obj = {"name": "John Doe", "email": "john@example.com"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj, anonymizers=[], fields_to_anonymize=set(), object_type="span"
+        )
+
+        expected = {"name": "John Doe", "email": "john@example.com"}
+        assert result == expected
+
+    def test_anonymize_encoded_obj__with_anonymizers_no_fields__no_error(self):
+        """Test that providing anonymizers with empty fields works."""
+        obj = {"name": "John Doe"}
+        mock_anonymizer = MockAnonymizer()
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=set(),
+            object_type="span",
+        )
+
+        expected = {"name": "John Doe"}
+        assert result == expected
+
+    def test_anonymize_encoded_obj__dict_with_matching_fields(self):
+        """Test anonymization of a dictionary with matching field names."""
+        obj = {
+            "name": "John Doe",
+            "email": "john@example.com",
+            "phone": "123-456-7890",
+            "age": 30,
+        }
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email", "phone"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        expected = {
+            "name": "John Doe",
+            "email": "[ANONYMIZED]",
+            "phone": "[ANONYMIZED]",
+            "age": 30,
+        }
+        assert result == expected
+
+    def test_anonymize_encoded_obj__dict_with_no_matching_fields(self):
+        """Test that fields not in dict are ignored."""
+        obj = {"name": "John Doe", "age": 30}
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email", "phone"}  # These fields don't exist in obj
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        expected = {"name": "John Doe", "age": 30}
+        assert result == expected
+
+    def test_anonymize_encoded_obj__dict_partial_field_match(self):
+        """Test anonymization when only some specified fields exist."""
+        obj = {"name": "John Doe", "email": "john@example.com", "age": 30}
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email", "phone", "ssn"}  # Only email exists
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        expected = {"name": "John Doe", "email": "[ANONYMIZED]", "age": 30}
+        assert result == expected
+
+    def test_anonymize_encoded_obj__non_dict_object__no_anonymization(self):
+        """Test that non-dict objects are not anonymized."""
+        obj = ["item1", "item2", "item3"]
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"item1"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="trace",
+        )
+
+        # Should return an encoded list without anonymization
+        assert result == ["item1", "item2", "item3"]
+
+    def test_anonymize_encoded_obj__string_object__no_anonymization(self):
+        """Test that string objects are not anonymized."""
+        obj = "This is a sensitive string"
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"field1"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="trace",
+        )
+
+        assert result == "This is a sensitive string"
+
+    def test_anonymize_encoded_obj__complex_nested_object(self):
+        """Test encoding complex nested objects before anonymization."""
+        import dataclasses
+        from opik import jsonable_encoder
+
+        @dataclasses.dataclass
+        class Person:
+            name: str
+            email: str
+            age: int
+            address: Dict[str, str] = dataclasses.field(default_factory=dict)
+
+        person = Person(name="John Doe", email="john@example.com", age=30)
+        person.address["street"] = "123 Main Street"
+        person.address["city"] = "New York"
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email"}
+
+        # Encode the object first, as expected by anonymize_encoded_obj
+        encoded_person = jsonable_encoder.encode(person)
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=encoded_person,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="trace",
+        )
+
+        expected = {
+            "name": "John Doe",
+            "email": "[ANONYMIZED]",
+            "age": 30,
+            "address": {
+                "city": "New York",
+                "street": "123 Main Street",
+            },
+        }
+        assert result == expected
+
+    def test_anonymize_encoded_obj__nested_dict_in_encoded_result(self):
+        """Test that only top-level fields are anonymized in nested structures."""
+        obj = {
+            "user_info": {"email": "nested@example.com", "name": "Nested User"},
+            "email": "top@example.com",
+            "id": "12345",
+        }
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email"}  # Only top-level email should be anonymized
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        expected = {
+            "user_info": {
+                "email": "nested@example.com",  # Not anonymized (nested)
+                "name": "Nested User",
+            },
+            "email": "[ANONYMIZED]",  # Anonymized (top-level)
+            "id": "12345",
+        }
+        assert result == expected
+
+    def test_anonymize_encoded_obj__empty_dict(self):
+        """Test handling of empty dictionary."""
+        obj = {}
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        assert result == {}
+
+    def test_anonymize_encoded_obj__empty_fields_set(self):
+        """Test with an empty fields_to_anonymize set."""
+        obj = {"name": "John", "email": "john@example.com"}
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = set()  # Empty set
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        # No fields should be anonymized
+        expected = {"name": "John", "email": "john@example.com"}
+        assert result == expected
+
+    def test_anonymize_encoded_obj__various_field_types(self):
+        """Test anonymization of fields with various data types."""
+        obj = {
+            "string_field": "test string",
+            "int_field": 42,
+            "float_field": 3.14,
+            "bool_field": True,
+            "none_field": None,
+            "list_field": [1, 2, 3],
+        }
+
+        # Create an anonymizer that just adds a prefix
+        class PrefixAnonymizer(anonymizer.Anonymizer):
+            def anonymize(self, data, **kwargs):
+                return f"ANON_{data}"
+
+        prefix_anonymizer = PrefixAnonymizer()
+        fields_to_anonymize = {"string_field", "int_field", "bool_field", "none_field"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[prefix_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        expected = {
+            "string_field": "ANON_test string",
+            "int_field": "ANON_42",
+            "float_field": 3.14,  # Not anonymized
+            "bool_field": "ANON_True",
+            "none_field": "ANON_None",
+            "list_field": [1, 2, 3],  # Not anonymized
+        }
+        assert result == expected
+
+    def test_anonymize_encoded_obj__integration_with_actual_encoder_features(self):
+        """Test integration with actual encoder features like datetime serialization."""
+        from datetime import datetime, timezone
+        from opik import jsonable_encoder
+
+        obj = {
+            "timestamp": datetime(2023, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+            "email": "test@example.com",
+            "data": {"nested": "value"},
+        }
+        mock_anonymizer = MockAnonymizer()
+        fields_to_anonymize = {"email"}
+
+        # Encode the object first, as expected by anonymize_encoded_obj
+        encoded_obj = jsonable_encoder.encode(obj)
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=encoded_obj,
+            anonymizers=[mock_anonymizer],
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        # Should have encoded datetime and anonymized email
+        assert "timestamp" in result
+        assert result["timestamp"] == "2023-01-01T12:00:00Z"  # Encoded datetime
+        assert result["email"] == "[ANONYMIZED]"  # Anonymized field
+        assert result["data"] == {"nested": "value"}  # Unchanged nested data
+
+    def test_anonymize_encoded_obj__remove_sensitive_dictionary_key(self):
+        """Test that sensitive keys can be removed from the result."""
+
+        class ApiKeyAnonymizer(anonymizer.Anonymizer):
+            def anonymize(self, data, **kwargs):
+                if "api_key" in data:
+                    del data["api_key"]
+                return data
+
+        obj = {
+            "metadata": {
+                "api_key": "12345",
+                "email": "test@example.com",
+                "data": {"nested": "value"},
+            },
+            "input": {"role": "user", "question": "What is LLM?"},
+        }
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[ApiKeyAnonymizer()],
+            fields_to_anonymize={"metadata"},
+            object_type="span",
+        )
+
+        # should remove api_key
+        assert "api_key" not in result["metadata"]
+
+    def test_anonymize_encoded_obj__field_name_passed_to_anonymizer(self):
+        """Test that sensitive field names and auxiliary information are passed to the anonymizer."""
+
+        class ApiKeyAnonymizer(anonymizer.Anonymizer):
+            def anonymize(self, data, **kwargs):
+                field_name = kwargs.get("field_name")
+                object_type = kwargs.get("object_type")
+                if (
+                    field_name == "metadata"
+                    and object_type == "span"
+                    and "api_key" in data
+                ):
+                    del data["api_key"]
+                return data
+
+        class SSNAnonymizer(recursive_anonymizer.RecursiveAnonymizer):
+            def anonymize_text(
+                self, data: str, field_name: Optional[str] = None, **kwargs: Any
+            ) -> str:
+                object_type = kwargs.get("object_type")
+                if field_name == "input.ssn" and object_type == "span":
+                    return "[SSN_REMOVED]"
+
+                return data
+
+        obj = {
+            "metadata": {
+                "api_key": "12345",
+                "email": "test@example.com",
+                "data": {"nested": "value"},
+            },
+            "input": {
+                "api_key": "12345",
+                "ssn": "123-4567-789",
+                "role": "user",
+                "question": "What is LLM?",
+            },
+        }
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=[ApiKeyAnonymizer(), SSNAnonymizer()],
+            fields_to_anonymize={"metadata", "input"},
+            object_type="span",
+        )
+
+        # should remove api_key from metadata
+        assert "api_key" not in result["metadata"]
+
+        # should not remove api_key from input
+        assert "api_key" in result["input"]
+
+        # should have SSN removed
+        assert result["input"]["ssn"] == "[SSN_REMOVED]"
+
+    def test_anonymize_encoded_obj__multiple_anonymizers(self):
+        """Test that multiple anonymizers are applied in a sequence."""
+
+        class PrefixAnonymizer(anonymizer.Anonymizer):
+            def anonymize(self, data, **kwargs):
+                if isinstance(data, str):
+                    return f"PREFIX_{data}"
+                return data
+
+        class SuffixAnonymizer(anonymizer.Anonymizer):
+            def anonymize(self, data, **kwargs):
+                if isinstance(data, str):
+                    return f"{data}_SUFFIX"
+                return data
+
+        obj = {"email": "test@example.com", "name": "John Doe"}
+        anonymizers = [PrefixAnonymizer(), SuffixAnonymizer()]
+        fields_to_anonymize = {"email"}
+
+        result = encoder_helpers.anonymize_encoded_obj(
+            obj=obj,
+            anonymizers=anonymizers,
+            fields_to_anonymize=fields_to_anonymize,
+            object_type="span",
+        )
+
+        # Should apply both anonymizers in order: first prefix, then suffix
+        expected = {"email": "PREFIX_test@example.com_SUFFIX", "name": "John Doe"}
+        assert result == expected

sdks/python/tests/unit/message_processing/test_jsonable_encoder.py

@@ -0,0 +1,197 @@
+import base64
+import dataclasses
+import decimal
+import importlib
+import uuid
+import warnings
+from datetime import date, datetime, timezone
+from threading import Lock
+from typing import Any, Optional
+
+import numpy as np
+
+import pytest
+
+import opik.jsonable_encoder as jsonable_encoder
+from opik.rest_api.core import pydantic_utilities
+
+
+@dataclasses.dataclass
+class Node:
+    value: int
+    child: Optional["Node"] = None
+
+
+def test_jsonable_encoder__cyclic_reference():
+    """
+    Test that the encoder detects cyclic references and does not infinitely recurse.
+    """
+    # Create a simple two-node cycle: A -> B -> A
+    node_a = Node(value=1)
+    node_b = Node(value=2)
+    node_a.child = node_b
+    node_b.child = node_a
+
+    encoded = jsonable_encoder.encode(node_a)
+    # The exact format of the cycle marker can vary; we check that:
+    # 1. We get some structure for node_a (like a dict).
+    # 2. Inside node_a, there's a reference to node_b (a dict).
+    # 3. Inside node_b, there's a "cyclic reference" marker instead of a full node_a object.
+    print("=" * 150)
+    print(encoded)
+    assert isinstance(encoded, dict)
+    assert "value" in encoded
+    assert "child" in encoded
+
+    # node_a.child (which is node_b) should be a dict
+    assert isinstance(encoded["child"], dict)
+    assert "value" in encoded["child"]
+    assert "child" in encoded["child"]
+
+    # node_b.child should be the cycle marker
+    cycle_marker = encoded["child"]["child"]
+    print("=" * 150)
+    print(cycle_marker)
+    assert isinstance(cycle_marker, str), (
+        "Expected a string marker for cyclic reference"
+    )
+    assert "<Cyclic reference to " in cycle_marker, (
+        "Should contain 'Cyclic reference' text"
+    )
+
+
+def test_jsonable_encoder__repeated_objects_in_list():
+    """
+    Test that the encoder handles a list of the same object repeated multiple times
+    without marking it as a cycle (because it isn't a cycle—just repeated references).
+    """
+    node = Node(value=42)
+
+    # Put the same node object in a list multiple times
+    repeated_list = [node, node, node]
+
+    encoded = jsonable_encoder.encode(repeated_list)
+    # We expect a list of three items, each being a dict with `value` = 42, `child` = None
+    assert isinstance(encoded, list)
+    assert len(encoded) == 3
+
+    for item in encoded:
+        assert isinstance(item, dict)
+        assert item.get("value") == 42
+        assert item.get("child") is None
+
+    # They are distinct dictionary objects, but there is no cycle reference marker
+    # because there's no actual cycle. It's just repeated references of the same object.
+    assert all("Cyclic reference" not in str(item) for item in encoded)
+
+
+@pytest.mark.parametrize(
+    "obj",
+    [
+        42,
+        42.000,
+        "42",
+        {
+            "p1": 42,
+            "p2": 42.000,
+            "p3": "42",
+        },
+        [
+            42,
+            42.000,
+            "42",
+        ],
+    ],
+)
+def test_jsonable_encoder__common_types(obj):
+    assert obj == jsonable_encoder.encode(obj)
+
+
+@pytest.mark.parametrize(
+    "obj",
+    [
+        {42, 42.000, "42"},
+        (42, 42.000, "42"),
+        np.array([42, 42.000, "42"]),
+    ],
+)
+def test_jsonable_encoder__converted_to_list(obj):
+    assert list(obj) == jsonable_encoder.encode(obj)
+
+
+@pytest.mark.parametrize(
+    "obj,expected",
+    [
+        (
+            date(2020, 1, 1),
+            "2020-01-01",
+        ),
+        (datetime(2020, 1, 1, 10, 20, 30, tzinfo=timezone.utc), "2020-01-01T10:20:30Z"),
+    ],
+)
+def test_jsonable_encoder__datetime_to_text(obj, expected):
+    assert expected == jsonable_encoder.encode(obj)
+
+
+def test_jsonable_encoder__non_serializable_to_text__class():
+    class SomeClass:
+        a = 1
+        b = 42.0
+
+    data = SomeClass()
+
+    assert "SomeClass object at 0x" in jsonable_encoder.encode(data)
+
+
+def test_jsonable_encoder__non_serializable_to_text__lock():
+    data = Lock()
+
+    assert jsonable_encoder.encode(data).startswith(
+        "<unlocked _thread.lock object at 0x"
+    )
+
+
+def test_jsonable_encoder__non_serializable_lock_inside_dataclass__lock_converted_to_text():
+    @dataclasses.dataclass
+    class SomeClass:
+        a: Any
+        b: Any
+
+    data = SomeClass(a=1, b=Lock())
+
+    encoded = jsonable_encoder.encode(data)
+    assert isinstance(encoded, dict)
+    assert encoded["a"] == 1
+    assert encoded["b"].startswith("<unlocked _thread.lock object at 0x")
+
+
+def test_jsonable_encoder__non_serializable_to_text__bytes():
+    data = b"deadbeef"
+
+    assert base64.b64encode(data).decode("utf-8") == jsonable_encoder.encode(data)
+
+
+def test_jsonable_encoder__uuid_and_decimal_remain_serializable():
+    payload = {
+        "id": uuid.UUID("12345678-1234-5678-1234-567812345678"),
+        "score": decimal.Decimal("1.23"),
+    }
+
+    encoded = jsonable_encoder.encode(payload)
+
+    assert encoded["id"] == "12345678-1234-5678-1234-567812345678"
+    assert encoded["score"] == "1.23"
+
+
+def test_pydantic_utilities__import_does_not_emit_v1_py314_warning():
+    pytest.importorskip("pydantic")
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        importlib.reload(pydantic_utilities)
+
+    assert not any(
+        "Core Pydantic V1 functionality isn't compatible with Python 3.14 or greater."
+        in str(w.message)
+        for w in caught
+    )

sdks/python/tests/unit/message_processing/test_message_queue.py

@@ -0,0 +1,213 @@
+import threading
+import time
+
+import pytest
+
+from opik.message_processing import message_queue
+
+
+def test_message_queue_accept_put_without_discarding___max_size_handling():
+    max_size = 10
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(max_size)
+
+    queue.put("the oldest message")
+    assert queue.accept_put_without_discarding() is True
+
+    for i in range(max_size - 1):
+        assert queue.accept_put_without_discarding() is True, f"Failed at {i}"
+        queue.put(f"line: {i}")
+
+    assert queue.accept_put_without_discarding() is False
+
+    # put one more message to have the oldest discarded
+    queue.put("the newest message")
+
+    # check that the oldest message ("the oldest message") was discarded
+    assert queue.get(0.0001) == "line: 0"
+
+
+def test_message_queue_accept_put_without_discarding___max_size_none():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(None)
+    for i in range(100):
+        queue.put(f"line: {i}")
+
+    assert queue.accept_put_without_discarding() is True
+
+
+def test_message_queue_task_done__fresh_queue__reports_all_tasks_done():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+    assert queue.all_tasks_done() is True
+
+
+def test_message_queue_task_done__put_then_task_done__quiescence_restored():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+
+    queue.put("msg-1")
+    assert queue.all_tasks_done() is False
+
+    queue.get(0.0001)
+    # Popping does not mark the task done — it's in flight until task_done().
+    assert queue.all_tasks_done() is False
+
+    queue.task_done()
+    assert queue.all_tasks_done() is True
+
+
+def test_message_queue_task_done__multiple_puts__decremented_one_by_one():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+
+    for i in range(5):
+        queue.put(f"msg-{i}")
+
+    assert queue.all_tasks_done() is False
+
+    for _ in range(4):
+        queue.get(0.0001)
+        queue.task_done()
+        assert queue.all_tasks_done() is False
+
+    queue.get(0.0001)
+    queue.task_done()
+    assert queue.all_tasks_done() is True
+
+
+def test_message_queue_task_done__put_back_does_not_increment_counter():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+
+    queue.put("msg")
+    msg = queue.get(0.0001)
+    # Re-enqueue (delivery-time deferral / rate-limit retry): the task is
+    # still in flight, so the counter must not change.
+    queue.put_back(msg)
+    assert queue.all_tasks_done() is False
+
+    queue.get(0.0001)
+    queue.task_done()
+    assert queue.all_tasks_done() is True
+
+
+def test_message_queue_task_done__called_too_many_times__raises():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+
+    queue.put("msg")
+    queue.get(0.0001)
+    queue.task_done()
+
+    with pytest.raises(ValueError):
+        queue.task_done()
+
+
+def test_message_queue_task_done__maxlen_eviction__does_not_increment_counter():
+    # When a bounded queue is full, appendleft silently evicts the oldest
+    # item. The displaced message's task carries over to the newcomer, so
+    # the unfinished-task count must not grow.
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(max_length=2)
+
+    queue.put("a")
+    queue.put("b")
+    # Now full — this put evicts the oldest end.
+    queue.put("c")
+
+    # Only two messages are reachable; only two task_done() calls should
+    # bring the queue back to quiescence.
+    queue.get(0.0001)
+    queue.task_done()
+    queue.get(0.0001)
+    queue.task_done()
+    assert queue.all_tasks_done() is True
+
+
+def test_message_queue_put_back__bounded_queue_full__evicted_task_is_released():
+    # Regression: on a bounded queue, `put_back` (used by the rate-limit
+    # retry path) silently evicts an item from the opposite end. That
+    # evicted item was already counted in `_unfinished_tasks` at its
+    # original `put()` time and will never see `task_done()` — so the
+    # counter must drop by 1 to avoid stalling `join()`.
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue(max_length=2)
+
+    queue.put("a")  # unfinished=1
+    queue.put("b")  # unfinished=2; deque [b, a]
+
+    # Simulate a consumer popping a message and being rate-limited mid-process.
+    msg = queue.get(0.0001)  # deque [b]; unfinished=2 (msg still in flight)
+    assert msg == "a"
+
+    queue.put("c")  # unfinished=3; deque [c, b]
+
+    # Re-enqueue the popped message: the deque is full, so `append` evicts
+    # "c" from the left. "c" will never be processed, so its task slot must
+    # be released.
+    queue.put_back(msg)
+
+    # Two reachable messages — two task_done() calls must restore quiescence.
+    queue.get(0.0001)
+    queue.task_done()
+    queue.get(0.0001)
+    queue.task_done()
+    assert queue.all_tasks_done() is True
+    assert queue.join(timeout=0.05) is True
+
+
+def test_message_queue_clear__resets_counter_and_unblocks_join():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+
+    for i in range(3):
+        queue.put(f"msg-{i}")
+
+    assert queue.all_tasks_done() is False
+
+    queue.clear()
+
+    # clear() drops messages that will never see task_done(); the counter
+    # must be reset, otherwise join()/all_tasks_done() would stall forever.
+    assert queue.all_tasks_done() is True
+    assert queue.join(timeout=0.05) is True
+
+
+def test_message_queue_join__no_pending_tasks__returns_true_immediately():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+    start = time.monotonic()
+    assert queue.join(timeout=1.0) is True
+    assert time.monotonic() - start < 0.1
+
+
+def test_message_queue_join__pending_tasks_with_timeout__returns_false():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+    queue.put("msg")
+
+    start = time.monotonic()
+    assert queue.join(timeout=0.1) is False
+    elapsed = time.monotonic() - start
+    assert 0.1 <= elapsed < 0.5
+
+
+def test_message_queue_join__task_done_in_other_thread__wakes_waiter():
+    queue: message_queue.MessageQueue[str] = message_queue.MessageQueue()
+    queue.put("msg")
+    queue.get(0.0001)
+
+    def _consumer() -> None:
+        time.sleep(0.05)
+        queue.task_done()
+
+    threading.Thread(target=_consumer, daemon=True).start()
+
+    start = time.monotonic()
+    assert queue.join(timeout=2.0) is True
+    elapsed = time.monotonic() - start
+    # Should wake up shortly after the 0.05s sleep, not wait the full timeout.
+    assert elapsed < 1.0
+
+
+@pytest.mark.parametrize(
+    "maximal_queue_size,batch_factor,expected",
+    [
+        (100, 10, 10),
+        (100, 0, 100),
+    ],
+)
+def test_calculate_max_queue_size(maximal_queue_size, batch_factor, expected):
+    assert (
+        message_queue.calculate_max_queue_size(maximal_queue_size, batch_factor)
+        == expected
+    )

sdks/python/tests/unit/message_processing/test_message_streaming.py

@@ -0,0 +1,168 @@
+import threading
+from unittest import mock
+from unittest.mock import sentinel
+
+import pytest
+
+from opik.message_processing import messages
+from opik.message_processing import streamer_constructors
+from ...testlib import fake_message_factory
+
+NOT_USED = sentinel.NOT_USED
+
+
+@pytest.fixture
+def batched_streamer_and_mock_message_processor(
+    fake_file_upload_manager, fake_replay_manager
+):
+    tested = None
+    try:
+        mock_message_processor = mock.Mock()
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_message_processor,
+            n_consumers=1,
+            use_batching=True,
+            use_attachment_extraction=False,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=fake_replay_manager,
+        )
+
+        yield tested, mock_message_processor
+    finally:
+        if tested is not None:
+            tested.close(flush=False)
+
+
+def test_streamer__drain_to_processors__waits_until_slow_process_returns(
+    fake_file_upload_manager, fake_replay_manager
+):
+    # Regression: previously `_all_done()` checked `workers_idling and queue.empty`
+    # — both could flip True in the gap between the consumer popping a message
+    # and entering `message_processor.process(...)`. With the unfinished-task
+    # counter, drain_to_processors must block until `process` actually returns.
+    process_started = threading.Event()
+    release_process = threading.Event()
+    process_returned = threading.Event()
+
+    def slow_process(message):
+        process_started.set()
+        release_process.wait(timeout=5.0)
+        process_returned.set()
+
+    mock_message_processor = mock.Mock()
+    mock_message_processor.process.side_effect = slow_process
+
+    tested = streamer_constructors.construct_streamer(
+        message_processor=mock_message_processor,
+        n_consumers=1,
+        use_batching=False,
+        use_attachment_extraction=False,
+        file_uploader=fake_file_upload_manager,
+        max_queue_size=None,
+        fallback_replay_manager=fake_replay_manager,
+    )
+    try:
+        tested.put(messages.BaseMessage())
+
+        # Wait until the consumer has popped the message and is inside
+        # `process` — i.e. queue is empty but a task is in flight.
+        assert process_started.wait(timeout=2.0) is True
+
+        # A very short drain budget should NOT report success: process is
+        # still running, so the queue is not quiescent.
+        assert tested.drain_to_processors(timeout=0.1) is False
+        assert process_returned.is_set() is False
+
+        # Release the processor; the next drain must wait for it to finish
+        # and then return True.
+        release_process.set()
+        assert tested.drain_to_processors(timeout=2.0) is True
+        assert process_returned.is_set() is True
+    finally:
+        tested.close(flush=False)
+
+
+def test_streamer__happy_flow(batched_streamer_and_mock_message_processor):
+    tested, mock_message_processor = batched_streamer_and_mock_message_processor
+
+    test_messages = [messages.BaseMessage(), messages.BaseMessage]
+
+    tested.put(test_messages[0])
+    tested.put(test_messages[1])
+    assert tested.flush(timeout=0.01) is True
+
+    mock_message_processor.process.assert_has_calls(
+        [
+            mock.call(test_messages[0]),
+            mock.call(test_messages[1]),
+        ]
+    )
+
+
+@pytest.mark.parametrize(
+    "objects",
+    [
+        fake_message_factory.fake_create_trace_message_batch(count=3),
+        fake_message_factory.fake_create_trace_message_batch(count=3),
+    ],
+)
+def test_streamer__batching_disabled__messages_that_support_batching_are_processed_independently(
+    objects, fake_file_upload_manager
+):
+    mock_message_processor = mock.Mock()
+    tested = None
+    try:
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_message_processor,
+            n_consumers=1,
+            use_batching=False,
+            use_attachment_extraction=False,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=mock.Mock(),
+        )
+
+        for obj in objects:
+            tested.put(obj)
+        assert tested.flush(0.1) is True
+
+        mock_message_processor.process.assert_has_calls(
+            [
+                mock.call(objects[0]),
+                mock.call(objects[1]),
+                mock.call(objects[2]),
+            ]
+        )
+    finally:
+        if tested is not None:
+            tested.close(flush=False)
+
+
+def test_streamer__span__batching_enabled__messages_that_support_batching_are_processed_in_batch(
+    batched_streamer_and_mock_message_processor,
+):
+    tested, mock_message_processor = batched_streamer_and_mock_message_processor
+
+    create_span_messages = fake_message_factory.fake_create_trace_message_batch(count=3)
+
+    for message in create_span_messages:
+        tested.put(message)
+
+    assert tested.flush(1.1) is True
+
+    mock_message_processor.process.assert_called_once()
+
+
+def test_streamer__trace__batching_enabled__messages_that_support_batching_are_processed_in_batch(
+    batched_streamer_and_mock_message_processor,
+):
+    tested, mock_message_processor = batched_streamer_and_mock_message_processor
+
+    create_trace_messages = fake_message_factory.fake_create_trace_message_batch(3)
+    for message in create_trace_messages:
+        tested.put(message)
+
+    assert tested.flush(1.1) is True
+
+    mock_message_processor.process.assert_called_once()

sdks/python/tests/unit/message_processing/test_streamer_replay_manager.py

@@ -0,0 +1,387 @@
+"""
+Unit tests for streamer and replay manager interaction including lifecycle management.
+
+Tests the integration between the Streamer and ReplayManager, covering:
+- Initialization: Streamer sets the replay callback and starts the ReplayManager
+- Close lifecycle: Streamer closes the ReplayManager during its own close
+- Flush: Streamer conditionally flushes the ReplayManager based on server connection
+- Replay callback: Messages replayed by ReplayManager flow back into the Streamer's queue
+"""
+
+import datetime
+import time
+from typing import Generator, Tuple
+from unittest import mock
+
+import pytest
+
+from opik.healthcheck import connection_monitor
+from opik.message_processing import messages, streamer_constructors, streamer
+from opik.message_processing.replay import db_manager, replay_manager
+
+
+def _create_trace_message(
+    message_id: int, trace_id: str = "trace-1"
+) -> messages.CreateTraceMessage:
+    msg = messages.CreateTraceMessage(
+        trace_id=trace_id,
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.datetime(2024, 1, 1, 12, 0, 0),
+        end_time=datetime.datetime(2024, 1, 1, 12, 0, 1),
+        input={"query": "test"},
+        output={"answer": "response"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    msg.message_id = message_id
+    return msg
+
+
+def _make_replay_manager(
+    has_connection: bool = True,
+    tick_interval: float = 0.05,
+) -> replay_manager.ReplayManager:
+    """Create a ReplayManager with a real DB but a mock connection monitor."""
+    monitor = mock.MagicMock(spec=connection_monitor.OpikConnectionMonitor)
+    monitor.has_server_connection = has_connection
+    monitor.tick.return_value = connection_monitor.ConnectionStatus.connection_ok
+    return replay_manager.ReplayManager(
+        monitor=monitor,
+        batch_size=10,
+        batch_replay_delay=0.01,
+        tick_interval_seconds=tick_interval,
+    )
+
+
+@pytest.fixture
+def mock_replay_manager() -> mock.MagicMock:
+    """A fully mocked ReplayManager for testing Streamer interactions."""
+    rm = mock.MagicMock(spec=replay_manager.ReplayManager)
+    rm.has_server_connection = True
+    # start() and close() must be callables
+    rm.start = mock.Mock()
+    rm.close = mock.Mock()
+    rm.flush = mock.Mock()
+    rm.set_replay_callback = mock.Mock()
+    return rm
+
+
+@pytest.fixture
+def streamer_with_mock_replay(
+    fake_file_upload_manager,
+    mock_replay_manager: mock.MagicMock,
+) -> Generator[Tuple[streamer.Streamer, mock.MagicMock, mock.Mock], None, None]:
+    """Create a Streamer with a mocked ReplayManager and a mock message processor."""
+    tested = None
+    try:
+        mock_message_processor = mock.Mock()
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_message_processor,
+            n_consumers=1,
+            use_batching=False,
+            use_attachment_extraction=False,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=mock_replay_manager,
+        )
+        yield tested, mock_replay_manager, mock_message_processor
+    finally:
+        if tested is not None:
+            tested.close(flush=False)
+
+
+@pytest.fixture
+def streamer_with_real_replay(
+    fake_file_upload_manager,
+) -> Generator[
+    Tuple[streamer.Streamer, replay_manager.ReplayManager, mock.Mock], None, None
+]:
+    """Create a Streamer with a real ReplayManager backed by an in-memory DB."""
+    tested = None
+    rm = None
+    try:
+        rm = _make_replay_manager(has_connection=True)
+        mock_message_processor = mock.Mock()
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_message_processor,
+            n_consumers=1,
+            use_batching=False,
+            use_attachment_extraction=False,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=rm,
+        )
+        yield tested, rm, mock_message_processor
+    finally:
+        if tested is not None:
+            tested.close(flush=False)
+        if rm is not None and rm.is_alive():
+            rm.close()
+            rm.join(timeout=2)
+
+
+class TestStreamerReplayManagerInitialization:
+    def test_init__sets_replay_callback_on_replay_manager(
+        self,
+        streamer_with_mock_replay: Tuple[streamer.Streamer, mock.MagicMock, mock.Mock],
+    ):
+        """Streamer.__init__ must call set_replay_callback with its put method."""
+        tested, mock_rm, _ = streamer_with_mock_replay
+
+        mock_rm.set_replay_callback.assert_called_once_with(tested.put)
+
+    def test_init__starts_replay_manager(
+        self,
+        streamer_with_mock_replay: Tuple[streamer.Streamer, mock.MagicMock, mock.Mock],
+    ):
+        """Streamer.__init__ must call start() on the ReplayManager."""
+        _, mock_rm, _ = streamer_with_mock_replay
+
+        mock_rm.start.assert_called_once()
+
+    def test_init__callback_set_before_start(
+        self,
+        streamer_with_mock_replay: Tuple[streamer.Streamer, mock.MagicMock, mock.Mock],
+    ):
+        """set_replay_callback must be called before start()."""
+        _, mock_rm, _ = streamer_with_mock_replay
+
+        # Both were called; verify ordering via mock's call list
+        calls = mock_rm.method_calls
+        callback_idx = next(
+            i for i, c in enumerate(calls) if c[0] == "set_replay_callback"
+        )
+        start_idx = next(i for i, c in enumerate(calls) if c[0] == "start")
+        assert callback_idx < start_idx
+
+
+class TestStreamerCloseLifecycle:
+    def test_close__closes_replay_manager(
+        self,
+        streamer_with_mock_replay: Tuple[streamer.Streamer, mock.MagicMock, mock.Mock],
+    ):
+        """Streamer.close() must call close() on the ReplayManager."""
+        tested, mock_rm, _ = streamer_with_mock_replay
+
+        tested.close(flush=False)
+
+        mock_rm.close.assert_called_once()
+
+    def test_close__replay_manager_closed_before_queue_consumers(
+        self,
+        fake_file_upload_manager,
+    ):
+        """ReplayManager should be closed before queue consumers are shut down."""
+        mock_rm = mock.MagicMock(spec=replay_manager.ReplayManager)
+        mock_rm.has_server_connection = True
+
+        call_order = []
+        mock_rm.close.side_effect = lambda: call_order.append("replay_manager_close")
+
+        mock_processor = mock.Mock()
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_processor,
+            n_consumers=1,
+            use_batching=False,
+            use_attachment_extraction=False,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=mock_rm,
+        )
+
+        # Patch _close_queue_consumers to track ordering
+        original_close_consumers = tested._close_queue_consumers
+
+        def tracked_close_consumers():
+            call_order.append("queue_consumers_close")
+            original_close_consumers()
+
+        tested._close_queue_consumers = tracked_close_consumers
+
+        tested.close(flush=False)
+
+        assert "replay_manager_close" in call_order
+        assert "queue_consumers_close" in call_order
+        rm_idx = call_order.index("replay_manager_close")
+        qc_idx = call_order.index("queue_consumers_close")
+        assert rm_idx < qc_idx
+
+    def test_close__with_real_replay_manager__thread_stops(
+        self,
+        streamer_with_real_replay: Tuple[
+            streamer.Streamer, replay_manager.ReplayManager, mock.Mock
+        ],
+    ):
+        """After Streamer.close(), the ReplayManager thread should stop."""
+        tested, rm, _ = streamer_with_real_replay
+
+        assert rm.is_alive()
+
+        tested.close(flush=False)
+        rm.join(timeout=2)
+
+        assert not rm.is_alive()
+
+
+class TestStreamerFlushReplayInteraction:
+    def test_flush__with_server_connection__flushes_replay_manager(
+        self,
+        streamer_with_mock_replay: Tuple[streamer.Streamer, mock.MagicMock, mock.Mock],
+    ):
+        """When has_server_connection is True, flush() should call replay_manager.flush()."""
+        tested, mock_rm, _ = streamer_with_mock_replay
+        mock_rm.has_server_connection = True
+        mock_rm.flush.reset_mock()
+
+        tested.flush(timeout=2)
+
+        mock_rm.flush.assert_called_once()
+
+    def test_flush__without_server_connection__skips_replay_flush(
+        self,
+        streamer_with_mock_replay: Tuple[streamer.Streamer, mock.MagicMock, mock.Mock],
+    ):
+        """When has_server_connection is False, flush() should NOT call replay_manager.flush()."""
+        tested, mock_rm, _ = streamer_with_mock_replay
+        mock_rm.has_server_connection = False
+        mock_rm.flush.reset_mock()
+
+        tested.flush(timeout=2)
+
+        mock_rm.flush.assert_not_called()
+
+
+class TestReplayCallbackFlow:
+    def test_replay_callback__replayed_message_reaches_processor(
+        self,
+        streamer_with_real_replay: Tuple[
+            streamer.Streamer, replay_manager.ReplayManager, mock.Mock
+        ],
+    ):
+        """Messages replayed by the ReplayManager should flow through the Streamer
+        and reach the message processor."""
+        tested, rm, mock_processor = streamer_with_real_replay
+
+        # Register a failed message directly in the DB
+        msg = _create_trace_message(message_id=1)
+        rm.database_manager.register_message(
+            msg, status=db_manager.MessageStatus.failed
+        )
+
+        # Wait for the message to be processed (requires connection) to trigger replay by the replay manager
+        tested.flush(timeout=5)
+
+        assert mock_processor.process.call_count >= 1
+
+    def test_replay_callback__multiple_messages__all_reach_processor(
+        self,
+        streamer_with_real_replay: Tuple[
+            streamer.Streamer, replay_manager.ReplayManager, mock.Mock
+        ],
+    ):
+        """Multiple replayed messages should all reach the processor."""
+        tested, rm, mock_processor = streamer_with_real_replay
+
+        for i in range(5):
+            msg = _create_trace_message(message_id=i + 1, trace_id=f"trace-{i}")
+            rm.database_manager.register_message(
+                msg, status=db_manager.MessageStatus.failed
+            )
+
+        tested.flush(timeout=5)
+
+        assert mock_processor.process.call_count >= 5
+
+    def test_put__after_drain__replayed_messages_are_dropped(
+        self,
+        streamer_with_real_replay: Tuple[
+            streamer.Streamer, replay_manager.ReplayManager, mock.Mock
+        ],
+    ):
+        """After close() sets _drain=True, put() should silently discard messages."""
+        tested, rm, mock_processor = streamer_with_real_replay
+
+        # Close the streamer to set drain mode
+        tested.close(flush=False)
+        mock_processor.process.reset_mock()
+
+        # Now try to put a message (simulates a late replay callback)
+        msg = _create_trace_message(message_id=99)
+        tested.put(msg)
+
+        # Give a moment for any processing
+        time.sleep(0.2)
+
+        mock_processor.process.assert_not_called()
+
+
+class TestStreamerReplayManagerLifecycleEndToEnd:
+    def test_full_lifecycle__init_put_flush_close(
+        self,
+        streamer_with_real_replay: Tuple[
+            streamer.Streamer, replay_manager.ReplayManager, mock.Mock
+        ],
+    ):
+        """Full lifecycle: init → put messages → flush → close.
+        Verify the message processor receives the messages put into streamer.
+        Verify replay manager is alive during operation and stopped after close."""
+        tested, rm, mock_processor = streamer_with_real_replay
+
+        # ReplayManager should be running
+        assert rm.is_alive()
+        assert rm.database_manager.initialized
+        assert not rm.database_manager.closed
+
+        # Put a regular message through the streamer
+        msg = _create_trace_message(message_id=1)
+        tested.put(msg)
+
+        # Flush to ensure processing
+        tested.flush(timeout=5)
+
+        # Processor should have received the message
+        assert mock_processor.process.call_count >= 1
+
+        # Close everything
+        tested.close(flush=False)
+        rm.join(timeout=2)
+
+        # Verify cleanup
+        assert not rm.is_alive()
+        assert rm.database_manager.closed
+
+    def test_lifecycle__failed_message_replayed_on_flush(
+        self,
+        streamer_with_real_replay: Tuple[
+            streamer.Streamer, replay_manager.ReplayManager, mock.Mock
+        ],
+    ):
+        """Simulate: message registered → marked failed → flush triggers replay → processor receives it."""
+        tested, rm, mock_processor = streamer_with_real_replay
+
+        # 1. Register a message and mark it as failed
+        msg = _create_trace_message(message_id=1)
+        rm.register_message(msg)
+        rm.message_sent_failed(1, failure_reason="server down")
+
+        # Verify it's in a failed state
+        db_msg = rm.database_manager.get_db_message(1)
+        assert db_msg is not None
+        assert db_msg.status == db_manager.MessageStatus.failed
+
+        # 2. Flush the streamer (this flushes the replay manager too since connection is mocked as OK)
+        tested.flush(timeout=5)
+
+        # 3. The replayed message should reach the processor
+        assert mock_processor.process.call_count >= 1
+
+        # 4. Clean up
+        tested.close(flush=False)
+        rm.join(timeout=2)
+
+        assert not rm.is_alive()

sdks/python/tests/unit/message_processing/test_streamer_with_attachments.py

@@ -0,0 +1,492 @@
+"""
+Unit tests for streamer handling of create spans and traces messages with embedded attachments.
+
+This module tests the complete flow of attachment processing through the streamer, including:
+- AttachmentsPreprocessor wrapping messages
+- Message flow through the queue
+- Integration with the message processor
+"""
+
+from datetime import datetime
+from unittest import mock
+
+import pytest
+
+from opik.message_processing import messages, streamer_constructors
+from opik.message_processing.preprocessing import constants
+
+
+@pytest.fixture
+def streamer_with_attachments_enabled(fake_file_upload_manager):
+    """Create a streamer with attachment extraction enabled."""
+    tested = None
+    try:
+        mock_message_processor = mock.Mock()
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_message_processor,
+            n_consumers=1,
+            use_batching=True,
+            use_attachment_extraction=True,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=mock.Mock(),
+        )
+
+        yield tested, mock_message_processor
+    finally:
+        if tested is not None:
+            tested.close(flush=False)
+
+
+@pytest.fixture
+def streamer_with_attachments_disabled(fake_file_upload_manager):
+    """Create a streamer with attachment extraction disabled."""
+    tested = None
+    try:
+        mock_message_processor = mock.Mock()
+        tested = streamer_constructors.construct_streamer(
+            message_processor=mock_message_processor,
+            n_consumers=1,
+            use_batching=True,
+            use_attachment_extraction=False,
+            file_uploader=fake_file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=mock.Mock(),
+        )
+
+        yield tested, mock_message_processor
+    finally:
+        if tested is not None:
+            tested.close(flush=False)
+
+
+def test_streamer__create_span_message__attachments_enabled__does_not_wrap_without_end_time(
+    streamer_with_attachments_enabled,
+):
+    """Test that CreateSpanMessage without end_time is NOT wrapped even when attachments are enabled."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    # Create a span message with potential attachment data but no end_time
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.now(),
+        end_time=None,  # No end_time - should not be wrapped
+        input={"image": "base64-encoded-image-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(span_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should NOT be wrapped in AttachmentSupportingMessage since end_time is None
+    # (it may be batched into CreateSpansBatchMessage by the batching preprocessor)
+    assert not isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+
+def test_streamer__create_trace_message__attachments_enabled__does_not_wrap_without_end_time(
+    streamer_with_attachments_enabled,
+):
+    """Test that CreateTraceMessage without end_time is NOT wrapped even when attachments are enabled."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    # Create a trace message with potential attachment data but no end_time
+    trace_message = messages.CreateTraceMessage(
+        trace_id="trace-123",
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.now(),
+        end_time=None,  # No end_time - should not be wrapped
+        input={"document": "base64-encoded-pdf-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(trace_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should NOT be wrapped in AttachmentSupportingMessage since end_time is None
+    # (it may be batched into CreateTraceBatchMessage by the batching preprocessor)
+    assert not isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+
+def test_streamer__create_span_message__attachments_disabled__does_not_wrap(
+    streamer_with_attachments_disabled,
+):
+    """Test that CreateSpanMessage is NOT wrapped when attachments are disabled (can be batched)."""
+    tested, mock_message_processor = streamer_with_attachments_disabled
+
+    span_message = messages.CreateSpanMessage(
+        span_id="span-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.now(),
+        end_time=None,
+        input={"image": "base64-encoded-image-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(span_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should NOT be wrapped in AttachmentSupportingMessage
+    # (it may be batched into CreateSpansBatchMessage by the batching preprocessor)
+    assert not isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+
+def test_streamer__create_trace_message__attachments_disabled__does_not_wrap(
+    streamer_with_attachments_disabled,
+):
+    """Test that CreateTraceMessage is NOT wrapped when attachments are disabled (may be batched)."""
+    tested, mock_message_processor = streamer_with_attachments_disabled
+
+    trace_message = messages.CreateTraceMessage(
+        trace_id="trace-123",
+        project_name="test-project",
+        name="test-trace",
+        start_time=datetime.now(),
+        end_time=None,
+        input={"document": "base64-encoded-pdf-data"},
+        output=None,
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(trace_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should NOT be wrapped in AttachmentSupportingMessage
+    # (it may be batched into CreateTraceBatchMessage by the batching preprocessor)
+    assert not isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+
+def test_streamer__update_span_message__attachments_enabled__wraps_in_attachment_supporting_message(
+    streamer_with_attachments_enabled,
+):
+    """Test that UpdateSpanMessage is wrapped in AttachmentSupportingMessage when attachments are enabled."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    update_span_message = messages.UpdateSpanMessage(
+        span_id="span-update-123",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        end_time=datetime.now(),
+        input=None,
+        output={"result": "base64-encoded-data"},
+        metadata=None,
+        tags=None,
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        source="sdk",
+    )
+
+    tested.put(update_span_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should be wrapped in AttachmentSupportingMessage
+    assert isinstance(processed_message, messages.AttachmentSupportingMessage)
+    assert processed_message.original_message is update_span_message
+
+
+def test_streamer__update_trace_message__attachments_enabled__wraps_in_attachment_supporting_message(
+    streamer_with_attachments_enabled,
+):
+    """Test that UpdateTraceMessage is wrapped in AttachmentSupportingMessage when attachments are enabled."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    update_trace_message = messages.UpdateTraceMessage(
+        trace_id="trace-update-123",
+        project_name="test-project",
+        end_time=datetime.now(),
+        input=None,
+        output={"final_result": "base64-encoded-data"},
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        source="sdk",
+    )
+
+    tested.put(update_trace_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should be wrapped in AttachmentSupportingMessage
+    assert isinstance(processed_message, messages.AttachmentSupportingMessage)
+    assert processed_message.original_message is update_trace_message
+
+
+def test_streamer__multiple_span_messages__attachments_enabled__not_wrapped_without_end_time(
+    streamer_with_attachments_enabled,
+):
+    """Test that multiple CreateSpanMessages without end_time are NOT wrapped even when attachments are enabled."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    span_messages = [
+        messages.CreateSpanMessage(
+            span_id=f"span-{i}",
+            trace_id="trace-456",
+            project_name="test-project",
+            parent_span_id=None,
+            name=f"test-span-{i}",
+            start_time=datetime.now(),
+            end_time=None,  # No end_time - should not be wrapped
+            input={"data": f"base64-data-{i}"},
+            output=None,
+            metadata=None,
+            tags=None,
+            type="general",
+            usage=None,
+            model=None,
+            provider=None,
+            error_info=None,
+            total_cost=None,
+            last_updated_at=datetime.now(),
+            source="sdk",
+        )
+        for i in range(3)
+    ]
+
+    for span_message in span_messages:
+        tested.put(span_message)
+
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify all messages were processed (can be batched)
+    # When batching is enabled, they might be combined into a single batch
+    assert mock_message_processor.process.call_count >= 1
+
+    # Check that none of the messages are wrapped in AttachmentSupportingMessage
+    for call in mock_message_processor.process.call_args_list:
+        processed_message = call[0][0]
+        assert not isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+
+def test_streamer__span_with_all_fields_populated__attachments_enabled__wraps_correctly(
+    streamer_with_attachments_enabled,
+):
+    """Test that span messages with all fields populated are wrapped correctly."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    span_message = messages.CreateSpanMessage(
+        span_id="span-full",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id="parent-span-789",
+        name="test-span-full",
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        input={"prompt": "test prompt", "image": "base64-image"},
+        output={"response": "test response", "chart": "base64-chart"},
+        metadata={"model": "gpt-4", "doc": "base64-pdf"},
+        tags=["test", "full"],
+        type="llm",
+        usage={"prompt_tokens": 10, "completion_tokens": 20},
+        model="gpt-4",
+        provider="openai",
+        error_info=None,
+        total_cost=0.05,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(span_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should be wrapped
+    assert isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+    # Verify the original message is preserved with all fields
+    original = processed_message.original_message
+    assert original.span_id == "span-full"
+    assert original.input == {"prompt": "test prompt", "image": "base64-image"}
+    assert original.output == {"response": "test response", "chart": "base64-chart"}
+    assert original.metadata == {"model": "gpt-4", "doc": "base64-pdf"}
+
+
+def test_streamer__trace_with_all_fields_populated__attachments_enabled__wraps_correctly(
+    streamer_with_attachments_enabled,
+):
+    """Test that trace messages with all fields populated are wrapped correctly."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    trace_message = messages.CreateTraceMessage(
+        trace_id="trace-full",
+        project_name="test-project",
+        name="test-trace-full",
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        input={"request": "test request", "audio": "base64-audio"},
+        output={"result": "test result", "video": "base64-video"},
+        metadata={"session": "abc123", "file": "base64-file"},
+        tags=["production", "important"],
+        error_info=None,
+        thread_id="thread-123",
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(trace_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should be wrapped
+    assert isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+    # Verify the original message is preserved with all fields
+    original = processed_message.original_message
+    assert original.trace_id == "trace-full"
+    assert original.input == {"request": "test request", "audio": "base64-audio"}
+    assert original.output == {"result": "test result", "video": "base64-video"}
+    assert original.metadata == {"session": "abc123", "file": "base64-file"}
+    assert original.thread_id == "thread-123"
+
+
+def test_streamer__message_already_marked__attachments_enabled__does_not_wrap_twice(
+    streamer_with_attachments_enabled,
+):
+    """Test that messages already marked are not wrapped again (prevents infinite recursion)."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    span_message = messages.CreateSpanMessage(
+        span_id="span-marked",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.now(),
+        end_time=None,
+        input={"data": "value"},
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    # Mark the message as already processed
+    setattr(span_message, constants.MARKER_ATTRIBUTE_NAME, True)
+
+    tested.put(span_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should NOT be wrapped in AttachmentSupportingMessage since it's already marked
+    # (it may still be batched into CreateSpansBatchMessage by the batching preprocessor)
+    assert not isinstance(processed_message, messages.AttachmentSupportingMessage)
+
+
+def test_streamer__null_input_output_metadata__attachments_enabled__doesnt_wraps(
+    streamer_with_attachments_enabled,
+):
+    """Test that messages with None input/output/metadata are not wrapped."""
+    tested, mock_message_processor = streamer_with_attachments_enabled
+
+    span_message = messages.CreateSpanMessage(
+        span_id="span-nulls",
+        trace_id="trace-456",
+        project_name="test-project",
+        parent_span_id=None,
+        name="test-span",
+        start_time=datetime.now(),
+        end_time=None,
+        input=None,
+        output=None,
+        metadata=None,
+        tags=None,
+        type="general",
+        usage=None,
+        model=None,
+        provider=None,
+        error_info=None,
+        total_cost=None,
+        last_updated_at=datetime.now(),
+        source="sdk",
+    )
+
+    tested.put(span_message)
+    assert tested.flush(timeout=1.0) is True
+
+    # Verify the message was processed
+    mock_message_processor.process.assert_called_once()
+    processed_message = mock_message_processor.process.call_args[0][0]
+
+    # The message should not be wrapped because it has no potential attachments
+    assert not isinstance(processed_message, messages.AttachmentSupportingMessage)

sdks/python/tests/unit/message_processing/test_uploads_streaming.py

@@ -0,0 +1,95 @@
+import time
+from typing import Tuple, Any, Generator
+from unittest import mock
+from unittest.mock import sentinel, patch
+
+import pytest
+
+from opik.file_upload import upload_manager
+from opik.message_processing import messages, streamer, streamer_constructors
+from opik.message_processing.processors import online_message_processor
+
+NOT_USED = sentinel.NOT_USED
+
+
+@pytest.fixture
+def streamer_with_file_upload_manager(
+    request,
+) -> Generator[Tuple[streamer.Streamer, upload_manager.FileUploadManager], Any, None]:
+    def upload_noop(**kwargs):
+        print("Uploading file")
+        time.sleep(request.param)
+
+    with patch("opik.file_upload.upload_manager.file_uploader") as uploader_mock:
+        uploader_mock.upload_attachment = upload_noop
+
+        file_upload_manager = upload_manager.FileUploadManager(
+            rest_client=mock.Mock(), httpx_client=mock.Mock(), worker_count=2
+        )
+        online = online_message_processor.OpikMessageProcessor(
+            rest_client=mock.Mock(),
+            file_upload_manager=file_upload_manager,
+            fallback_replay_manager=mock.Mock(),
+            unauthorized_message_types_registry=mock.Mock(),
+        )
+        streamer_ = streamer_constructors.construct_streamer(
+            message_processor=online,
+            n_consumers=1,
+            use_batching=True,
+            use_attachment_extraction=False,
+            file_uploader=file_upload_manager,
+            max_queue_size=None,
+            fallback_replay_manager=mock.Mock(),
+        )
+
+        yield streamer_, file_upload_manager
+
+
+@pytest.mark.parametrize("streamer_with_file_upload_manager", [0.5], indirect=True)
+def test_streamer__flush__attachment_uploads__ok(
+    streamer_with_file_upload_manager, temp_file_15mb
+):
+    message_streamer, file_upload_manager = streamer_with_file_upload_manager
+
+    attachment = messages.CreateAttachmentMessage(
+        file_path=temp_file_15mb.name,
+        file_name="test_file",
+        mime_type=None,
+        entity_type="span",
+        entity_id=NOT_USED,
+        project_name=NOT_USED,
+        encoded_url_override=NOT_USED,
+    )
+
+    message_streamer.put(attachment)
+    message_streamer.put(attachment)
+
+    # we have timeout greater than upload time 1 > 0.5, thus all uploads will be completed
+    assert message_streamer.flush(timeout=1, upload_sleep_time=1) is True
+
+    assert file_upload_manager.remaining_data().uploads == 0
+
+
+@pytest.mark.parametrize("streamer_with_file_upload_manager", [0.5], indirect=True)
+def test_streamer__flush__attachment_uploads__timeout(
+    streamer_with_file_upload_manager, temp_file_15mb
+):
+    message_streamer, file_upload_manager = streamer_with_file_upload_manager
+
+    attachment = messages.CreateAttachmentMessage(
+        file_path=temp_file_15mb.name,
+        file_name="test_file",
+        mime_type=None,
+        entity_type="span",
+        entity_id=NOT_USED,
+        project_name=NOT_USED,
+        encoded_url_override=NOT_USED,
+    )
+
+    message_streamer.put(attachment)
+    message_streamer.put(attachment)
+
+    # we have timeout less than upload time 0.1 < 0.5, thus not all uploads will be completed
+    assert message_streamer.flush(timeout=0.1, upload_sleep_time=0.1) is False
+
+    assert file_upload_manager.remaining_data().uploads == 2