chore: import upstream snapshot with attribution

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()