chore: import upstream snapshot with attribution

sdks/python/tests/unit/anonymizer/test_factory.py

@@ -0,0 +1,275 @@
+import pytest
+
+from opik.anonymizer import factory
+from opik.anonymizer import rules_anonymizer
+
+
+class TestCreateAnonymizer:
+    """Test suite for the create_anonymizer factory function."""
+
+    def test_create_anonymizer__single_dict_rule(self):
+        """Test creating anonymizer with a single dictionary rule."""
+        rules = {"regex": r"\d+", "replace": "***"}
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 1
+        assert anonymizer.max_depth == 10
+
+        result = anonymizer.anonymize("Phone: 123-456")
+        assert result == "Phone: ***-***"
+
+    def test_create_anonymizer__single_tuple_rule(self):
+        """Test creating anonymizer with a single tuple rule."""
+        rules = (r"\d+", "***")
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 1
+
+        result = anonymizer.anonymize("Phone: 123-456")
+        assert result == "Phone: ***-***"
+
+    def test_create_anonymizer__single_function_rule(self):
+        """Test creating anonymizer with a single function rule."""
+
+        def uppercase_rule(text: str) -> str:
+            return text.upper()
+
+        rules = uppercase_rule
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 1
+
+        result = anonymizer.anonymize("hello world")
+        assert result == "HELLO WORLD"
+
+    def test_create_anonymizer__list_of_dict_rules(self):
+        """Test creating anonymizer with a list of dictionary rules."""
+        rules = [
+            {"regex": r"\d+", "replace": "[NUMBER]"},
+            {
+                "regex": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
+                "replace": "[EMAIL]",
+            },
+        ]
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 2
+
+        result = anonymizer.anonymize("Contact: test@example.com or 123")
+        assert result == "Contact: [EMAIL] or [NUMBER]"
+
+    def test_create_anonymizer__list_of_tuple_rules(self):
+        """Test creating anonymizer with a list of tuple rules."""
+        rules = [
+            (r"\d+", "[NUMBER]"),
+            (r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}", "[EMAIL]"),
+        ]
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 2
+
+        result = anonymizer.anonymize("Contact: test@example.com or 123")
+        assert result == "Contact: [EMAIL] or [NUMBER]"
+
+    def test_create_anonymizer__list_of_function_rules(self):
+        """Test creating anonymizer with a list of function rules."""
+
+        def uppercase_rule(text: str) -> str:
+            return text.upper()
+
+        def replace_spaces(text: str) -> str:
+            return text.replace(" ", "_")
+
+        rules = [uppercase_rule, replace_spaces]
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 2
+
+        result = anonymizer.anonymize("hello world")
+        assert result == "HELLO_WORLD"
+
+    def test_create_anonymizer__mixed_rules_list(self):
+        """Test creating anonymizer with mixed rule types in a list."""
+
+        def prefix_rule(text: str) -> str:
+            return f"[PROCESSED] {text}"
+
+        rules = [{"regex": r"\d+", "replace": "***"}, (r"test", "TEST"), prefix_rule]
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 3
+
+        result = anonymizer.anonymize("test 123")
+        assert result == "[PROCESSED] TEST ***"
+
+    def test_create_anonymizer__custom_max_depth(self):
+        """Test creating anonymizer with custom max_depth."""
+        rules = {"regex": r"\d+", "replace": "***"}
+        anonymizer = factory.create_anonymizer(rules, max_depth=5)
+
+        assert anonymizer.max_depth == 5
+
+    def test_create_anonymizer__complex_nested_data(self):
+        """Test created anonymizer with complex nested data."""
+        rules = [
+            {"regex": r"\d{3}-\d{3}-\d{4}", "replace": "[PHONE]"},
+            {
+                "regex": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
+                "replace": "[EMAIL]",
+            },
+        ]
+        anonymizer = factory.create_anonymizer(rules)
+
+        data = {
+            "users": [
+                {"name": "John", "contact": "john@test.com", "phone": "123-456-7890"},
+                {"name": "Jane", "contact": "jane@test.org", "phone": "987-654-3210"},
+            ]
+        }
+
+        result = anonymizer.anonymize(data)
+
+        expected = {
+            "users": [
+                {"name": "John", "contact": "[EMAIL]", "phone": "[PHONE]"},
+                {"name": "Jane", "contact": "[EMAIL]", "phone": "[PHONE]"},
+            ]
+        }
+        assert result == expected
+
+    def test_create_anonymizer__invalid_dict_rule_missing_regex__error_raised(self):
+        """Test error handling for a dictionary rule missing the 'regex' key."""
+        rules = {"replace": "***"}
+        with pytest.raises(
+            ValueError, match="Dictionary rule must have 'regex' and 'replace' keys"
+        ):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__invalid_dict_rule_missing_replace__error_raised(self):
+        """Test error handling for a dictionary rule missing the 'replace' key."""
+        rules = {"regex": r"\d+"}
+        with pytest.raises(
+            ValueError, match="Dictionary rule must have 'regex' and 'replace' keys"
+        ):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__invalid_tuple_rule_wrong_length__error_raised(self):
+        """Test error handling for tuple rule with the wrong length."""
+        rules = (r"\d+",)  # Only one element
+        with pytest.raises(ValueError, match="Tuple rule must have exactly 2 elements"):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__invalid_tuple_rule_too_many_elements__error_raised(
+        self,
+    ):
+        """Test error handling for tuple rule with too many elements."""
+        rules = (r"\d+", "***", "extra")
+        with pytest.raises(ValueError, match="Tuple rule must have exactly 2 elements"):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__invalid_dict_rule_in_list__error_raised(self):
+        """Test error handling for invalid dictionary rule in a list."""
+        rules = [
+            {"regex": r"\d+", "replace": "***"},
+            {"regex": r"test"},  # Missing 'replace'
+        ]
+        with pytest.raises(
+            ValueError, match="Dictionary rule must have 'regex' and 'replace' keys"
+        ):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__invalid_tuple_rule_in_list__error_raised(self):
+        """Test error handling for invalid tuple rule in a list."""
+        rules = [
+            (r"\d+", "***"),
+            (r"test",),  # Wrong length
+        ]
+        with pytest.raises(ValueError, match="Tuple rule must have exactly 2 elements"):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__unsupported_rule_type_in_list__error_raised(self):
+        """Test error handling for an unsupported rule type in the list."""
+        rules = [
+            {"regex": r"\d+", "replace": "***"},
+            123,  # Invalid type
+        ]
+        with pytest.raises(ValueError, match="Unsupported rule type in list"):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__unsupported_rules_type__error_raised(self):
+        """Test error handling for completely unsupported rules type."""
+        rules = 123  # Invalid type
+        with pytest.raises(ValueError, match="Unsupported rules type"):
+            factory.create_anonymizer(rules)
+
+    def test_create_anonymizer__empty_list(self):
+        """Test creating anonymizer with an empty rules list."""
+        rules = []
+        anonymizer = factory.create_anonymizer(rules)
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 0
+
+        # Should return original text with no rules
+        result = anonymizer.anonymize("test data")
+        assert result == "test data"
+
+    def test_create_anonymizer__lambda_function(self):
+        """Test creating anonymizer with lambda function rule."""
+        anonymizer = factory.create_anonymizer(
+            lambda text: text.replace("secret", "[REDACTED]")
+        )
+
+        assert isinstance(anonymizer, rules_anonymizer.RulesAnonymizer)
+        assert len(anonymizer.rules) == 1
+
+        result = anonymizer.anonymize("This is secret information")
+        assert result == "This is [REDACTED] information"
+
+    def test_create_anonymizer__real_world_pii_rules(self):
+        """Test creating anonymizer with realistic PII anonymization rules."""
+        rules = [
+            # Phone numbers
+            {"regex": r"\b\d{3}-\d{3}-\d{4}\b", "replace": "[PHONE]"},
+            # Email addresses
+            {
+                "regex": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
+                "replace": "[EMAIL]",
+            },
+            # SSN
+            {"regex": r"\b\d{3}-\d{2}-\d{4}\b", "replace": "[SSN]"},
+            # Credit card (simplified)
+            {"regex": r"\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b", "replace": "[CARD]"},
+        ]
+        anonymizer = factory.create_anonymizer(rules)
+
+        text = """
+        Personal Information:
+        Name: John Doe
+        Email: john.doe@company.com
+        Phone: 555-123-4567
+        SSN: 123-45-6789
+        Credit Card: 1234 5678 9012 3456
+        """
+
+        result = anonymizer.anonymize(text)
+
+        assert result != text
+        assert (
+            result
+            == """
+        Personal Information:
+        Name: John Doe
+        Email: [EMAIL]
+        Phone: [PHONE]
+        SSN: [SSN]
+        Credit Card: [CARD]
+        """
+        )

sdks/python/tests/unit/anonymizer/test_recursive_anonymizer.py

@@ -0,0 +1,560 @@
+from typing import Any, Dict, List
+
+from unittest.mock import Mock
+
+from opik.anonymizer.recursive_anonymizer import RecursiveAnonymizer
+
+
+class TestRecursiveAnonymizer:
+    """Test suite for RecursiveAnonymizer parameter handling and nested structure processing."""
+
+    def test_recursive_anonymizer__simple_string__calls_anonymize_text_with_correct_parameters(
+        self,
+    ):
+        """Test that anonymize_text is called with correct parameters for a simple string."""
+
+        class MockRecursiveAnonymizer(RecursiveAnonymizer):
+            def __init__(self):
+                super().__init__()
+                self.anonymize_text = Mock(return_value="[ANONYMIZED]")
+
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                return self.anonymize_text(data, **kwargs)
+
+        anonymizer = MockRecursiveAnonymizer()
+
+        # Test with initial parameters
+        result = anonymizer.anonymize(
+            "sensitive text", field_name="input", object_type=dict
+        )
+
+        # Verify anonymize_text was called correctly
+        anonymizer.anonymize_text.assert_called_once_with(
+            "sensitive text", field_name="input", object_type=dict
+        )
+        assert result == "[ANONYMIZED]"
+
+    def test_recursive_anonymizer__nested_dict__preserves_field_path_in_parameters(
+        self,
+    ):
+        """Test that field paths are correctly built and passed for nested dictionaries."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(
+                    {
+                        "data": data,
+                        "field_name": kwargs.get("field_name"),
+                        "object_type": kwargs.get("object_type"),
+                    }
+                )
+                return f"[ANON:{data}]"
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        nested_data = {
+            "user": {
+                "email": "user@example.com",
+                "profile": {"name": "John Doe", "phone": "555-1234"},
+            },
+            "metadata": {"api_key": "secret123"},
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(
+            nested_data, field_name="trace", object_type="TraceMessage"
+        )
+
+        # Verify the correct field paths were generated
+        expected_calls = [
+            {
+                "data": "user@example.com",
+                "field_name": "trace.user.email",
+                "object_type": "TraceMessage",
+            },
+            {
+                "data": "John Doe",
+                "field_name": "trace.user.profile.name",
+                "object_type": "TraceMessage",
+            },
+            {
+                "data": "555-1234",
+                "field_name": "trace.user.profile.phone",
+                "object_type": "TraceMessage",
+            },
+            {
+                "data": "secret123",
+                "field_name": "trace.metadata.api_key",
+                "object_type": "TraceMessage",
+            },
+        ]
+
+        assert len(calls_log) == 4
+        for expected_call in expected_calls:
+            assert expected_call in calls_log
+
+        # Verify the structure was preserved with anonymized content
+        assert result["user"]["email"] == "[ANON:user@example.com]"
+        assert result["user"]["profile"]["name"] == "[ANON:John Doe]"
+        assert result["user"]["profile"]["phone"] == "[ANON:555-1234]"
+        assert result["metadata"]["api_key"] == "[ANON:secret123]"
+
+    def test_recursive_anonymizer__nested_list__preserves_field_path_with_indices(self):
+        """Test that field paths include list indices for nested lists."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(
+                    {
+                        "data": data,
+                        "field_name": kwargs.get("field_name"),
+                        "object_type": kwargs.get("object_type"),
+                    }
+                )
+                return f"[ANON:{data}]"
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        list_data = [
+            "first item",
+            {"nested": "nested value", "list": ["inner item 1", "inner item 2"]},
+            ["list item 1", "list item 2"],
+        ]
+
+        result: List[Any] = anonymizer.anonymize(
+            list_data, field_name="input", object_type="SpanMessage"
+        )
+
+        # Verify the correct field paths were generated with indices
+        expected_calls = [
+            {
+                "data": "first item",
+                "field_name": "input.0",
+                "object_type": "SpanMessage",
+            },
+            {
+                "data": "nested value",
+                "field_name": "input.1.nested",
+                "object_type": "SpanMessage",
+            },
+            {
+                "data": "inner item 1",
+                "field_name": "input.1.list.0",
+                "object_type": "SpanMessage",
+            },
+            {
+                "data": "inner item 2",
+                "field_name": "input.1.list.1",
+                "object_type": "SpanMessage",
+            },
+            {
+                "data": "list item 1",
+                "field_name": "input.2.0",
+                "object_type": "SpanMessage",
+            },
+            {
+                "data": "list item 2",
+                "field_name": "input.2.1",
+                "object_type": "SpanMessage",
+            },
+        ]
+
+        assert len(calls_log) == 6
+        for expected_call in expected_calls:
+            assert expected_call in calls_log
+
+        # Verify the structure was preserved with anonymized content
+        assert result[0] == "[ANON:first item]"
+        assert result[1]["nested"] == "[ANON:nested value]"
+        assert result[1]["list"] == ["[ANON:inner item 1]", "[ANON:inner item 2]"]
+        assert result[2] == ["[ANON:list item 1]", "[ANON:list item 2]"]
+
+    def test_recursive_anonymizer__mixed_complex_structure__handles_all_parameter_combinations(
+        self,
+    ):
+        """Test a complex nested structure with mixed dictionaries, lists, and strings."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(
+                    {
+                        "data": data,
+                        "field_name": kwargs.get("field_name"),
+                        "object_type": kwargs.get("object_type"),
+                        "custom_param": kwargs.get("custom_param"),
+                    }
+                )
+                return f"[{kwargs.get('field_name', 'UNKNOWN')}:{data}]"
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        complex_data = {
+            "messages": [
+                {"role": "user", "content": "Hello, my email is john@example.com"},
+                {
+                    "role": "assistant",
+                    "content": "I can help you with that",
+                    "attachments": ["file1.txt", "file2.pdf"],
+                },
+            ],
+            "metadata": {
+                "session_id": "sess_12345",
+                "user_data": {
+                    "preferences": ["pref1", "pref2"],
+                    "settings": {"theme": "dark", "language": "en"},
+                },
+            },
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(
+            complex_data,
+            field_name="output",
+            object_type="TraceMessage",
+            custom_param="test_value",
+        )
+
+        # Verify all expected calls were made with correct parameters
+        expected_field_paths = [
+            "output.messages.0.role",
+            "output.messages.0.content",
+            "output.messages.1.role",
+            "output.messages.1.content",
+            "output.messages.1.attachments.0",
+            "output.messages.1.attachments.1",
+            "output.metadata.session_id",
+            "output.metadata.user_data.preferences.0",
+            "output.metadata.user_data.preferences.1",
+            "output.metadata.user_data.settings.theme",
+            "output.metadata.user_data.settings.language",
+        ]
+
+        assert len(calls_log) == len(expected_field_paths)
+
+        # Verify each call has the correct parameters
+        for call in calls_log:
+            assert call["object_type"] == "TraceMessage"
+            assert call["custom_param"] == "test_value"
+            assert call["field_name"] in expected_field_paths
+
+        # Verify specific anonymization results
+        assert (
+            result["messages"][0]["content"]
+            == "[output.messages.0.content:Hello, my email is john@example.com]"
+        )
+        assert (
+            result["metadata"]["session_id"]
+            == "[output.metadata.session_id:sess_12345]"
+        )
+        assert (
+            result["metadata"]["user_data"]["settings"]["theme"]
+            == "[output.metadata.user_data.settings.theme:dark]"
+        )
+
+    def test_recursive_anonymizer__max_depth_limiting__stops_recursion_at_limit(self):
+        """Test that max_depth parameter properly limits recursion depth."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def __init__(self, max_depth: int = 2):
+                super().__init__(max_depth=max_depth)
+
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(
+                    {
+                        "data": data,
+                        "field_name": kwargs.get("field_name"),
+                    }
+                )
+                return f"[ANON:{data}]"
+
+        anonymizer = ParameterTrackingAnonymizer(max_depth=2)
+
+        # Create a structure where strings at different depths can be tested
+        deeply_nested = {
+            "level1_text": "depth 1 - should be processed",
+            "level1": {
+                "level2_text": "depth 2 - should be processed",
+                "level2": {"level3_text": "depth 3 - should NOT be processed"},
+            },
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(deeply_nested, field_name="root")
+
+        field_names = [call["field_name"] for call in calls_log]
+
+        # The recursion depth starts at 0 for the initial call, so with max_depth=2:
+        # - root.level1_text is at depth 1 (should be processed)
+        # - root.level1.level2_text is at depth 2 (exceeds max_depth=2, should NOT be processed)
+        # - root.level1.level2.level3_text is at depth 3+ (exceeds max_depth=2, should NOT be processed)
+
+        assert "root.level1_text" in field_names
+        assert len([name for name in field_names if "level2_text" in name]) == 0
+        assert len([name for name in field_names if "level3_text" in name]) == 0
+
+        # Verify the results
+        assert result["level1_text"] == "[ANON:depth 1 - should be processed]"
+        assert (
+            result["level1"]["level2_text"] == "depth 2 - should be processed"
+        )  # Unchanged due to depth limit
+        assert (
+            result["level1"]["level2"]["level3_text"]
+            == "depth 3 - should NOT be processed"
+        )  # Unchanged
+
+    def test_recursive_anonymizer__non_string_types__preserves_unchanged(self):
+        """Test that non-string types are preserved without calling anonymize_text."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(data)
+                return f"[ANON:{data}]"
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        mixed_types_data = {
+            "string_field": "text to anonymize",
+            "int_field": 42,
+            "float_field": 3.14,
+            "bool_field": True,
+            "none_field": None,
+            "nested": {
+                "another_string": "another text",
+                "number": 100,
+                "list_with_mixed": ["string in list", 123, False],
+            },
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(
+            mixed_types_data, field_name="data"
+        )
+
+        # Should only anonymize strings (there are 3: "text to anonymize", "another text", "string in list")
+        assert len(calls_log) == 3
+        assert "text to anonymize" in calls_log
+        assert "another text" in calls_log
+        assert "string in list" in calls_log
+
+        # Non-string types should be preserved
+        assert result["int_field"] == 42
+        assert result["float_field"] == 3.14
+        assert result["bool_field"] is True
+        assert result["none_field"] is None
+        assert result["nested"]["number"] == 100
+        assert result["nested"]["list_with_mixed"][1] == 123
+        assert result["nested"]["list_with_mixed"][2] is False
+
+        # Strings should be anonymized
+        assert result["string_field"] == "[ANON:text to anonymize]"
+        assert result["nested"]["another_string"] == "[ANON:another text]"
+        assert result["nested"]["list_with_mixed"][0] == "[ANON:string in list]"
+
+    def test_recursive_anonymizer__empty_structures__handles_gracefully(self):
+        """Test that empty dictionaries and lists are handled gracefully."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(data)
+                return f"[ANON:{data}]"
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        empty_structures = {
+            "empty_dict": {},
+            "empty_list": [],
+            "mixed": {
+                "nested_empty_dict": {},
+                "nested_empty_list": [],
+                "text": "some text",
+            },
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(
+            empty_structures, field_name="test"
+        )
+
+        # Should only call anonymize_text for the one string
+        assert len(calls_log) == 1
+        assert "some text" in calls_log
+
+        # Empty structures should be preserved
+        assert result["empty_dict"] == {}
+        assert result["empty_list"] == []
+        assert result["mixed"]["nested_empty_dict"] == {}
+        assert result["mixed"]["nested_empty_list"] == []
+        assert result["mixed"]["text"] == "[ANON:some text]"
+
+    def test_recursive_anonymizer__field_specific_anonymization__uses_field_path_for_logic(
+        self,
+    ):
+        """Test that anonymizers can use field paths to implement field-specific logic."""
+
+        class FieldSpecificAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                field_name = kwargs.get("field_name", "")
+
+                # Different anonymization based on a field path
+                if "email" in field_name:
+                    return "[EMAIL_REDACTED]"
+                elif "phone" in field_name:
+                    return "[PHONE_REDACTED]"
+                elif "api_key" in field_name:
+                    return "[API_KEY_REDACTED]"
+                elif field_name.endswith(".name"):
+                    return "[NAME_REDACTED]"
+                else:
+                    return data  # Leave unchanged for other fields
+
+        anonymizer = FieldSpecificAnonymizer()
+
+        user_data = {
+            "user": {
+                "email": "john.doe@example.com",
+                "name": "John Doe",
+                "phone": "555-1234",
+                "notes": "Regular user notes",
+            },
+            "config": {
+                "api_key": "secret123",
+                "description": "Configuration description",
+            },
+            "contacts": [
+                {
+                    "name": "Contact One",
+                    "email": "contact1@example.com",
+                    "other_info": "Some other information",
+                }
+            ],
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(user_data, field_name="input")
+
+        # Verify field-specific anonymization
+        assert result["user"]["email"] == "[EMAIL_REDACTED]"
+        assert result["user"]["name"] == "[NAME_REDACTED]"
+        assert result["user"]["phone"] == "[PHONE_REDACTED]"
+        assert result["user"]["notes"] == "Regular user notes"  # Unchanged
+
+        assert result["config"]["api_key"] == "[API_KEY_REDACTED]"
+        assert (
+            result["config"]["description"] == "Configuration description"
+        )  # Unchanged
+
+        assert result["contacts"][0]["name"] == "[NAME_REDACTED]"
+        assert result["contacts"][0]["email"] == "[EMAIL_REDACTED]"
+        assert (
+            result["contacts"][0]["other_info"] == "Some other information"
+        )  # Unchanged
+
+    def test_recursive_anonymizer__parameter_propagation__all_kwargs_preserved(self):
+        """Test that all custom kwargs are properly propagated to anonymize_text."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(kwargs.copy())
+                return data
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        test_data = {"nested": {"text": "sample text"}}
+
+        # Pass multiple custom parameters
+        anonymizer.anonymize(
+            test_data,
+            field_name="test_field",
+            object_type="TestMessage",
+            custom_param1="value1",
+            custom_param2=42,
+            custom_param3={"nested": "param"},
+        )
+
+        # Should have one call for the text
+        assert len(calls_log) == 1
+        kwargs = calls_log[0]
+
+        # Verify all parameters were preserved
+        assert kwargs["field_name"] == "test_field.nested.text"
+        assert kwargs["object_type"] == "TestMessage"
+        assert kwargs["custom_param1"] == "value1"
+        assert kwargs["custom_param2"] == 42
+        assert kwargs["custom_param3"] == {"nested": "param"}
+
+    def test_recursive_anonymizer__circular_reference_protection__respects_max_depth(
+        self,
+    ):
+        """Test that max_depth prevents infinite recursion even with circular references."""
+
+        calls_count = 0
+
+        class CountingAnonymizer(RecursiveAnonymizer):
+            def __init__(self):
+                super().__init__(max_depth=3)
+
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                nonlocal calls_count
+                calls_count += 1
+                return f"[CALL_{calls_count}:{data}]"
+
+        anonymizer = CountingAnonymizer()
+
+        # Create a structure that tests depth limiting with strings at different levels
+        deep_structure = {
+            "text_at_level1": "depth 1 - should be processed",
+            "level1": {
+                "text_at_level2": "depth 2 - should be processed",
+                "level2": {
+                    "text_at_level3": "depth 3 - should be processed",
+                    "level3": {"text_at_level4": "depth 4 - should NOT be processed"},
+                },
+            },
+        }
+
+        result: Dict[str, Any] = anonymizer.anonymize(deep_structure, field_name="root")
+
+        # With max_depth=3, strings at depth 1, 2, 3 should be processed, but depth 4+ should not
+
+        # Verify the structure - strings beyond max_depth should remain unchanged
+        assert (
+            "depth 4 - should NOT be processed"
+            in result["level1"]["level2"]["level3"]["text_at_level4"]
+        )
+
+        # The strings within max_depth should be processed
+        # With max_depth=3, only strings at depth 1 and 2 get processed
+        assert calls_count == 2  # Should process the first 2 strings within max_depth
+
+    def test_recursive_anonymizer__no_field_name_provided__uses_empty_string_as_base(
+        self,
+    ):
+        """Test behavior when no field_name is provided in initial kwargs."""
+
+        calls_log = []
+
+        class ParameterTrackingAnonymizer(RecursiveAnonymizer):
+            def anonymize_text(self, data: str, **kwargs: Any) -> str:
+                calls_log.append(kwargs.get("field_name"))
+                return data
+
+        anonymizer = ParameterTrackingAnonymizer()
+
+        test_data = {"key1": "value1", "nested": {"key2": "value2"}}
+
+        # Call without field_name
+        anonymizer.anonymize(test_data)
+
+        # Should use empty string as a base and build paths from there
+        expected_field_names = [".key1", ".nested.key2"]
+        assert len(calls_log) == 2
+        for field_name in calls_log:
+            assert field_name in expected_field_names

sdks/python/tests/unit/anonymizer/test_rules_anonymizer.py

@@ -0,0 +1,318 @@
+import pytest
+
+from opik.anonymizer import rules
+from opik.anonymizer import rules_anonymizer
+
+
+class TestRulesAnonymizer:
+    """Test suite for RulesAnonymizer functionality."""
+
+    def test_init(self):
+        """Test RulesAnonymizer initialization."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        rules_ = [regex_rule]
+
+        anonymizer = rules_anonymizer.RulesAnonymizer(rules_, max_depth=5)
+
+        assert anonymizer.rules == rules_
+        assert anonymizer.max_depth == 5
+
+    def test_anonymize_text__single_regex_rule(self):
+        """Test anonymizing text with a single regex rule."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        result = anonymizer.anonymize_text("Phone: 123-456-7890")
+
+        assert result == "Phone: ***-***-***"
+
+    def test_anonymize_text__multiple_regex_rules(self):
+        """Test anonymizing text with multiple regex rules applied sequentially."""
+        email_rule = rules.RegexRule(
+            r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}", "[EMAIL]"
+        )
+        phone_rule = rules.RegexRule(r"\d{3}-\d{3}-\d{4}", "[PHONE]")
+
+        anonymizer = rules_anonymizer.RulesAnonymizer([email_rule, phone_rule])
+
+        result = anonymizer.anonymize_text("Contact: user@example.com or 123-456-7890")
+
+        assert result == "Contact: [EMAIL] or [PHONE]"
+
+    def test_anonymize_text__function_rule(self):
+        """Test anonymizing text with a function rule."""
+
+        def uppercase_anonymizer(text: str) -> str:
+            return text.upper()
+
+        function_rule = rules.FunctionRule(uppercase_anonymizer)
+        anonymizer = rules_anonymizer.RulesAnonymizer([function_rule])
+
+        result = anonymizer.anonymize_text("hello world")
+
+        assert result == "HELLO WORLD"
+
+    def test_anonymize_text__mixed_rules(self):
+        """Test anonymizing text with mixed rule types."""
+        regex_rule = rules.RegexRule(r"\d+", "XXX")
+        function_rule = rules.FunctionRule(lambda text: text.replace(" ", "_"))
+
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule, function_rule])
+
+        result = anonymizer.anonymize_text("Phone 123 456")
+
+        assert result == "Phone_XXX_XXX"
+
+    def test_anonymize_text__no_rules(self):
+        """Test anonymizing text with no rules returns original text."""
+        anonymizer = rules_anonymizer.RulesAnonymizer([])
+
+        result = anonymizer.anonymize_text("sensitive data")
+
+        assert result == "sensitive data"
+
+    def test_anonymize__string_data(self):
+        """Test anonymizing string data."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        result = anonymizer.anonymize("Phone: 123-456-7890")
+
+        assert result == "Phone: ***-***-***"
+
+    def test_anonymize__dict_data(self):
+        """Test anonymizing dictionary data."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        data = {"name": "John Doe", "phone": "123-456-7890", "address": "123 Main St"}
+
+        result = anonymizer.anonymize(data)
+
+        expected = {
+            "name": "John Doe",
+            "phone": "***-***-***",
+            "address": "*** Main St",
+        }
+        assert result == expected
+
+    def test_anonymize__list_data(self):
+        """Test anonymizing list data."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        data = ["Contact: 123-456-7890", "Another: 987-654-3210", "No numbers here"]
+
+        result = anonymizer.anonymize(data)
+
+        expected = ["Contact: ***-***-***", "Another: ***-***-***", "No numbers here"]
+        assert result == expected
+
+    def test_anonymize__nested_dict_data(self):
+        """Test anonymizing nested dictionary data."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        data = {
+            "user": {
+                "name": "John Doe",
+                "contact": {"phone": "123-456-7890", "email": "john@example.com"},
+            },
+            "id": "12345",
+        }
+
+        result = anonymizer.anonymize(data)
+
+        expected = {
+            "user": {
+                "name": "John Doe",
+                "contact": {"phone": "***-***-***", "email": "john@example.com"},
+            },
+            "id": "***",
+        }
+        assert result == expected
+
+    def test_anonymize__nested_list_data(self):
+        """Test anonymizing nested list data."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        data = [
+            ["Phone: 123", "Email: user@test.com"],
+            {"contact": "456-789-0123"},
+            "Simple string with 999",
+        ]
+
+        result = anonymizer.anonymize(data)
+
+        expected = [
+            ["Phone: ***", "Email: user@test.com"],
+            {"contact": "***-***-***"},
+            "Simple string with ***",
+        ]
+        assert result == expected
+
+    def test_anonymize__mixed_nested_data(self):
+        """Test anonymizing complex mixed nested data structures."""
+        email_rule = rules.RegexRule(
+            r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}", "[EMAIL]"
+        )
+        phone_rule = rules.RegexRule(r"\d{3}-\d{3}-\d{4}", "[PHONE]")
+
+        anonymizer = rules_anonymizer.RulesAnonymizer([email_rule, phone_rule])
+
+        data = {
+            "users": [
+                {"name": "John Doe", "contacts": ["john@example.com", "123-456-7890"]},
+                {"name": "Jane Smith", "contacts": ["jane@test.org", "987-654-3210"]},
+            ],
+            "admin_contact": "admin@company.com",
+        }
+
+        result = anonymizer.anonymize(data)
+
+        expected = {
+            "users": [
+                {"name": "John Doe", "contacts": ["[EMAIL]", "[PHONE]"]},
+                {"name": "Jane Smith", "contacts": ["[EMAIL]", "[PHONE]"]},
+            ],
+            "admin_contact": "[EMAIL]",
+        }
+        assert result == expected
+
+    def test_anonymize__non_string_values_unchanged(self):
+        """Test that non-string values are not processed."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        data = {
+            "number": 12345,
+            "float_val": 123.45,
+            "bool_val": True,
+            "none_val": None,
+            "string_val": "123",
+        }
+
+        result = anonymizer.anonymize(data)
+
+        expected = {
+            "number": 12345,
+            "float_val": 123.45,
+            "bool_val": True,
+            "none_val": None,
+            "string_val": "***",
+        }
+        assert result == expected
+
+    def test_anonymize__max_depth_limiting(self):
+        """Test that max_depth limits recursion depth."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule], max_depth=2)
+
+        # Create deeply nested data that exceeds max_depth
+        data = {"level1": {"level2": {"level3": {"phone": "123-456-7890"}}}}
+
+        result = anonymizer.anonymize(data)
+
+        # At max_depth=2, level3 should not be processed
+        expected = {
+            "level1": {
+                "level2": {
+                    "level3": {
+                        "phone": "123-456-7890"  # Should remain unchanged
+                    }
+                }
+            }
+        }
+        assert result == expected
+
+    def test_anonymize__max_depth_exact_limit(self):
+        """Test that max_depth allows processing exactly at the limit."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule], max_depth=4)
+
+        # Create data where string is processed at depth=3 (within limit)
+        data = {"level1": {"level2": {"phone": "123-456-7890"}}}
+
+        result = anonymizer.anonymize(data)
+
+        # At max_depth=4, the string at depth 3 should be processed
+        expected = {"level1": {"level2": {"phone": "***-***-***"}}}
+        assert result == expected
+
+    def test_anonymize__circular_reference_protection(self):
+        """Test protection against circular references through depth limiting."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule], max_depth=3)
+
+        # Create circular reference
+        data = {"phone": "123"}
+        data["self"] = data
+
+        # Should not cause infinite recursion due to max_depth
+        result = anonymizer.anonymize(data)
+
+        # The result should have processed the string but stopped at max depth
+        assert result["phone"] == "***"
+        assert "self" in result
+
+    @pytest.mark.parametrize("rule_count", [1, 5, 10])
+    def test_anonymize_text__performance_with_multiple_rules(self, rule_count):
+        """Test performance with varying numbers of rules."""
+        rules_ = []
+        for i in range(rule_count):
+            rules_.append(rules.RegexRule(f"pattern{i}", f"replacement{i}"))
+
+        anonymizer = rules_anonymizer.RulesAnonymizer(rules_)
+
+        # Test with a simple string
+        result = anonymizer.anonymize_text("test pattern0 and pattern5")
+
+        # Should apply all applicable rules
+        assert "replacement0" in result or "pattern0" in result
+
+    def test_anonymize_text__empty_string_handling(self):
+        """Test handling of empty strings."""
+        regex_rule = rules.RegexRule(r"\d+", "***")
+        anonymizer = rules_anonymizer.RulesAnonymizer([regex_rule])
+
+        result = anonymizer.anonymize_text("")
+
+        assert result == ""
+
+    def test_anonymize_text__complex_regex_patterns(self):
+        """Test complex regex patterns for real-world scenarios."""
+        # Email pattern
+        email_rule = rules.RegexRule(
+            r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", "[EMAIL_REDACTED]"
+        )
+
+        # Credit card pattern (simplified)
+        cc_rule = rules.RegexRule(
+            r"\b\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}\b", "[CARD_REDACTED]"
+        )
+
+        # SSN pattern
+        ssn_rule = rules.RegexRule(r"\b\d{3}-\d{2}-\d{4}\b", "[SSN_REDACTED]")
+
+        anonymizer = rules_anonymizer.RulesAnonymizer([email_rule, cc_rule, ssn_rule])
+
+        text = """
+        Contact info:
+        Email: john.doe@example.com
+        Credit Card: 1234 5678 9012 3456
+        SSN: 123-45-6789
+        """
+
+        result = anonymizer.anonymize_text(text)
+
+        assert result != text
+        assert (
+            result
+            == """
+        Contact info:
+        Email: [EMAIL_REDACTED]
+        Credit Card: [CARD_REDACTED]
+        SSN: [SSN_REDACTED]
+        """
+        )

sdks/python/tests/unit/api_key/test_opik_api_key.py

@@ -0,0 +1,102 @@
+from opik.api_key.opik_api_key import DELIMITER_CHAR, parse_api_key
+from opik.logging_messages import (
+    PARSE_API_KEY_EMPTY_EXPECTED_ATTRIBUTES,
+    PARSE_API_KEY_TOO_MANY_PARTS,
+)
+
+import pytest
+
+
+@pytest.mark.parametrize("raw_key", ["", None])
+def test_parse_api_key__empty_key(raw_key, capture_log):
+    comet_api_key = parse_api_key(raw_key)
+
+    assert comet_api_key is None
+
+
+def test_parse_api_key__one_part():
+    raw_key = "some API key"
+    comet_api_key = parse_api_key(raw_key)
+
+    assert comet_api_key is not None
+    assert comet_api_key.api_key == raw_key
+    assert comet_api_key.short_api_key == raw_key
+
+
+def test_parse_api_key__no_expected_attributes(capture_log):
+    raw_key = "some API key"
+    comet_api_key = parse_api_key(raw_key + DELIMITER_CHAR)
+
+    assert comet_api_key is not None
+    assert comet_api_key.api_key == raw_key
+    assert comet_api_key.short_api_key == raw_key
+
+    assert (
+        PARSE_API_KEY_EMPTY_EXPECTED_ATTRIBUTES % (raw_key + DELIMITER_CHAR)
+        in capture_log.messages
+    )
+
+
+def test_parse_api_key__too_many_parts(capture_log):
+    raw_key = "some API key" + DELIMITER_CHAR + "one" + DELIMITER_CHAR + "two"
+    comet_api_key = parse_api_key(raw_key)
+
+    assert comet_api_key is None
+    assert PARSE_API_KEY_TOO_MANY_PARTS % (3, raw_key) in capture_log.messages
+
+
+def test_parse_api_key__happy_path__with_padding():
+    # attributes: {"baseUrl": "https://www.comet.com"}
+    raw_key = (
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0="
+    )
+    comet_api_key = parse_api_key(raw_key)
+
+    assert comet_api_key is not None
+    assert comet_api_key.api_key == raw_key
+    assert comet_api_key.short_api_key == "Et1RBc4nd1ef3LfyJvhyB34Po"
+
+    assert comet_api_key.base_url == "https://www.comet.com"
+    assert comet_api_key["baseUrl"] == "https://www.comet.com"
+
+
+def test_parse_api_key__happy_path__no_padding():
+    # attributes: {"baseUrl": "https://www.comet.com"}
+    raw_key = (
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0"
+    )
+    comet_api_key = parse_api_key(raw_key)
+
+    assert comet_api_key is not None
+    assert comet_api_key.api_key == raw_key
+    assert comet_api_key.short_api_key == "Et1RBc4nd1ef3LfyJvhyB34Po"
+
+    assert comet_api_key.base_url == "https://www.comet.com"
+    assert comet_api_key["baseUrl"] == "https://www.comet.com"
+
+
+@pytest.mark.parametrize(
+    "raw_key",
+    [
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0===",
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0==",
+    ],
+)
+def test_parse_api_key__happy_path__wrong_padding(raw_key):
+    # attributes: {"baseUrl": "https://www.comet.com"}
+    comet_api_key = parse_api_key(raw_key)
+
+    assert comet_api_key is not None
+    assert comet_api_key.api_key == raw_key
+    assert comet_api_key.short_api_key == "Et1RBc4nd1ef3LfyJvhyB34Po"
+
+    assert comet_api_key.base_url == "https://www.comet.com"
+    assert comet_api_key["baseUrl"] == "https://www.comet.com"

sdks/python/tests/unit/api_objects/agent_config/conftest.py

@@ -0,0 +1,62 @@
+from unittest import mock
+
+import pytest
+
+from opik.api_objects import opik_client
+from opik.rest_api import core as rest_api_core
+from opik.rest_api.types.agent_blueprint_public import AgentBlueprintPublic
+from opik.decorator.context_manager import start_as_current_trace
+
+
+def make_raw_blueprint(blueprint_id="bp-1", name=None, values=None, description=None):
+    if values is None:
+        values = []
+    return AgentBlueprintPublic(
+        id=blueprint_id,
+        name=name,
+        type="blueprint",
+        values=values,
+        description=description,
+    )
+
+
+@pytest.fixture
+def mock_rest_client():
+    client = mock.Mock()
+    client.agent_configs = mock.Mock()
+    client.agent_configs.create_agent_config.return_value = None
+    client.agent_configs.get_latest_blueprint.side_effect = rest_api_core.ApiError(
+        status_code=404, body="not found"
+    )
+    client.agent_configs.get_blueprint_by_env.side_effect = rest_api_core.ApiError(
+        status_code=404, body="not found"
+    )
+    client.agent_configs.get_blueprint_by_id.return_value = make_raw_blueprint()
+    client.projects.retrieve_project.return_value = mock.Mock(id="proj-test")
+    return client
+
+
+@pytest.fixture
+def mock_opik_client(mock_rest_client):
+    client = opik_client.Opik.__new__(opik_client.Opik)
+    client._rest_client = mock_rest_client
+    client._project_name = "test-project"
+    return client
+
+
+@pytest.fixture(autouse=True)
+def clear_caches():
+    yield
+    from opik.api_objects.agent_config.cache import get_global_registry
+
+    get_global_registry().clear()
+
+
+@pytest.fixture(autouse=True)
+def fake_track_context():
+    """Push a fake trace so _get_or_create_from_backend's @track guard passes in all unit tests."""
+    with mock.patch.object(
+        opik_client, "get_global_client", return_value=mock.Mock(spec=opik_client.Opik)
+    ):
+        with start_as_current_trace(name="test-trace"):
+            yield

sdks/python/tests/unit/api_objects/agent_config/test_blueprint.py

@@ -0,0 +1,470 @@
+from typing import List
+from unittest import mock
+
+import pytest
+
+from opik.api_objects.agent_config.blueprint import Blueprint
+from opik.api_objects.prompt.text.prompt import Prompt
+from opik.api_objects.prompt.chat.chat_prompt import ChatPrompt
+from opik.api_objects.prompt.base_prompt import BasePrompt
+from opik.rest_api.types.agent_blueprint_public import AgentBlueprintPublic
+from opik.rest_api.types.agent_config_value_public import AgentConfigValuePublic
+from opik.rest_api.types.prompt_version_detail import PromptVersionDetail
+
+
+def _make_raw_blueprint(
+    blueprint_id="bp-1",
+    values=None,
+    description=None,
+    bp_type="blueprint",
+    envs=None,
+    created_by=None,
+    created_at=None,
+):
+    if values is None:
+        values = [
+            AgentConfigValuePublic(key="temperature", type="float", value="0.6"),
+            AgentConfigValuePublic(key="name", type="string", value="agent"),
+        ]
+    return AgentBlueprintPublic(
+        id=blueprint_id,
+        type=bp_type,
+        values=values,
+        description=description,
+        envs=envs,
+        created_by=created_by,
+        created_at=created_at,
+    )
+
+
+class TestBlueprintProperties:
+    @pytest.mark.parametrize(
+        "kwargs,attr,expected",
+        [
+            ({"blueprint_id": "bp-42"}, "id", "bp-42"),
+            ({"description": "test desc"}, "description", "test desc"),
+            ({"bp_type": "mask"}, "type", "mask"),
+            ({"envs": ["prod", "staging"]}, "envs", ["prod", "staging"]),
+            ({"created_by": "user-1"}, "created_by", "user-1"),
+        ],
+    )
+    def test_property__returns_value(self, kwargs, attr, expected):
+        bp = Blueprint(_make_raw_blueprint(**kwargs))
+        assert getattr(bp, attr) == expected
+
+
+class TestBlueprintValueResolution:
+    def test_without_field_types__values_are_inferred_from_backend_type(self):
+        bp = Blueprint(_make_raw_blueprint())
+        assert bp["temperature"] == 0.6
+        assert isinstance(bp["temperature"], float)
+        assert bp["name"] == "agent"
+        assert isinstance(bp["name"], str)
+
+    def test_with_field_types__deserializes_values(self):
+        bp = Blueprint(
+            _make_raw_blueprint(),
+            field_types={"temperature": float, "name": str},
+        )
+        assert bp["temperature"] == 0.6
+        assert bp["name"] == "agent"
+
+    @pytest.mark.parametrize(
+        "backend_type,value_str,py_type,expected",
+        [
+            ("string", "true", bool, True),
+            ("integer", "42", int, 42),
+        ],
+    )
+    def test_with_field_types__type_deserialization(
+        self, backend_type, value_str, py_type, expected
+    ):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="val", type=backend_type, value=value_str)
+            ]
+        )
+        bp = Blueprint(raw, field_types={"val": py_type})
+        assert bp["val"] == expected
+
+
+class TestBlueprintDictLikeAccess:
+    def test_get__existing_key(self):
+        bp = Blueprint(_make_raw_blueprint())
+        assert bp.get("temperature") == 0.6
+
+    def test_get__missing_key__returns_default(self):
+        bp = Blueprint(_make_raw_blueprint())
+        assert bp.get("missing", 42) == 42
+
+    def test_get__missing_key_no_default__returns_none(self):
+        bp = Blueprint(_make_raw_blueprint())
+        assert bp.get("missing") is None
+
+    def test_getitem__existing_key(self):
+        bp = Blueprint(_make_raw_blueprint())
+        assert bp["name"] == "agent"
+
+    def test_getitem__missing_key__raises_key_error(self):
+        bp = Blueprint(_make_raw_blueprint())
+        with pytest.raises(KeyError):
+            _ = bp["missing"]
+
+    def test_keys__returns_all_keys(self):
+        bp = Blueprint(_make_raw_blueprint())
+        assert set(bp.keys()) == {"temperature", "name"}
+
+    def test_values__returns_deep_copy(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="items", type="string", value="[1, 2, 3]"),
+            ]
+        )
+        bp = Blueprint(raw, field_types={"items": List[int]})
+        vals = bp.values
+        vals["items"].append(4)
+        assert bp.values["items"] == [1, 2, 3]
+
+
+class TestBlueprintPromptResolution:
+    def test_prompt_field__resolves_to_prompt_object(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="system_prompt", type="prompt", value="abc12345"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        version_detail.template_structure = "text"
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "my-prompt"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_prompt = mock.Mock(spec=Prompt)
+        with mock.patch(
+            "opik.api_objects.prompt.text.prompt.Prompt.from_fern_prompt_version",
+            return_value=fake_prompt,
+        ):
+            bp = Blueprint(
+                raw,
+                field_types={"system_prompt": Prompt},
+                rest_client_=mock_rest,
+            )
+
+        assert bp["system_prompt"] is fake_prompt
+        mock_rest.prompts.get_prompt_by_commit.assert_called_once_with("abc12345")
+
+    def test_chat_prompt_field__resolves_to_chat_prompt_object(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="messages", type="prompt", value="bcd23456"),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        version_detail.type = "mustache"
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "chat-prompt"
+        prompt_detail.template_structure = "chat"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_chat_prompt = mock.Mock(spec=ChatPrompt)
+        with mock.patch(
+            "opik.api_objects.prompt.chat.chat_prompt.ChatPrompt.from_fern_prompt_version",
+            return_value=fake_chat_prompt,
+        ):
+            bp = Blueprint(
+                raw,
+                field_types={"messages": ChatPrompt},
+                rest_client_=mock_rest,
+            )
+
+        assert bp["messages"] is fake_chat_prompt
+
+    def test_base_prompt_annotation__chat_structure__resolves_to_chat_prompt(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="p", type="prompt", value="cde34567"),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        version_detail.type = "mustache"
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "any-prompt"
+        prompt_detail.template_structure = "chat"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_chat_prompt = mock.Mock(spec=ChatPrompt)
+        with mock.patch(
+            "opik.api_objects.prompt.chat.chat_prompt.ChatPrompt.from_fern_prompt_version",
+            return_value=fake_chat_prompt,
+        ):
+            bp = Blueprint(
+                raw,
+                field_types={"p": BasePrompt},
+                rest_client_=mock_rest,
+            )
+
+        assert bp["p"] is fake_chat_prompt
+
+    def test_base_prompt_annotation__text_structure__resolves_to_prompt(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="p", type="prompt", value="def45678"),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        version_detail.template_structure = "text"
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "any-prompt"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_prompt = mock.Mock(spec=Prompt)
+        with mock.patch(
+            "opik.api_objects.prompt.text.prompt.Prompt.from_fern_prompt_version",
+            return_value=fake_prompt,
+        ):
+            bp = Blueprint(
+                raw,
+                field_types={"p": BasePrompt},
+                rest_client_=mock_rest,
+            )
+
+        assert bp["p"] is fake_prompt
+
+    def test_prompt_resolution_fails__raises(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="system_prompt", type="prompt", value="badbad00"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        mock_rest.prompts.get_prompt_by_commit.side_effect = Exception("network error")
+
+        with pytest.raises(Exception, match="network error"):
+            Blueprint(
+                raw,
+                field_types={"system_prompt": Prompt},
+                rest_client_=mock_rest,
+            )
+
+    def test_prompt_version_field__resolves_to_prompt_version_detail(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="version", type="prompt_commit", value="pv111111"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        fake_version_detail = mock.Mock(spec=PromptVersionDetail)
+        prompt_detail = mock.Mock()
+        prompt_detail.requested_version = fake_version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        bp = Blueprint(
+            raw,
+            field_types={"version": PromptVersionDetail},
+            rest_client_=mock_rest,
+        )
+
+        assert bp["version"] is fake_version_detail
+        mock_rest.prompts.get_prompt_by_commit.assert_called_once_with("pv111111")
+
+    def test_prompt_version_field__resolution_fails__raises(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="version", type="prompt_commit", value="badbad00"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        mock_rest.prompts.get_prompt_by_commit.side_effect = Exception("not found")
+
+        with pytest.raises(Exception, match="not found"):
+            Blueprint(
+                raw,
+                field_types={"version": PromptVersionDetail},
+                rest_client_=mock_rest,
+            )
+
+    def test_prompt_field_type_declared_as_chatprompt__chat_template_structure__resolves_to_chat_prompt(
+        self,
+    ):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="p", type="prompt", value="aaa11111"),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "my-prompt"
+        prompt_detail.template_structure = "chat"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_chat_prompt = mock.Mock(spec=ChatPrompt)
+        with mock.patch(
+            "opik.api_objects.prompt.chat.chat_prompt.ChatPrompt.from_fern_prompt_version",
+            return_value=fake_chat_prompt,
+        ):
+            bp = Blueprint(
+                raw,
+                field_types={"p": ChatPrompt},
+                rest_client_=mock_rest,
+            )
+
+        assert bp["p"] is fake_chat_prompt
+        assert isinstance(bp["p"], ChatPrompt)
+
+    def test_two_prompt_fields__makes_exactly_one_api_call_per_prompt(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="p1", type="prompt", value="aaaaaaaa"),
+                AgentConfigValuePublic(key="p2", type="prompt", value="bbbbbbbb"),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+
+        def _commit_side_effect(commit):
+            v = mock.Mock()
+            d = mock.Mock()
+            d.name = f"name-{commit}"
+            d.template_structure = "text"
+            d.requested_version = v
+            return d
+
+        mock_rest.prompts.get_prompt_by_commit.side_effect = _commit_side_effect
+
+        with mock.patch(
+            "opik.api_objects.prompt.text.prompt.Prompt.from_fern_prompt_version",
+            return_value=mock.Mock(spec=Prompt),
+        ):
+            Blueprint(
+                raw,
+                field_types={"p1": Prompt, "p2": Prompt},
+                rest_client_=mock_rest,
+            )
+
+        assert mock_rest.prompts.get_prompt_by_commit.call_count == 2
+
+
+class TestBlueprintPromptResolutionWithoutFieldTypes:
+    def test_prompt_type__resolves_to_prompt_object(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="system_prompt", type="prompt", value="abc12345"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "my-prompt"
+        prompt_detail.template_structure = "text"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_prompt = mock.Mock(spec=Prompt)
+        with mock.patch(
+            "opik.api_objects.prompt.text.prompt.Prompt.from_fern_prompt_version",
+            return_value=fake_prompt,
+        ):
+            bp = Blueprint(raw, rest_client_=mock_rest)
+
+        assert bp["system_prompt"] is fake_prompt
+
+    def test_chat_prompt_type__resolves_to_chat_prompt_object(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(key="messages", type="prompt", value="bcd23456"),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        version_detail = mock.Mock()
+        prompt_detail = mock.Mock()
+        prompt_detail.name = "chat-prompt"
+        prompt_detail.template_structure = "chat"
+        prompt_detail.requested_version = version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        fake_chat_prompt = mock.Mock(spec=ChatPrompt)
+        with mock.patch(
+            "opik.api_objects.prompt.chat.chat_prompt.ChatPrompt.from_fern_prompt_version",
+            return_value=fake_chat_prompt,
+        ):
+            bp = Blueprint(raw, rest_client_=mock_rest)
+
+        assert bp["messages"] is fake_chat_prompt
+
+    def test_prompt_commit_type__resolves_to_prompt_version_detail(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="version", type="prompt_commit", value="pv111111"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        fake_version_detail = mock.Mock(spec=PromptVersionDetail)
+        prompt_detail = mock.Mock()
+        prompt_detail.requested_version = fake_version_detail
+        mock_rest.prompts.get_prompt_by_commit.return_value = prompt_detail
+
+        bp = Blueprint(raw, rest_client_=mock_rest)
+
+        assert bp["version"] is fake_version_detail
+        mock_rest.prompts.get_prompt_by_commit.assert_called_once_with("pv111111")
+
+    def test_prompt_resolution_fails__raises(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="system_prompt", type="prompt", value="badbad00"
+                ),
+            ]
+        )
+
+        mock_rest = mock.Mock()
+        mock_rest.prompts.get_prompt_by_commit.side_effect = Exception("network error")
+
+        with pytest.raises(Exception, match="network error"):
+            Blueprint(raw, rest_client_=mock_rest)
+
+    def test_without_rest_client__prompt_stays_raw_string(self):
+        raw = _make_raw_blueprint(
+            values=[
+                AgentConfigValuePublic(
+                    key="system_prompt", type="prompt", value="ver-111"
+                ),
+            ]
+        )
+
+        bp = Blueprint(raw)
+
+        assert bp["system_prompt"] == "ver-111"

sdks/python/tests/unit/api_objects/agent_config/test_cache.py

@@ -0,0 +1,357 @@
+import threading
+import time
+from unittest import mock
+
+import pytest
+
+from opik.api_objects.agent_config.cache import (
+    CacheRefreshThread,
+    SharedCacheRegistry,
+    SharedConfigCache,
+)
+
+
+@pytest.fixture
+def registry():
+    r = SharedCacheRegistry()
+    yield r
+    r.clear()
+
+
+class TestSharedCacheRegistry:
+    def test_get__same_key__returns_same_instance(self, registry):
+        a = registry.get("proj", None, None)
+        b = registry.get("proj", None, None)
+        assert a is b
+
+    def test_get__different_key__returns_different_instance(self, registry):
+        a = registry.get("proj-a", None, None)
+        b = registry.get("proj-b", None, None)
+        assert a is not b
+
+    def test_clear__empties_registry(self, registry):
+        registry.get("proj", None, None)
+        registry.clear()
+        # After clear, a new call returns a fresh instance
+        fresh = registry.get("proj", None, None)
+        assert fresh.blueprint_id is None
+
+    def test_clear__stops_thread(self, registry):
+        registry.ensure_refresh_thread_started()
+        assert registry._thread is not None and registry._thread.is_alive()
+        registry.clear()
+        assert registry._thread is None
+
+    def test_ensure_refresh_thread_started__starts_thread(self, registry):
+        assert registry._thread is None
+        registry.ensure_refresh_thread_started()
+        assert registry._thread is not None
+        assert registry._thread.is_alive()
+
+    def test_ensure_refresh_thread_started__second_call_noop(self, registry):
+        registry.ensure_refresh_thread_started()
+        first_thread = registry._thread
+        registry.ensure_refresh_thread_started()
+        assert registry._thread is first_thread
+
+    def test_stop_refresh_thread__stops_and_nulls(self, registry):
+        registry.ensure_refresh_thread_started()
+        thread = registry._thread
+        registry.stop_refresh_thread()
+        thread.join(timeout=2)
+        assert not thread.is_alive()
+        assert registry._thread is None
+
+    def test_concurrent_get__returns_same_instance(self, registry):
+        results = [None] * 10
+        barrier = threading.Barrier(10)
+
+        def fetch(idx):
+            barrier.wait()
+            results[idx] = registry.get("proj", None, None)
+
+        threads = [threading.Thread(target=fetch, args=(i,)) for i in range(10)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        assert all(r is results[0] for r in results)
+
+
+class TestSharedConfigCacheThreadSafety:
+    def test_apply__concurrent_reads_see_consistent_dict(self):
+        cache = SharedConfigCache(ttl_seconds=300)
+        errors = []
+        stop = threading.Event()
+
+        def writer():
+            for i in range(100):
+                bp = mock.Mock()
+                bp.id = f"bp-{i}"
+                bp._values = {f"key-{j}": f"val-{i}" for j in range(5)}
+                cache.update(bp)
+
+        def reader():
+            while not stop.is_set():
+                vals = cache.values
+                unique_vals = set(vals.values())
+                if len(unique_vals) > 1:
+                    errors.append(f"Inconsistent values: {vals}")
+                    break
+
+        writer_thread = threading.Thread(target=writer)
+        reader_thread = threading.Thread(target=reader)
+        reader_thread.start()
+        writer_thread.start()
+        writer_thread.join()
+        stop.set()
+        reader_thread.join()
+
+        assert errors == [], f"Found inconsistencies: {errors}"
+
+    def test_register_fields__concurrent__all_fields_present(self):
+        cache = SharedConfigCache()
+        barrier = threading.Barrier(4)
+
+        def register(prefix: str):
+            barrier.wait()
+            cache.register_fields({f"{prefix}.f1": str, f"{prefix}.f2": int})
+
+        threads = [
+            threading.Thread(target=register, args=(f"Class{i}",)) for i in range(4)
+        ]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join()
+
+        all_types = cache.all_field_types
+        assert len(all_types) == 8
+        for i in range(4):
+            assert f"Class{i}.f1" in all_types
+            assert f"Class{i}.f2" in all_types
+
+
+class TestRefreshCallback:
+    def test_set_refresh_callback__first_writer_wins(self):
+        cache = SharedConfigCache()
+        cb1 = mock.Mock()
+        cb2 = mock.Mock()
+
+        cache.set_refresh_callback(cb1)
+        cache.set_refresh_callback(cb2)
+
+        cache.try_background_refresh()
+        cb1.assert_called_once()
+        cb2.assert_not_called()
+
+    def test_try_background_refresh__no_callback__noop(self):
+        cache = SharedConfigCache()
+        cache.try_background_refresh()
+
+    def test_try_background_refresh__callback_returns_blueprint__applies(self):
+        cache = SharedConfigCache(ttl_seconds=300)
+        bp = mock.Mock()
+        bp.id = "bp-new"
+        bp._values = {"A.x": 42}
+
+        cache.set_refresh_callback(lambda: bp)
+        cache.try_background_refresh()
+
+        assert cache.blueprint_id == "bp-new"
+        assert cache.values == {"A.x": 42}
+        assert not cache.is_stale()
+
+    def test_try_background_refresh__callback_returns_none__no_change(self):
+        cache = SharedConfigCache()
+        cache.set_refresh_callback(lambda: None)
+        cache.try_background_refresh()
+
+        assert cache.blueprint_id is None
+        assert cache.values == {}
+
+    def test_try_background_refresh__callback_raises__no_crash(self):
+        cache = SharedConfigCache()
+        cache.set_refresh_callback(mock.Mock(side_effect=RuntimeError("boom")))
+        cache.try_background_refresh()
+
+        assert cache.blueprint_id is None
+
+
+class TestRefreshPolicy:
+    """Verify which cache lookups get a background refresh callback and which do not.
+
+    Tests use a local SharedCacheRegistry to avoid interference with the global
+    singleton used by init_cache_entry.
+    """
+
+    def _make_bp(self, bp_id: str, values: dict) -> mock.Mock:
+        bp = mock.Mock()
+        bp.id = bp_id
+        bp._values = values
+        return bp
+
+    def test_latest_lookup__refresh_callback_registered(self, registry):
+        manager = mock.Mock()
+        bp = self._make_bp("bp-latest", {"K.v": "v1"})
+        cache = registry.get("proj", None, None, None)
+        cache.update(bp)
+        cache.set_refresh_callback(
+            lambda: manager.get_blueprint(env=None, mask_id=None, field_types={})
+        )
+        registry.ensure_refresh_thread_started()
+        assert cache._refresh_callback is not None
+
+    def test_env_lookup__refresh_callback_registered(self, registry):
+        manager = mock.Mock()
+        bp = self._make_bp("bp-env", {"K.v": "v1"})
+        cache = registry.get("proj", "prod", None, None)
+        cache.update(bp)
+        cache.set_refresh_callback(
+            lambda: manager.get_blueprint(env="prod", mask_id=None, field_types={})
+        )
+        registry.ensure_refresh_thread_started()
+        assert cache._refresh_callback is not None
+
+    def test_version_lookup__no_refresh_callback(self, registry):
+        bp = self._make_bp("bp-v1", {"K.v": "v1"})
+        cache = registry.get("proj", None, None, "v1")
+        cache.update(bp)
+        # version-pinned: no refresh callback registered
+        assert cache._refresh_callback is None
+
+    def test_masked_lookup__no_refresh_callback(self, registry):
+        bp = self._make_bp("bp-masked", {"K.v": "v1"})
+        cache = registry.get("proj", None, "mask-abc", None)
+        cache.update(bp)
+        # masked: no refresh callback registered
+        assert cache._refresh_callback is None
+
+    def test_latest_and_version__separate_cache_entries(self, registry):
+        bp_latest = self._make_bp("bp-latest", {"K.v": "latest"})
+        bp_v1 = self._make_bp("bp-v1", {"K.v": "v1"})
+
+        latest_cache = registry.get("proj", None, None, None)
+        latest_cache.update(bp_latest)
+
+        version_cache = registry.get("proj", None, None, "v1")
+        version_cache.update(bp_v1)
+
+        assert latest_cache is not version_cache
+        assert latest_cache.blueprint_id == "bp-latest"
+        assert version_cache.blueprint_id == "bp-v1"
+
+    def test_latest__background_refresh_updates_cache(self, registry):
+        bp_new = self._make_bp("bp-refreshed", {"K.v": "refreshed"})
+        cache = registry.get("proj", None, None, None)
+        cache._ttl_seconds = 0
+        cache.set_refresh_callback(lambda: bp_new)
+        cache.try_background_refresh()
+        assert cache.blueprint_id == "bp-refreshed"
+        assert cache.values == {"K.v": "refreshed"}
+
+    def test_version__no_background_refresh_even_when_stale(self):
+        bp = self._make_bp("bp-v1", {"K.v": "v1"})
+        cache = SharedConfigCache(ttl_seconds=0)
+        cache.update(bp)
+        # No refresh callback registered for version-pinned cache
+        assert cache._refresh_callback is None
+        cache.try_background_refresh()
+        assert cache.blueprint_id == "bp-v1"
+
+    def test_init_cache_entry__latest__registers_refresh(self, registry):
+        """init_cache_entry with version=None, mask_id=None must register a refresh callback."""
+        import opik.api_objects.agent_config.cache as cache_mod
+
+        manager = mock.Mock()
+        bp = self._make_bp("bp-latest", {"K.v": "v1"})
+        with mock.patch.object(cache_mod, "_registry", registry):
+            cache_mod.init_cache_entry(
+                "proj", None, None, {}, manager, blueprint=bp, version=None
+            )
+        cache = registry.get("proj", None, None, None)
+        assert cache._refresh_callback is not None
+
+    def test_init_cache_entry__version__no_refresh(self, registry):
+        """init_cache_entry with version set must NOT register a refresh callback."""
+        import opik.api_objects.agent_config.cache as cache_mod
+
+        manager = mock.Mock()
+        bp = self._make_bp("bp-v1", {"K.v": "v1"})
+        with mock.patch.object(cache_mod, "_registry", registry):
+            cache_mod.init_cache_entry(
+                "proj", None, None, {}, manager, blueprint=bp, version="v1"
+            )
+        cache = registry.get("proj", None, None, "v1")
+        assert cache._refresh_callback is None
+
+    def test_init_cache_entry__masked__no_refresh(self, registry):
+        """init_cache_entry with mask_id set must NOT register a refresh callback."""
+        import opik.api_objects.agent_config.cache as cache_mod
+
+        manager = mock.Mock()
+        bp = self._make_bp("bp-masked", {"K.v": "v1"})
+        with mock.patch.object(cache_mod, "_registry", registry):
+            cache_mod.init_cache_entry(
+                "proj", None, "mask-abc", {}, manager, blueprint=bp, version=None
+            )
+        cache = registry.get("proj", None, "mask-abc", None)
+        assert cache._refresh_callback is None
+
+
+class TestCacheRefreshThread:
+    def test_stops_on_close(self):
+        thread = CacheRefreshThread(get_caches=list, interval_seconds=0.01)
+        thread.start()
+        assert thread.is_alive()
+
+        thread.close()
+        thread.join(timeout=2)
+        assert not thread.is_alive()
+
+    def test_refreshes_stale_cache(self, registry):
+        bp = mock.Mock()
+        bp.id = "bp-bg"
+        bp._values = {"K.v": "refreshed"}
+
+        callback = mock.Mock(return_value=bp)
+
+        cache = registry.get("bg-proj", None, None)
+        cache._ttl_seconds = 0
+        cache.set_refresh_callback(callback)
+
+        thread = CacheRefreshThread(
+            get_caches=lambda: list(registry._caches.values()),
+            interval_seconds=0.05,
+        )
+        thread.start()
+        try:
+            time.sleep(0.3)
+            assert callback.call_count >= 1
+            assert cache.values == {"K.v": "refreshed"}
+        finally:
+            thread.close()
+            thread.join(timeout=2)
+
+    def test_skips_non_stale_caches(self, registry):
+        bp = mock.Mock()
+        bp.id = "bp-init"
+        bp._values = {"K.v": "initial"}
+
+        cache = registry.get("fresh-proj", None, None)
+        cache.update(bp)
+        callback = mock.Mock()
+        cache.set_refresh_callback(callback)
+
+        thread = CacheRefreshThread(
+            get_caches=lambda: list(registry._caches.values()),
+            interval_seconds=0.05,
+        )
+        thread.start()
+        try:
+            time.sleep(0.2)
+            callback.assert_not_called()
+        finally:
+            thread.close()
+            thread.join(timeout=2)

sdks/python/tests/unit/api_objects/agent_config/test_config.py

@@ -0,0 +1,406 @@
+from unittest import mock
+
+import pytest
+
+from opik.api_objects.agent_config import types
+from opik.api_objects.agent_config.config import ConfigManager
+from opik.api_objects.agent_config.blueprint import Blueprint
+from opik.rest_api import core as rest_api_core
+from opik.rest_api.types.agent_blueprint_public import AgentBlueprintPublic
+from opik.rest_api.types.agent_config_value_public import AgentConfigValuePublic
+
+
+def _make_raw_blueprint(blueprint_id="bp-1", values=None, description=None):
+    if values is None:
+        values = [
+            AgentConfigValuePublic(key="temp", type="float", value="0.6"),
+            AgentConfigValuePublic(key="name", type="string", value="agent"),
+        ]
+    return AgentBlueprintPublic(
+        id=blueprint_id, type="blueprint", values=values, description=description
+    )
+
+
+@pytest.fixture
+def mock_rest_client():
+    client = mock.Mock()
+    client.agent_configs = mock.Mock()
+    client.agent_configs.create_agent_config.return_value = None
+    client.agent_configs.get_latest_blueprint.return_value = _make_raw_blueprint()
+    client.projects.retrieve_project.return_value = mock.Mock(id="proj-default")
+    return client
+
+
+@pytest.fixture
+def agent_config(mock_rest_client):
+    return ConfigManager(
+        project_name="my-project",
+        rest_client_=mock_rest_client,
+    )
+
+
+class TestConfigManagerProperties:
+    def test_project_name(self, agent_config):
+        assert agent_config.project_name == "my-project"
+
+
+class TestConfigManagerGetBlueprint:
+    def test_get_blueprint__returns_blueprint(self, agent_config, mock_rest_client):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_latest_blueprint.return_value = (
+            _make_raw_blueprint()
+        )
+
+        result = agent_config.get_blueprint()
+
+        assert isinstance(result, Blueprint)
+        assert result.id == "bp-1"
+
+    def test_get_blueprint__with_env__routes_to_get_blueprint_by_env(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_blueprint_by_env.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.get_blueprint(env="prod")
+
+        mock_rest_client.agent_configs.get_blueprint_by_env.assert_called_once_with(
+            env_name="prod",
+            project_id="proj-1",
+            mask_id=None,
+            request_options=None,
+        )
+        mock_rest_client.agent_configs.get_latest_blueprint.assert_not_called()
+
+    def test_get_blueprint__with_mask_id__passes_mask_id(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_latest_blueprint.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.get_blueprint(mask_id="mask-1")
+
+        mock_rest_client.agent_configs.get_latest_blueprint.assert_called_once_with(
+            project_id="proj-1",
+            mask_id="mask-1",
+            request_options=None,
+        )
+
+    def test_get_blueprint__with_field_types__resolves_values(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_latest_blueprint.return_value = (
+            _make_raw_blueprint()
+        )
+
+        result = agent_config.get_blueprint(field_types={"temp": float, "name": str})
+
+        assert result["temp"] == 0.6
+        assert result["name"] == "agent"
+
+    def test_get_blueprint__without_field_types__infers_types_from_backend(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_latest_blueprint.return_value = (
+            _make_raw_blueprint()
+        )
+
+        result = agent_config.get_blueprint()
+
+        assert result["temp"] == 0.6
+        assert isinstance(result["temp"], float)
+        assert result["name"] == "agent"
+        assert isinstance(result["name"], str)
+
+    def test_get_blueprint__not_found__returns_none(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_latest_blueprint.side_effect = (
+            rest_api_core.ApiError(status_code=404, body="not found")
+        )
+
+        result = agent_config.get_blueprint()
+
+        assert result is None
+
+    @pytest.mark.parametrize(
+        "mask_id",
+        ["mask-1", "mask-2", None],
+        ids=["mask_1", "mask_2", "no_mask"],
+    )
+    def test_get_blueprint__mask_id__passed_to_backend(
+        self, agent_config, mock_rest_client, mask_id
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_latest_blueprint.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.get_blueprint(mask_id=mask_id)
+
+        mock_rest_client.agent_configs.get_latest_blueprint.assert_called_once_with(
+            project_id="proj-1",
+            mask_id=mask_id,
+            request_options=None,
+        )
+
+    def test_get_blueprint__env_with_mask_id__routes_to_get_blueprint_by_env(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_blueprint_by_env.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.get_blueprint(env="prod", mask_id="mask-1")
+
+        mock_rest_client.agent_configs.get_blueprint_by_env.assert_called_once_with(
+            env_name="prod",
+            project_id="proj-1",
+            mask_id="mask-1",
+            request_options=None,
+        )
+
+
+class TestConfigManagerGetBlueprintByName:
+    def test_get_blueprint_by_name__returns_blueprint(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_blueprint_by_name.return_value = (
+            _make_raw_blueprint(blueprint_id="bp-specific")
+        )
+
+        result = agent_config.get_blueprint(name="v1")
+
+        assert isinstance(result, Blueprint)
+        assert result.id == "bp-specific"
+        mock_rest_client.agent_configs.get_blueprint_by_name.assert_called_once_with(
+            project_id="proj-1", name="v1", mask_id=None, request_options=None
+        )
+
+    def test_get_blueprint_by_name__not_found__returns_none(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.projects.retrieve_project.return_value = mock.Mock(id="proj-1")
+        mock_rest_client.agent_configs.get_blueprint_by_name.side_effect = (
+            rest_api_core.ApiError(status_code=404, body="not found")
+        )
+
+        result = agent_config.get_blueprint(name="nonexistent")
+
+        assert result is None
+
+
+class TestConfigManagerCreateBlueprint:
+    def test_create_blueprint__happy_path__calls_backend_and_returns_blueprint(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint(blueprint_id="bp-new")
+        )
+
+        result = agent_config.create_blueprint(
+            fields_with_values={
+                "temperature": types.FieldValueSpec(float, 0.6),
+                "name": types.FieldValueSpec(str, "agent"),
+            }
+        )
+
+        mock_rest_client.agent_configs.create_agent_config.assert_called_once()
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        blueprint = call_kwargs["blueprint"]
+        assert blueprint.type == "blueprint"
+        assert blueprint.values is not None
+        assert blueprint.id is not None
+        assert isinstance(result, Blueprint)
+
+    def test_create_blueprint__bool_field__serialized_as_boolean_type(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.create_blueprint(
+            fields_with_values={"flag": types.FieldValueSpec(bool, False)}
+        )
+
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        blueprint = call_kwargs["blueprint"]
+        flag_param = [v for v in blueprint.values if v.key == "flag"][0]
+        assert flag_param.type == "boolean"
+        assert flag_param.value == "false"
+
+    def test_create_blueprint__with_project_name__passes_project_to_backend(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.create_blueprint(
+            fields_with_values={"temp": types.FieldValueSpec(float, 0.5)}
+        )
+
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        assert call_kwargs["project_name"] == "my-project"
+
+    def test_create_blueprint__with_parameters__returns_blueprint(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint(blueprint_id="bp-new")
+        )
+
+        result = agent_config.create_blueprint(
+            parameters={"temp": 0.6, "name": "agent"}
+        )
+
+        assert isinstance(result, Blueprint)
+        assert result.id == "bp-new"
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        bp = call_kwargs["blueprint"]
+        keys = {v.key for v in bp.values}
+        assert "temp" in keys
+        assert "name" in keys
+
+    def test_create_blueprint__with_fields_with_values__returns_blueprint(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint()
+        )
+
+        result = agent_config.create_blueprint(
+            fields_with_values={"temp": types.FieldValueSpec(float, 0.6)}
+        )
+
+        assert isinstance(result, Blueprint)
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        keys = {v.key for v in call_kwargs["blueprint"].values}
+        assert "temp" in keys
+
+    def test_create_blueprint__with_description__passes_description(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.create_blueprint(parameters={"temp": 0.6}, description="v1")
+
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        assert call_kwargs["blueprint"].description == "v1"
+
+
+class TestConfigManagerCreateMask:
+    def test_create_mask__happy_path__calls_backend_with_mask_type(
+        self, agent_config, mock_rest_client
+    ):
+        agent_config.create_mask(
+            fields_with_values={"temperature": types.FieldValueSpec(float, 0.3)}
+        )
+
+        call_kwargs = mock_rest_client.agent_configs.update_agent_config.call_args[1]
+        blueprint = call_kwargs["blueprint"]
+        assert blueprint.type == "mask"
+        assert blueprint.values is not None
+
+    def test_create_mask__returns_mask_id(self, agent_config, mock_rest_client):
+        result = agent_config.create_mask(
+            fields_with_values={"temperature": types.FieldValueSpec(float, 0.3)}
+        )
+
+        assert isinstance(result, str)
+        mock_rest_client.agent_configs.get_blueprint_by_id.assert_not_called()
+
+    def test_create_mask__sends_under_blueprint_key(
+        self, agent_config, mock_rest_client
+    ):
+        agent_config.create_mask(
+            fields_with_values={"temperature": types.FieldValueSpec(float, 0.3)}
+        )
+
+        call_kwargs = mock_rest_client.agent_configs.update_agent_config.call_args[1]
+        assert "blueprint" in call_kwargs
+        assert call_kwargs["blueprint"].id is not None
+
+    def test_create_mask__with_parameters__returns_mask_id(
+        self, agent_config, mock_rest_client
+    ):
+        result = agent_config.create_mask(parameters={"temp": 0.3})
+
+        assert isinstance(result, str)
+
+    def test_create_mask__with_description__passes_description(
+        self, agent_config, mock_rest_client
+    ):
+        agent_config.create_mask(
+            fields_with_values={"temperature": types.FieldValueSpec(float, 0.3)},
+            description="variant-A",
+        )
+
+        call_kwargs = mock_rest_client.agent_configs.update_agent_config.call_args[1]
+        assert call_kwargs["blueprint"].description == "variant-A"
+
+    def test_create_mask__with_project_name__passes_project_to_backend(
+        self, agent_config, mock_rest_client
+    ):
+        agent_config.create_mask(
+            fields_with_values={"temp": types.FieldValueSpec(float, 0.5)}
+        )
+
+        call_kwargs = mock_rest_client.agent_configs.update_agent_config.call_args[1]
+        assert call_kwargs["project_name"] == "my-project"
+
+
+class TestResolveFieldsWithValues:
+    def test_none_value__included_with_str_type(self):
+        result = ConfigManager._resolve_fields_with_values(
+            parameters={"temp": 0.5, "name": None},
+            fields_with_values=None,
+        )
+        assert result["temp"] == types.FieldValueSpec(float, 0.5)
+        assert result["name"] == types.FieldValueSpec(str, None)
+
+    def test_all_none_parameters__included_with_str_type(self):
+        result = ConfigManager._resolve_fields_with_values(
+            parameters={"a": None, "b": None},
+            fields_with_values=None,
+        )
+        assert result == {
+            "a": types.FieldValueSpec(str, None),
+            "b": types.FieldValueSpec(str, None),
+        }
+
+    def test_fields_with_values_takes_precedence_over_parameters(self):
+        explicit = {"x": types.FieldValueSpec(int, 1)}
+        result = ConfigManager._resolve_fields_with_values(
+            parameters={"x": 99},
+            fields_with_values=explicit,
+        )
+        assert result is explicit
+
+    def test_create_blueprint__none_parameter__included_in_payload_with_string_type(
+        self, agent_config, mock_rest_client
+    ):
+        mock_rest_client.agent_configs.get_blueprint_by_id.return_value = (
+            _make_raw_blueprint()
+        )
+
+        agent_config.create_blueprint(parameters={"temp": 0.6, "name": None})
+
+        call_kwargs = mock_rest_client.agent_configs.create_agent_config.call_args[1]
+        values_by_key = {v.key: v for v in call_kwargs["blueprint"].values}
+        assert "temp" in values_by_key
+        assert "name" in values_by_key
+        assert values_by_key["name"].type == "string"
+        assert values_by_key["name"].value is None

sdks/python/tests/unit/api_objects/agent_config/test_context.py

@@ -0,0 +1,35 @@
+from opik.api_objects.agent_config.context import (
+    agent_config_context,
+    get_active_config_mask,
+)
+
+
+class TestConfigContext:
+    def test_no_context__returns_none(self):
+        assert get_active_config_mask() is None
+
+    def test_inside_context__returns_active_mask(self):
+        with agent_config_context("mask-1"):
+            assert get_active_config_mask() == "mask-1"
+
+    def test_after_context_exit__returns_none(self):
+        with agent_config_context("mask-1"):
+            pass
+        assert get_active_config_mask() is None
+
+    def test_nested_contexts__inner_wins_then_restores_outer(self):
+        with agent_config_context("outer"):
+            assert get_active_config_mask() == "outer"
+            with agent_config_context("inner"):
+                assert get_active_config_mask() == "inner"
+            assert get_active_config_mask() == "outer"
+        assert get_active_config_mask() is None
+
+    def test_exception_inside_context__mask_still_resets(self):
+        try:
+            with agent_config_context("mask-err"):
+                assert get_active_config_mask() == "mask-err"
+                raise ValueError("test error")
+        except ValueError:
+            pass
+        assert get_active_config_mask() is None

sdks/python/tests/unit/api_objects/agent_config/test_type_helpers.py

@@ -0,0 +1,445 @@
+import json
+from typing import List, Dict, Callable, Optional
+from unittest import mock
+
+import pytest
+
+from opik.api_objects import type_helpers
+from opik.api_objects.prompt.base_prompt import BasePrompt
+from opik.api_objects.prompt.text.prompt import Prompt
+from opik.api_objects.prompt.chat.chat_prompt import ChatPrompt
+from opik.rest_api.types.prompt_version_detail import PromptVersionDetail
+
+
+class TestIsPromptType:
+    @pytest.mark.parametrize(
+        "py_type",
+        [BasePrompt, Prompt, ChatPrompt],
+        ids=["BasePrompt", "Prompt", "ChatPrompt"],
+    )
+    def test_prompt_classes__returns_true(self, py_type):
+        assert type_helpers.is_prompt_type(py_type) is True
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [str, int, float, bool, object, list],
+        ids=["str", "int", "float", "bool", "object", "list"],
+    )
+    def test_non_prompt_types__returns_false(self, py_type):
+        assert type_helpers.is_prompt_type(py_type) is False
+
+    def test_custom_subclass_of_base_prompt__returns_true(self):
+        class MyPrompt(BasePrompt):
+            @property
+            def name(self):
+                return ""
+
+            @property
+            def commit(self):
+                return None
+
+            @property
+            def version_id(self):
+                return ""
+
+            @property
+            def metadata(self):
+                return None
+
+            @property
+            def type(self):
+                return None
+
+            @property
+            def id(self):
+                return None
+
+            @property
+            def description(self):
+                return None
+
+            @property
+            def change_description(self):
+                return None
+
+            @property
+            def tags(self):
+                return None
+
+            def format(self, *args, **kwargs):
+                return ""
+
+            def __internal_api__to_info_dict__(self):
+                return {}
+
+        assert type_helpers.is_prompt_type(MyPrompt) is True
+
+
+class TestBackendTypeToPythonType:
+    @pytest.mark.parametrize(
+        "backend_type, expected",
+        [
+            ("string", str),
+            ("integer", int),
+            ("float", float),
+            ("boolean", bool),
+        ],
+        ids=["string", "integer", "float", "boolean"],
+    )
+    def test_primitive_types__returns_python_type(self, backend_type, expected):
+        assert type_helpers.backend_type_to_python_type(backend_type) is expected
+
+    @pytest.mark.parametrize(
+        "backend_type",
+        ["prompt", "prompt_commit", "unknown", ""],
+        ids=["prompt", "prompt_commit", "unknown", "empty"],
+    )
+    def test_non_primitive_types__returns_none(self, backend_type):
+        assert type_helpers.backend_type_to_python_type(backend_type) is None
+
+
+class TestIsSupportedType:
+    @pytest.mark.parametrize(
+        "py_type, expected",
+        [
+            (str, True),
+            (int, True),
+            (float, True),
+            (bool, True),
+            (List[str], True),
+            (List[int], True),
+            (List[float], True),
+            (List[bool], True),
+            (Dict[str, float], True),
+            (Dict[str, bool], True),
+            (Dict[str, int], True),
+            (Dict[str, str], True),
+            (Prompt, True),
+            (ChatPrompt, True),
+            (BasePrompt, True),
+            (PromptVersionDetail, True),
+        ],
+        ids=[
+            "str",
+            "int",
+            "float",
+            "bool",
+            "List[str]",
+            "List[int]",
+            "List[float]",
+            "List[bool]",
+            "Dict[str,float]",
+            "Dict[str,bool]",
+            "Dict[str,int]",
+            "Dict[str,str]",
+            "Prompt",
+            "ChatPrompt",
+            "BasePrompt",
+            "PromptVersionDetail",
+        ],
+    )
+    def test_supported_types__returns_true(self, py_type, expected):
+        assert type_helpers.is_supported_type(py_type) is expected
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [
+            list,
+            dict,
+            Callable,
+            object,
+            Dict[int, str],
+        ],
+        ids=["bare_list", "bare_dict", "Callable", "object", "Dict[int,str]"],
+    )
+    def test_unsupported_types__returns_false(self, py_type):
+        assert type_helpers.is_supported_type(py_type) is False
+
+    def test_custom_class__returns_false(self):
+        class Foo:
+            pass
+
+        assert type_helpers.is_supported_type(Foo) is False
+
+    def test_list_of_custom_class__returns_false(self):
+        class Bar:
+            pass
+
+        assert type_helpers.is_supported_type(List[Bar]) is False
+
+
+class TestPythonTypeToBackendType:
+    @pytest.mark.parametrize(
+        "py_type, expected_backend_type",
+        [
+            (str, "string"),
+            (int, "integer"),
+            (float, "float"),
+            (bool, "boolean"),
+            (List[str], "string"),
+            (Dict[str, int], "string"),
+        ],
+        ids=["str", "int", "float", "bool", "List", "Dict"],
+    )
+    def test_known_types__returns_correct_backend_type(
+        self, py_type, expected_backend_type
+    ):
+        assert (
+            type_helpers.python_type_to_backend_type(py_type) == expected_backend_type
+        )
+
+    def test_unsupported_type__raises_type_error(self):
+        with pytest.raises(TypeError):
+            type_helpers.python_type_to_backend_type(object)
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [Prompt, ChatPrompt, BasePrompt],
+        ids=["Prompt", "ChatPrompt", "BasePrompt"],
+    )
+    def test_prompt_types__return_prompt(self, py_type):
+        assert type_helpers.python_type_to_backend_type(py_type) == "prompt"
+
+    def test_prompt_version_type__returns_prompt_commit(self):
+        assert (
+            type_helpers.python_type_to_backend_type(PromptVersionDetail)
+            == "prompt_commit"
+        )
+
+
+class TestPythonValueToBackendValue:
+    @pytest.mark.parametrize(
+        "value, py_type, expected",
+        [
+            ("hello", str, "hello"),
+            (42, int, "42"),
+            (0, int, "0"),
+            (0.6, float, "0.6"),
+            (0.0, float, "0.0"),
+            (True, bool, "true"),
+            (False, bool, "false"),
+        ],
+        ids=[
+            "str",
+            "int",
+            "int_zero",
+            "float",
+            "float_zero",
+            "bool_true",
+            "bool_false",
+        ],
+    )
+    def test_primitives__serialized_correctly(self, value, py_type, expected):
+        assert type_helpers.python_value_to_backend_value(value, py_type) == expected
+
+    @pytest.mark.parametrize(
+        "value, py_type, expected_parsed",
+        [
+            ([1, 2, 3], List[int], [1, 2, 3]),
+            ([], List[str], []),
+            ({"a": 1, "b": 2}, Dict[str, int], {"a": 1, "b": 2}),
+            ({}, Dict[str, str], {}),
+        ],
+        ids=["list", "empty_list", "dict", "empty_dict"],
+    )
+    def test_collections__serialized_as_json(self, value, py_type, expected_parsed):
+        result = type_helpers.python_value_to_backend_value(value, py_type)
+        assert json.loads(result) == expected_parsed
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [Prompt, ChatPrompt, BasePrompt],
+        ids=["Prompt", "ChatPrompt", "BasePrompt"],
+    )
+    def test_prompt_value__returns_commit(self, py_type):
+        prompt = mock.Mock()
+        prompt.commit = "abc12345"
+        assert type_helpers.python_value_to_backend_value(prompt, py_type) == "abc12345"
+
+    def test_prompt_version_value__returns_commit(self):
+        version = mock.Mock(spec=PromptVersionDetail)
+        version.commit = "pv123456"
+        assert (
+            type_helpers.python_value_to_backend_value(version, PromptVersionDetail)
+            == "pv123456"
+        )
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [str, int, float, bool, List[str], Dict[str, int]],
+        ids=["str", "int", "float", "bool", "List[str]", "Dict[str,int]"],
+    )
+    def test_none_value__returns_none(self, py_type):
+        assert type_helpers.python_value_to_backend_value(None, py_type) is None
+
+
+class TestBackendValueToPythonValue:
+    @pytest.mark.parametrize(
+        "value, py_type, expected",
+        [
+            ("hello", str, "hello"),
+            ("42", int, 42),
+            ("42.0", int, 42),
+            (42, int, 42),
+            ("0.6", float, 0.6),
+            (0.6, float, 0.6),
+            ("true", bool, True),
+            ("1", bool, True),
+            ("yes", bool, True),
+            ("false", bool, False),
+            ("0", bool, False),
+            (True, bool, True),
+        ],
+        ids=[
+            "str",
+            "int_from_str",
+            "int_from_float_str",
+            "int_native",
+            "float_from_str",
+            "float_native",
+            "bool_true_str",
+            "bool_one_str",
+            "bool_yes_str",
+            "bool_false_str",
+            "bool_zero_str",
+            "bool_native",
+        ],
+    )
+    def test_primitives__deserialized_correctly(self, value, py_type, expected):
+        assert type_helpers.backend_value_to_python_value(value, py_type) == expected
+
+    @pytest.mark.parametrize(
+        "value, py_type, expected",
+        [
+            ("[1, 2, 3]", List[int], [1, 2, 3]),
+            ([1, 2, 3], List[int], [1, 2, 3]),
+            ('{"a": 1}', Dict[str, int], {"a": 1}),
+            ({"a": 1}, Dict[str, int], {"a": 1}),
+        ],
+        ids=["list_from_json", "list_native", "dict_from_json", "dict_native"],
+    )
+    def test_collections__deserialized_correctly(self, value, py_type, expected):
+        assert type_helpers.backend_value_to_python_value(value, py_type) == expected
+
+    def test_none__returns_none(self):
+        assert type_helpers.backend_value_to_python_value(None, str) is None
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [Prompt, ChatPrompt, BasePrompt],
+        ids=["Prompt", "ChatPrompt", "BasePrompt"],
+    )
+    def test_prompt_type__returns_raw_version_id_string(self, py_type):
+        result = type_helpers.backend_value_to_python_value("ver-xyz", py_type)
+        assert result == "ver-xyz"
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [Prompt, ChatPrompt, BasePrompt],
+        ids=["Prompt", "ChatPrompt", "BasePrompt"],
+    )
+    def test_prompt_type__none_value__returns_none(self, py_type):
+        assert type_helpers.backend_value_to_python_value(None, py_type) is None
+
+    def test_prompt_version_type__returns_raw_version_id_string(self):
+        result = type_helpers.backend_value_to_python_value(
+            "ver-pv-xyz", PromptVersionDetail
+        )
+        assert result == "ver-pv-xyz"
+
+    def test_prompt_version_type__none_value__returns_none(self):
+        assert (
+            type_helpers.backend_value_to_python_value(None, PromptVersionDetail)
+            is None
+        )
+
+
+class TestRoundTrip:
+    @pytest.mark.parametrize(
+        "value, py_type",
+        [
+            ("hello", str),
+            (42, int),
+            (0, int),
+            (0.6, float),
+            (0.0, float),
+            (True, bool),
+            (False, bool),
+        ],
+        ids=[
+            "str",
+            "int",
+            "int_zero",
+            "float",
+            "float_zero",
+            "bool_true",
+            "bool_false",
+        ],
+    )
+    def test_serialize_then_deserialize__recovers_original(self, value, py_type):
+        backend_value = type_helpers.python_value_to_backend_value(value, py_type)
+        restored = type_helpers.backend_value_to_python_value(backend_value, py_type)
+        assert restored == value
+        assert isinstance(restored, py_type)
+
+
+class TestUnwrapOptional:
+    @pytest.mark.parametrize(
+        "py_type, expected",
+        [
+            (Optional[str], str),
+            (Optional[int], int),
+            (Optional[float], float),
+            (Optional[bool], bool),
+        ],
+        ids=["Optional[str]", "Optional[int]", "Optional[float]", "Optional[bool]"],
+    )
+    def test_optional_primitives__returns_inner_type(self, py_type, expected):
+        assert type_helpers.unwrap_optional(py_type) is expected
+
+    @pytest.mark.parametrize(
+        "py_type",
+        [str, int, float, bool, List[str], Dict[str, int]],
+        ids=["str", "int", "float", "bool", "List[str]", "Dict[str,int]"],
+    )
+    def test_non_optional__returns_none(self, py_type):
+        assert type_helpers.unwrap_optional(py_type) is None
+
+    def test_union_with_multiple_types__returns_none(self):
+        # Union[str, int] is not Optional so must not be unwrapped
+        import typing
+
+        assert type_helpers.unwrap_optional(typing.Union[str, int]) is None
+
+
+class TestIsSupportedTypeOptional:
+    @pytest.mark.parametrize(
+        "py_type",
+        [Optional[str], Optional[int], Optional[float], Optional[bool]],
+        ids=["Optional[str]", "Optional[int]", "Optional[float]", "Optional[bool]"],
+    )
+    def test_optional_primitives__returns_true(self, py_type):
+        assert type_helpers.is_supported_type(py_type) is True
+
+    def test_optional_unsupported__returns_false(self):
+        assert type_helpers.is_supported_type(Optional[object]) is False
+
+    def test_union_with_multiple_types__returns_false(self):
+        import typing
+
+        assert type_helpers.is_supported_type(typing.Union[str, int]) is False
+
+
+class TestPythonTypeToBackendTypeOptional:
+    @pytest.mark.parametrize(
+        "py_type, expected",
+        [
+            (Optional[str], "string"),
+            (Optional[int], "integer"),
+            (Optional[float], "float"),
+            (Optional[bool], "boolean"),
+        ],
+        ids=["Optional[str]", "Optional[int]", "Optional[float]", "Optional[bool]"],
+    )
+    def test_optional_primitives__returns_inner_backend_type(self, py_type, expected):
+        assert type_helpers.python_type_to_backend_type(py_type) == expected

sdks/python/tests/unit/api_objects/annotation_queue/test_annotation_queue.py

@@ -0,0 +1,391 @@
+import json
+from unittest.mock import Mock
+import pytest
+
+from opik.api_objects.annotation_queue.annotation_queue import (
+    TracesAnnotationQueue,
+    ThreadsAnnotationQueue,
+)
+from opik.rest_api.types import (
+    trace_public,
+    trace_thread,
+)
+from opik.exceptions import OpikException
+
+
+class TestTracesAnnotationQueueProperties:
+    def test_id__returns_id(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        assert queue.id == "queue-123"
+
+    def test_name__returns_name(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        assert queue.name == "test_queue"
+
+    def test_scope__returns_trace(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        assert queue.scope == "trace"
+
+    def test_items_count__cached_value__returns_cached_count(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            items_count=10,
+        )
+
+        assert queue.items_count == 10
+        mock_rest_client.annotation_queues.get_annotation_queue_by_id.assert_not_called()
+
+    def test_items_count__no_cached_value__fetches_from_backend(self):
+        mock_rest_client = Mock()
+        mock_rest_client.annotation_queues.get_annotation_queue_by_id.return_value = (
+            Mock(items_count=25)
+        )
+
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            items_count=None,
+        )
+
+        assert queue.items_count == 25
+        mock_rest_client.annotation_queues.get_annotation_queue_by_id.assert_called_once_with(
+            "queue-123"
+        )
+
+
+class TestThreadsAnnotationQueueProperties:
+    def test_scope__returns_thread(self):
+        mock_rest_client = Mock()
+        queue = ThreadsAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        assert queue.scope == "thread"
+
+
+class TestTracesAnnotationQueueUpdate:
+    def test_update__happyflow(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            description="old description",
+        )
+
+        queue.update(
+            name="new_name",
+            description="new description",
+            instructions="new instructions",
+        )
+
+        mock_rest_client.annotation_queues.update_annotation_queue.assert_called_once_with(
+            id="queue-123",
+            name="new_name",
+            description="new description",
+            instructions="new instructions",
+            comments_enabled=None,
+            feedback_definition_names=None,
+        )
+
+        assert queue.name == "new_name"
+        assert queue.description == "new description"
+        assert queue.instructions == "new instructions"
+
+    def test_update__partial_update__only_specified_fields_updated(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            description="old description",
+            instructions="old instructions",
+        )
+
+        queue.update(description="new description")
+
+        mock_rest_client.annotation_queues.update_annotation_queue.assert_called_once_with(
+            id="queue-123",
+            name=None,
+            description="new description",
+            instructions=None,
+            comments_enabled=None,
+            feedback_definition_names=None,
+        )
+
+        assert queue.name == "test_queue"
+        assert queue.description == "new description"
+        assert queue.instructions == "old instructions"
+
+
+class TestTracesAnnotationQueueDelete:
+    def test_delete__happyflow(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        queue.delete()
+
+        mock_rest_client.annotation_queues.delete_annotation_queue_batch.assert_called_once_with(
+            ids=["queue-123"]
+        )
+
+
+class TestTracesAnnotationQueueAddTraces:
+    def test_add_traces__single_trace_in_list__happyflow(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            items_count=5,
+        )
+
+        trace = Mock(spec=trace_public.TracePublic)
+        trace.id = "trace-1"
+
+        queue.add_traces([trace])
+
+        mock_rest_client.annotation_queues.add_items_to_annotation_queue.assert_called_once_with(
+            id="queue-123", ids=["trace-1"]
+        )
+        assert queue._items_count is None
+
+    def test_add_traces__multiple_traces__happyflow(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        traces = [
+            Mock(spec=trace_public.TracePublic, id="trace-1"),
+            Mock(spec=trace_public.TracePublic, id="trace-2"),
+            Mock(spec=trace_public.TracePublic, id="trace-3"),
+        ]
+
+        queue.add_traces(traces)
+
+        mock_rest_client.annotation_queues.add_items_to_annotation_queue.assert_called_once_with(
+            id="queue-123", ids=["trace-1", "trace-2", "trace-3"]
+        )
+
+    def test_add_traces__trace_without_id__raises_exception(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        trace = Mock(spec=trace_public.TracePublic)
+        trace.id = None
+
+        with pytest.raises(OpikException) as exc_info:
+            queue.add_traces([trace])
+
+        assert "no id" in str(exc_info.value)
+
+    def test_add_traces__empty_list__no_api_call(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        queue.add_traces([])
+
+        mock_rest_client.annotation_queues.add_items_to_annotation_queue.assert_not_called()
+
+
+class TestThreadsAnnotationQueueAddThreads:
+    def test_add_threads__single_thread_in_list__happyflow(self):
+        mock_rest_client = Mock()
+        queue = ThreadsAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            items_count=5,
+        )
+
+        thread = Mock(spec=trace_thread.TraceThread)
+        thread.thread_model_id = "thread-model-1"
+
+        queue.add_threads([thread])
+
+        mock_rest_client.annotation_queues.add_items_to_annotation_queue.assert_called_once_with(
+            id="queue-123", ids=["thread-model-1"]
+        )
+        assert queue._items_count is None
+
+    def test_add_threads__multiple_threads__happyflow(self):
+        mock_rest_client = Mock()
+        queue = ThreadsAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        thread1 = Mock(spec=trace_thread.TraceThread)
+        thread1.thread_model_id = "thread-model-1"
+        thread2 = Mock(spec=trace_thread.TraceThread)
+        thread2.thread_model_id = "thread-model-2"
+        threads = [thread1, thread2]
+
+        queue.add_threads(threads)
+
+        mock_rest_client.annotation_queues.add_items_to_annotation_queue.assert_called_once_with(
+            id="queue-123", ids=["thread-model-1", "thread-model-2"]
+        )
+
+    def test_add_threads__thread_without_thread_model_id__raises_exception(self):
+        mock_rest_client = Mock()
+        queue = ThreadsAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        thread = Mock(spec=trace_thread.TraceThread)
+        thread.thread_model_id = None
+
+        with pytest.raises(OpikException) as exc_info:
+            queue.add_threads([thread])
+
+        assert "thread_model_id" in str(exc_info.value)
+
+
+class TestTracesAnnotationQueueRemoveTraces:
+    def test_remove_traces__single_trace_in_list__happyflow(self):
+        mock_rest_client = Mock()
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            items_count=5,
+        )
+
+        trace = Mock(spec=trace_public.TracePublic)
+        trace.id = "trace-1"
+
+        queue.remove_traces([trace])
+
+        mock_rest_client.annotation_queues.remove_items_from_annotation_queue.assert_called_once_with(
+            id="queue-123", ids=["trace-1"]
+        )
+        assert queue._items_count is None
+
+
+class TestThreadsAnnotationQueueRemoveThreads:
+    def test_remove_threads__single_thread_in_list__happyflow(self):
+        mock_rest_client = Mock()
+        queue = ThreadsAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+            items_count=5,
+        )
+
+        thread = Mock(spec=trace_thread.TraceThread)
+        thread.thread_model_id = "thread-model-1"
+
+        queue.remove_threads([thread])
+
+        mock_rest_client.annotation_queues.remove_items_from_annotation_queue.assert_called_once_with(
+            id="queue-123", ids=["thread-model-1"]
+        )
+        assert queue._items_count is None
+
+
+class TestTracesAnnotationQueueGetItems:
+    def test_get_items__happyflow(self):
+        mock_rest_client = Mock()
+        mock_rest_client.traces.search_traces.return_value = [
+            json.dumps({"id": "trace-1", "start_time": "2024-01-01T00:00:00Z"}).encode(
+                "utf-8"
+            )
+            + b"\n",
+            json.dumps({"id": "trace-2", "start_time": "2024-01-01T00:00:00Z"}).encode(
+                "utf-8"
+            ),
+        ]
+
+        queue = TracesAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        items = queue.get_items()
+
+        assert [item.id for item in items] == ["trace-1", "trace-2"]
+        assert all(isinstance(item, trace_public.TracePublic) for item in items)
+
+
+class TestThreadsAnnotationQueueGetItems:
+    def test_get_items__happyflow(self):
+        mock_rest_client = Mock()
+        mock_rest_client.traces.search_trace_threads.return_value = [
+            json.dumps({"id": "thread-1"}).encode("utf-8") + b"\n",
+            json.dumps({"id": "thread-2"}).encode("utf-8"),
+        ]
+
+        queue = ThreadsAnnotationQueue(
+            id="queue-123",
+            name="test_queue",
+            project_id="project-123",
+            rest_client=mock_rest_client,
+        )
+
+        items = queue.get_items()
+
+        assert [item.id for item in items] == ["thread-1", "thread-2"]
+        assert all(isinstance(item, trace_thread.TraceThread) for item in items)

sdks/python/tests/unit/api_objects/api_key/test_opik_api_key.py

@@ -0,0 +1,102 @@
+from opik.api_key.opik_api_key import DELIMITER_CHAR, parse_api_key
+from opik.logging_messages import (
+    PARSE_API_KEY_EMPTY_EXPECTED_ATTRIBUTES,
+    PARSE_API_KEY_TOO_MANY_PARTS,
+)
+
+import pytest
+
+
+@pytest.mark.parametrize("raw_key", ["", None])
+def test_parse_api_key__empty_key(raw_key, capture_log):
+    opik_api_key = parse_api_key(raw_key)
+
+    assert opik_api_key is None
+
+
+def test_parse_api_key__one_part():
+    raw_key = "some API key"
+    opik_api_key = parse_api_key(raw_key)
+
+    assert opik_api_key is not None
+    assert opik_api_key.api_key == raw_key
+    assert opik_api_key.short_api_key == raw_key
+
+
+def test_parse_api_key__no_expected_attributes(capture_log):
+    raw_key = "some API key"
+    opik_api_key = parse_api_key(raw_key + DELIMITER_CHAR)
+
+    assert opik_api_key is not None
+    assert opik_api_key.api_key == raw_key
+    assert opik_api_key.short_api_key == raw_key
+
+    assert (
+        PARSE_API_KEY_EMPTY_EXPECTED_ATTRIBUTES % (raw_key + DELIMITER_CHAR)
+        in capture_log.messages
+    )
+
+
+def test_parse_api_key__too_many_parts(capture_log):
+    raw_key = "some API key" + DELIMITER_CHAR + "one" + DELIMITER_CHAR + "two"
+    opik_api_key = parse_api_key(raw_key)
+
+    assert opik_api_key is None
+    assert PARSE_API_KEY_TOO_MANY_PARTS % (3, raw_key) in capture_log.messages
+
+
+def test_parse_api_key__happy_path__with_padding():
+    # attributes: {"baseUrl": "https://www.comet.com"}
+    raw_key = (
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0="
+    )
+    opik_api_key = parse_api_key(raw_key)
+
+    assert opik_api_key is not None
+    assert opik_api_key.api_key == raw_key
+    assert opik_api_key.short_api_key == "Et1RBc4nd1ef3LfyJvhyB34Po"
+
+    assert opik_api_key.base_url == "https://www.comet.com"
+    assert opik_api_key["baseUrl"] == "https://www.comet.com"
+
+
+def test_parse_api_key__happy_path__no_padding():
+    # attributes: {"baseUrl": "https://www.comet.com"}
+    raw_key = (
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0"
+    )
+    opik_api_key = parse_api_key(raw_key)
+
+    assert opik_api_key is not None
+    assert opik_api_key.api_key == raw_key
+    assert opik_api_key.short_api_key == "Et1RBc4nd1ef3LfyJvhyB34Po"
+
+    assert opik_api_key.base_url == "https://www.comet.com"
+    assert opik_api_key["baseUrl"] == "https://www.comet.com"
+
+
+@pytest.mark.parametrize(
+    "raw_key",
+    [
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0===",
+        "Et1RBc4nd1ef3LfyJvhyB34Po"
+        + DELIMITER_CHAR
+        + "eyJiYXNlVXJsIjoiaHR0cHM6Ly93d3cuY29tZXQuY29tIn0==",
+    ],
+)
+def test_parse_api_key__happy_path__wrong_padding(raw_key):
+    # attributes: {"baseUrl": "https://www.comet.com"}
+    opik_api_key = parse_api_key(raw_key)
+
+    assert opik_api_key is not None
+    assert opik_api_key.api_key == raw_key
+    assert opik_api_key.short_api_key == "Et1RBc4nd1ef3LfyJvhyB34Po"
+
+    assert opik_api_key.base_url == "https://www.comet.com"
+    assert opik_api_key["baseUrl"] == "https://www.comet.com"

sdks/python/tests/unit/api_objects/attachment/constants.py

@@ -0,0 +1,51 @@
+"""
+Common test data constants for attachment tests.
+
+This module contains base64-encoded and binary test data for various file formats
+used across attachment decoder tests.
+"""
+
+import base64
+
+
+# PNG Test Data
+# 1x1 transparent PNG image
+PNG_BASE64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+
+# JPEG Test Data
+# Minimal valid JPEG image (1x1 red pixel)
+JPEG_BASE64 = "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/2wBDAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQEBAQH/wAARCAABAAEDASIAAhEBAxEB/8QAFQABAQAAAAAAAAAAAAAAAAAAAAv/xAAUEAEAAAAAAAAAAAAAAAAAAAAA/8QAFQEBAQAAAAAAAAAAAAAAAAAAAAX/xAAUEQEAAAAAAAAAAAAAAAAAAAAA/9oADAMBAAIRAxEAPwCRAP/Z"
+
+# PDF Test Data
+# Minimal valid PDF document with "Hello World"
+PDF_BYTES: bytes = b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n1 0 obj\n<<\n/Type /Catalog\n/Pages 2 0 R\n>>\nendobj\n2 0 obj\n<<\n/Type /Pages\n/Kids [3 0 R]\n/Count 1\n>>\nendobj\n3 0 obj\n<<\n/Type /Page\n/Parent 2 0 R\n/Resources <<\n/Font <<\n/F1 4 0 R\n>>\n>>\n/MediaBox [0 0 612 792]\n/Contents 5 0 R\n>>\nendobj\n4 0 obj\n<<\n/Type /Font\n/Subtype /Type1\n/BaseFont /Helvetica\n>>\nendobj\n5 0 obj\n<<\n/Length 44\n>>\nstream\nBT\n/F1 24 Tf\n100 700 Td\n(Hello World) Tj\nET\nendstream\nendobj\nxref\n0 6\n0000000000 65535 f\n0000000015 00000 n\n0000000074 00000 n\n0000000131 00000 n\n0000000277 00000 n\n0000000356 00000 n\ntrailer\n<<\n/Size 6\n/Root 1 0 R\n>>\nstartxref\n448\n%%EOF"
+PDF_BASE64 = base64.b64encode(PDF_BYTES).decode("utf-8")
+
+# GIF Test Data
+# Minimal GIF89a image
+GIF89_BYTES: bytes = (
+    b"GIF89a\x01\x00\x01\x00\x00\xff\x00,\x00\x00\x00\x00\x01\x00\x01\x00\x00\x02\x00;"
+)
+GIF89_BASE64 = base64.b64encode(GIF89_BYTES).decode("utf-8")
+
+# WebP Test Data
+# Minimal WebP image
+WEBP_BYTES: bytes = b"RIFF\x1a\x00\x00\x00WEBPVP8 \x0e\x00\x00\x000\x01\x00\x9d\x01*\x01\x00\x01\x00\x00\x00\x00\x00"
+WEBP_BASE64 = base64.b64encode(WEBP_BYTES).decode("utf-8")
+
+# SVG Test Data
+# Simple SVG image with a red circle
+SVG_BYTES: bytes = b'<svg xmlns="http://www.w3.org/2000/svg" width="100" height="100"><circle cx="50" cy="50" r="40" fill="red"/></svg>'
+SVG_BASE64 = base64.b64encode(SVG_BYTES).decode("utf-8")
+
+# JSON Test Data
+JSON_BYTES: bytes = b'{"key": "value", "number": 123, "array": [1, 2, 3]}'
+JSON_BASE64 = base64.b64encode(JSON_BYTES).decode("utf-8")
+
+# Plain Text Test Data
+PLAIN_TEXT_BYTES: bytes = b"This is just plain text without special markers"
+PLAIN_TEXT_BASE64 = base64.b64encode(PLAIN_TEXT_BYTES).decode("utf-8")
+
+# Random Binary Data (octet-stream)
+RANDOM_BINARY_BYTES: bytes = b"\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a"
+RANDOM_BINARY_BASE64 = base64.b64encode(RANDOM_BINARY_BYTES).decode("utf-8")

sdks/python/tests/unit/api_objects/attachment/test_base64_normalizer.py

@@ -0,0 +1,240 @@
+import base64
+import uuid
+
+import pytest
+
+from opik.api_objects.attachment import base64_normalizer
+from . import constants
+
+
+def _to_urlsafe(standard_b64: str) -> str:
+    """Re-encode a standard-base64 string into URL-safe form for fixtures."""
+    return base64.urlsafe_b64encode(base64.b64decode(standard_b64)).decode("utf-8")
+
+
+# Bytes whose base64 encoding contains '+' / '/' (so URL-safe encoding contains
+# '-' / '_'). Appended to fixtures to guarantee we actually exercise the
+# URL-safe path rather than falling through the "no '-'/'_'" early exit.
+_URLSAFE_FORCING_TAIL = b"\xfb\xff\xfe" * 8
+
+
+def _urlsafe_image_fixture(signature: bytes) -> str:
+    """Build a URL-safe base64 string that (a) decodes to bytes starting with
+    ``signature`` and (b) is guaranteed to contain '-' or '_' so the
+    detector's URL-safe-alphabet pre-check actually runs.
+    """
+    encoded = base64.urlsafe_b64encode(signature + _URLSAFE_FORCING_TAIL).decode(
+        "utf-8"
+    )
+    assert "-" in encoded or "_" in encoded, "fixture failed to force URL-safe chars"
+    return encoded
+
+
+# ---------------------------------------------------------------------------
+# urlsafe_to_standard_base64
+# ---------------------------------------------------------------------------
+
+
+class TestUrlsafeToStandardBase64:
+    def test_urlsafe_to_standard_base64__dash_and_underscore_chars__replaced_with_plus_and_slash(
+        self,
+    ):
+        assert (
+            base64_normalizer.urlsafe_to_standard_base64("ab-cd_ef==") == "ab+cd/ef=="
+        )
+
+    def test_urlsafe_to_standard_base64__already_standard_alphabet__returns_unchanged(
+        self,
+    ):
+        value = "iVBORw0KGgo+AB/CD=="
+        assert base64_normalizer.urlsafe_to_standard_base64(value) == value
+
+    def test_urlsafe_to_standard_base64__urlsafe_chars_with_padding__padding_preserved(
+        self,
+    ):
+        assert base64_normalizer.urlsafe_to_standard_base64("ab-_==") == "ab+/=="
+
+    def test_urlsafe_to_standard_base64__empty_string__returns_empty(self):
+        assert base64_normalizer.urlsafe_to_standard_base64("") == ""
+
+
+# ---------------------------------------------------------------------------
+# is_urlsafe_base64_image
+# ---------------------------------------------------------------------------
+
+
+class TestIsUrlsafeBase64Image:
+    @pytest.mark.parametrize(
+        "signature",
+        [
+            b"\x89PNG\r\n\x1a\n",
+            b"\xff\xd8\xff",
+            b"GIF87a",
+            b"GIF89a",
+            b"RIFF\x00\x00\x00\x00WEBP",
+        ],
+        ids=["png", "jpeg", "gif87a", "gif89a", "webp"],
+    )
+    def test_is_urlsafe_base64_image__urlsafe_encoded_image_signature__returns_true(
+        self, signature
+    ):
+        """Every supported image signature, when URL-safe-base64-encoded, is detected."""
+        urlsafe = _urlsafe_image_fixture(signature)
+        assert base64_normalizer.is_urlsafe_base64_image(urlsafe) is True
+
+    def test_is_urlsafe_base64_image__urlsafe_encoded_real_png__returns_true(self):
+        """End-to-end sanity check using a real PNG fixture from constants.py."""
+        urlsafe = _to_urlsafe(constants.PNG_BASE64)
+        # The real PNG fixture happens to contain '+' / '/' in its standard
+        # encoding, so its URL-safe form contains '-' / '_' — exercise that.
+        assert "-" in urlsafe or "_" in urlsafe
+        assert base64_normalizer.is_urlsafe_base64_image(urlsafe) is True
+
+    def test_is_urlsafe_base64_image__standard_alphabet_image__returns_false(self):
+        """Standard-alphabet base64 should short-circuit to False — nothing to rewrite."""
+        assert base64_normalizer.is_urlsafe_base64_image(constants.PNG_BASE64) is False
+
+    def test_is_urlsafe_base64_image__string_below_min_length__returns_false(self):
+        assert base64_normalizer.is_urlsafe_base64_image("ab-_cd==") is False
+
+    def test_is_urlsafe_base64_image__uuid_string__returns_false(self):
+        """UUIDs share the URL-safe alphabet and contain '-', but aren't images."""
+        for _ in range(5):
+            assert base64_normalizer.is_urlsafe_base64_image(str(uuid.uuid4())) is False
+
+    def test_is_urlsafe_base64_image__urlsafe_non_image_binary__returns_false(self):
+        """Non-image binary in URL-safe base64 (e.g. PDF) must be left alone."""
+        urlsafe_pdf = _to_urlsafe(constants.PDF_BASE64)
+        assert base64_normalizer.is_urlsafe_base64_image(urlsafe_pdf) is False
+
+    def test_is_urlsafe_base64_image__plain_text_with_dashes__returns_false(self):
+        text = "hello-world this is not base64-encoded at all"
+        assert base64_normalizer.is_urlsafe_base64_image(text) is False
+
+    def test_is_urlsafe_base64_image__non_base64_chars__returns_false(self):
+        # Contains '$' which is not in either alphabet
+        assert (
+            base64_normalizer.is_urlsafe_base64_image("iVBORw0KGgo$$$abcdefgh") is False
+        )
+
+    def test_is_urlsafe_base64_image__empty_string__returns_false(self):
+        assert base64_normalizer.is_urlsafe_base64_image("") is False
+
+
+# ---------------------------------------------------------------------------
+# normalize_urlsafe_base64_images_in_place
+# ---------------------------------------------------------------------------
+
+
+class TestNormalizeUrlsafeBase64ImagesInPlace:
+    def test_normalize_urlsafe_base64_images_in_place__top_level_dict_with_image__image_rewritten(
+        self,
+    ):
+        urlsafe = _to_urlsafe(constants.PNG_BASE64)
+        payload = {"data": urlsafe}
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert "-" not in payload["data"]
+        assert "_" not in payload["data"]
+        # Round-trips back to the original PNG bytes
+        original_bytes = base64.b64decode(constants.PNG_BASE64)
+        assert base64.b64decode(payload["data"]) == original_bytes
+
+    def test_normalize_urlsafe_base64_images_in_place__list_with_image__image_rewritten(
+        self,
+    ):
+        urlsafe = _to_urlsafe(constants.JPEG_BASE64)
+        payload = ["leading text", urlsafe]
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert payload[0] == "leading text"
+        assert "-" not in payload[1] and "_" not in payload[1]
+
+    def test_normalize_urlsafe_base64_images_in_place__deeply_nested_image__image_rewritten(
+        self,
+    ):
+        urlsafe = _to_urlsafe(constants.PNG_BASE64)
+        payload = {
+            "role": "user",
+            "parts": [
+                {"text": "hello"},
+                {"inline_data": {"data": urlsafe, "mime_type": "image/png"}},
+            ],
+        }
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        normalized = payload["parts"][1]["inline_data"]["data"]
+        assert "-" not in normalized and "_" not in normalized
+        assert base64.b64decode(normalized) == base64.b64decode(constants.PNG_BASE64)
+
+    def test_normalize_urlsafe_base64_images_in_place__uuid_values__left_untouched(
+        self,
+    ):
+        uid = str(uuid.uuid4())
+        payload = {"trace_id": uid, "nested": [{"span_id": uid}]}
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert payload["trace_id"] == uid
+        assert payload["nested"][0]["span_id"] == uid
+
+    def test_normalize_urlsafe_base64_images_in_place__standard_base64_image__left_untouched(
+        self,
+    ):
+        """If the value is already standard-alphabet base64, it stays byte-for-byte identical."""
+        payload = {"data": constants.PNG_BASE64}
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert payload["data"] == constants.PNG_BASE64
+
+    def test_normalize_urlsafe_base64_images_in_place__non_image_binary__left_untouched(
+        self,
+    ):
+        """URL-safe non-image binary (e.g. PDF) is not rewritten — only images are."""
+        urlsafe_pdf = _to_urlsafe(constants.PDF_BASE64)
+        payload = {"data": urlsafe_pdf}
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert payload["data"] == urlsafe_pdf
+
+    def test_normalize_urlsafe_base64_images_in_place__non_string_leaves__ignored(self):
+        payload = {"a": 1, "b": None, "c": True, "d": 1.5, "e": [None, 2, False]}
+        snapshot = {"a": 1, "b": None, "c": True, "d": 1.5, "e": [None, 2, False]}
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert payload == snapshot
+
+    def test_normalize_urlsafe_base64_images_in_place__mutated_payload__returns_none(
+        self,
+    ):
+        urlsafe = _to_urlsafe(constants.PNG_BASE64)
+        payload = {"data": urlsafe}
+
+        result = base64_normalizer.normalize_urlsafe_base64_images_in_place(payload)
+
+        assert result is None
+        assert payload["data"] != urlsafe  # was mutated
+
+    def test_normalize_urlsafe_base64_images_in_place__empty_containers__no_change(
+        self,
+    ):
+        empty_dict: dict = {}
+        empty_list: list = []
+
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(empty_dict)
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(empty_list)
+
+        assert empty_dict == {}
+        assert empty_list == []
+
+    def test_normalize_urlsafe_base64_images_in_place__scalar_root__no_raise(self):
+        """Scalar roots are valid inputs (no-op) — the walker should not raise."""
+        base64_normalizer.normalize_urlsafe_base64_images_in_place("plain string")
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(42)
+        base64_normalizer.normalize_urlsafe_base64_images_in_place(None)

sdks/python/tests/unit/api_objects/attachment/test_converters.py

@@ -0,0 +1,481 @@
+import os
+import tempfile
+from typing import Generator
+from unittest import mock
+
+import pytest
+
+from opik.api_objects import attachment
+from opik.api_objects.attachment import converters
+from opik.message_processing import messages
+
+
+@pytest.fixture
+def original_file() -> Generator[str, None, None]:
+    """Create a file with test content and clean up after test."""
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False) as f:
+        f.write("test content")
+        path = f.name
+
+    yield path
+
+    if os.path.exists(path):
+        os.unlink(path)
+
+
+@pytest.mark.parametrize(
+    "attachment_data,expected",
+    [
+        (attachment.Attachment(data="test.png", file_name=None), "image/png"),
+        (attachment.Attachment(data="test.png", file_name="test.jpg"), "image/jpeg"),
+        (
+            attachment.Attachment(
+                data="test.pdf", file_name=None, content_type="image/jpeg"
+            ),
+            "image/jpeg",
+        ),
+    ],
+)
+def test_guess_attachment_type(attachment_data: attachment.Attachment, expected: str):
+    mimetype = converters.guess_attachment_type(attachment_data)
+    assert mimetype == expected
+
+
+def test_attachment_to_message__no_temp_copy(original_file: str):
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    attachment_data = attachment.Attachment(data=original_file, create_temp_copy=False)
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    assert message == messages.CreateAttachmentMessage(
+        file_path=original_file,
+        file_name=os.path.basename(original_file),
+        mime_type="text/plain",
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        encoded_url_override="aHR0cHM6Ly9leGFtcGxlLmNvbQ==",
+    )
+
+
+def test_attachment_to_message__file_name(original_file: str):
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    attachment_data = attachment.Attachment(
+        data=original_file, file_name="test.jpg", create_temp_copy=False
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    assert message == messages.CreateAttachmentMessage(
+        file_path=original_file,
+        file_name="test.jpg",
+        mime_type="image/jpeg",
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        encoded_url_override="aHR0cHM6Ly9leGFtcGxlLmNvbQ==",
+    )
+
+
+def test_attachment_to_message__content_type(original_file: str):
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    attachment_data = attachment.Attachment(
+        data=original_file, content_type="image/jpeg", create_temp_copy=False
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    assert message == messages.CreateAttachmentMessage(
+        file_path=original_file,
+        file_name=os.path.basename(original_file),
+        mime_type="image/jpeg",
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        encoded_url_override="aHR0cHM6Ly9leGFtcGxlLmNvbQ==",
+    )
+
+
+def test_attachment_to_message__create_temp_copy(original_file: str):
+    """Test that create_temp_copy creates a temporary file and sets delete_after_upload."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+
+    attachment_data = attachment.Attachment(
+        data=original_file,
+        file_name="test.txt",
+        content_type="text/plain",
+        create_temp_copy=True,
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    # The file_path should be different from the original (it's a temp copy)
+    assert message.file_path != original_file
+    assert message.file_path.endswith(".txt")
+    assert os.path.exists(message.file_path)
+
+    # delete_after_upload should be True when create_temp_copy is used
+    assert message.delete_after_upload is True
+
+    # Other fields should be preserved
+    assert message.file_name == "test.txt"
+    assert message.mime_type == "text/plain"
+    assert message.entity_type == "trace"
+    assert message.entity_id == entity_id
+    assert message.project_name == project_name
+
+    # Verify the temp file has the same content
+    with open(message.file_path, "r") as f:
+        assert f.read() == "test content"
+
+    # Clean up the temp copy
+    os.unlink(message.file_path)
+
+
+def test_attachment_to_message__create_temp_copy_fails_on_open__uses_original_file_and_does_not_delete(
+    original_file: str,
+):
+    """
+    Test that when create_temp_copy is True but opening the source file fails,
+    the original file is used and delete_after_upload is False.
+    """
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+
+    attachment_data = attachment.Attachment(
+        data=original_file,
+        file_name="test.txt",
+        content_type="text/plain",
+        create_temp_copy=True,
+    )
+
+    # Patch open to fail when trying to read the source file for copying
+    original_open = open
+
+    def mock_open_fail_on_rb(file, mode="r", *args, **kwargs):
+        if mode == "rb" and file == original_file:
+            raise PermissionError("Cannot read source file")
+        return original_open(file, mode, *args, **kwargs)
+
+    with mock.patch("builtins.open", side_effect=mock_open_fail_on_rb):
+        message = converters.attachment_to_message(
+            attachment_data=attachment_data,
+            entity_type="trace",
+            entity_id=entity_id,
+            project_name=project_name,
+            url_override=url_override,
+        )
+
+    # The file_path should be the original file (fallback behavior)
+    assert message.file_path == original_file
+
+    # delete_after_upload should be False to prevent deleting the original file
+    assert message.delete_after_upload is False
+
+    # Other fields should be preserved
+    assert message.file_name == "test.txt"
+    assert message.mime_type == "text/plain"
+    assert message.entity_type == "trace"
+    assert message.entity_id == entity_id
+    assert message.project_name == project_name
+
+    # Original file should still exist
+    assert os.path.exists(original_file)
+
+    # Clean up the temp copy
+    os.unlink(message.file_path)
+
+
+def test_attachment_to_message__create_temp_copy_fails_with_delete_after_upload_true__still_does_not_delete(
+    original_file: str,
+):
+    """
+    Test that when create_temp_copy fails and delete_after_upload was explicitly True,
+    delete_after_upload is still set to False to protect the original file.
+    """
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+
+    attachment_data = attachment.Attachment(
+        data=original_file,
+        file_name="test.txt",
+        content_type="text/plain",
+        create_temp_copy=True,
+    )
+
+    with mock.patch("shutil.copyfileobj", side_effect=IOError("Disk full")):
+        message = converters.attachment_to_message(
+            attachment_data=attachment_data,
+            entity_type="trace",
+            entity_id=entity_id,
+            project_name=project_name,
+            url_override=url_override,
+            delete_after_upload=True,  # Explicitly set to True
+        )
+
+    # Even though delete_after_upload was True, it should be False after failure
+    assert message.delete_after_upload is False
+    assert message.file_path == original_file
+
+
+def test_attachment_to_message__bytes_data():
+    """Test that bytes data is written to a temp file and marked for deletion."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    data = b"binary content here"
+
+    attachment_data = attachment.Attachment(
+        data=data,
+        file_name="test.bin",
+        content_type="application/octet-stream",
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    # A temp file should be created
+    assert os.path.exists(message.file_path)
+    assert message.file_path != "test.bin"
+
+    # delete_after_upload should be True for bytes data
+    assert message.delete_after_upload is True
+
+    # Verify file content matches the bytes
+    with open(message.file_path, "rb") as f:
+        assert f.read() == data
+
+    # Other fields should be preserved
+    assert message.file_name == "test.bin"
+    assert message.mime_type == "application/octet-stream"
+    assert message.entity_type == "trace"
+    assert message.entity_id == entity_id
+    assert message.project_name == project_name
+    assert message.encoded_url_override == "aHR0cHM6Ly9leGFtcGxlLmNvbQ=="
+
+    # Clean up
+    os.unlink(message.file_path)
+
+
+def test_attachment_to_message__bytes_data_without_file_name():
+    """Test bytes data without file_name uses temp file basename."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    data = b"some binary data"
+
+    attachment_data = attachment.Attachment(
+        data=data,
+        content_type="image/png",
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="span",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    # file_name should be the basename of the temp file
+    assert message.file_name == os.path.basename(message.file_path)
+    assert message.delete_after_upload is True
+    assert message.mime_type == "image/png"
+    assert message.entity_type == "span"
+
+    # Clean up
+    os.unlink(message.file_path)
+
+
+def test_attachment_to_message__bytes_data_infers_mime_type_from_file_name():
+    """Test that a mime type is inferred from file_name when data is bytes."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    data = b"\x00\x01\x02\x03"
+
+    attachment_data = attachment.Attachment(
+        data=data,
+        file_name="image.png",
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    assert message.file_name == "image.png"
+    assert message.mime_type == "image/png"
+    assert message.delete_after_upload is True
+
+    # Clean up
+    os.unlink(message.file_path)
+
+
+def test_attachment_to_message__bytes_data_default_mime_type():
+    """Test that bytes data without file_name or content_type use a binary mime type."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    data = b"\x00\x01\x02\x03"
+
+    attachment_data = attachment.Attachment(
+        data=data,
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    # Should use default binary mime type
+    assert message.mime_type == "application/octet-stream"
+    assert message.delete_after_upload is True
+
+    # Clean up
+    os.unlink(message.file_path)
+
+
+def test_attachment_to_message__bytes_data_write_fails__error_reraised():
+    """Test behavior when _write_file_like_to_temp_file fails and error reraised."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    data = b"test data"
+
+    attachment_data = attachment.Attachment(
+        data=data,
+        file_name="test.bin",
+        content_type="application/octet-stream",
+    )
+
+    with mock.patch(
+        "opik.api_objects.attachment.converters._write_file_like_to_temp_file",
+        side_effect=OSError("Disk full"),
+    ):
+        with pytest.raises(OSError):
+            converters.attachment_to_message(
+                attachment_data=attachment_data,
+                entity_type="trace",
+                entity_id=entity_id,
+                project_name=project_name,
+                url_override=url_override,
+            )
+
+
+def test_attachment_to_message__base64_string(request):
+    """Test that base64-encoded string is decoded and written to temp file."""
+    import base64
+
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+    original_content = b"This is test content for base64 encoding"
+    base64_string = base64.b64encode(original_content).decode("utf-8")
+
+    attachment_data = attachment.Attachment(
+        data=base64_string,
+        file_name="test.txt",
+        content_type="text/plain",
+    )
+
+    message = converters.attachment_to_message(
+        attachment_data=attachment_data,
+        entity_type="trace",
+        entity_id=entity_id,
+        project_name=project_name,
+        url_override=url_override,
+    )
+
+    request.addfinalizer(
+        lambda: os.path.exists(message.file_path) and os.unlink(message.file_path)
+    )
+
+    # A temp file should be created
+    assert os.path.exists(message.file_path)
+    assert message.file_path != base64_string
+
+    # delete_after_upload should be True for base64 data
+    assert message.delete_after_upload is True
+
+    # Verify file content matches the decoded bytes
+    with open(message.file_path, "rb") as f:
+        assert f.read() == original_content
+
+    # Other fields should be preserved
+    assert message.file_name == "test.txt"
+    assert message.mime_type == "text/plain"
+    assert message.entity_type == "trace"
+    assert message.entity_id == entity_id
+    assert message.project_name == project_name
+
+
+def test_attachment_to_message__invalid_data_raises_error():
+    """Test that invalid attachment data (not bytes, file, or base64) raises ValueError."""
+    url_override = "https://example.com"
+    entity_id = "123"
+    project_name = "test-project"
+
+    invalid_data = "not-a-file-and-not-base64!@#$%"
+
+    attachment_data = attachment.Attachment(
+        data=invalid_data,
+        file_name="test.txt",
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="Attachment data must be bytes, an existing file path, or a valid base64-encoded string",
+    ):
+        converters.attachment_to_message(
+            attachment_data=attachment_data,
+            entity_type="trace",
+            entity_id=entity_id,
+            project_name=project_name,
+            url_override=url_override,
+        )

sdks/python/tests/unit/api_objects/attachment/test_decoder_base64.py

@@ -0,0 +1,286 @@
+import base64
+import os
+import re
+
+import pytest
+
+from opik.api_objects.attachment import attachment, decoder_base64
+from opik.api_objects.attachment import decoder_helpers
+from . import constants
+
+
+@pytest.fixture
+def decoder():
+    """Create a Base64AttachmentDecoder instance."""
+    return decoder_base64.Base64AttachmentDecoder()
+
+
+def test_decode_png_success(decoder, files_to_remove):
+    """Test successful decoding of PNG image."""
+    result = decoder.decode(constants.PNG_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert isinstance(result, attachment.Attachment)
+    assert result.content_type == "image/png"
+    assert result.file_name.startswith("input-attachment-")
+    assert result.file_name.endswith(".png")
+    assert os.path.exists(result.data)
+
+
+def test_decode_jpeg_success(decoder, files_to_remove):
+    """Test successful decoding of JPEG image."""
+    result = decoder.decode(constants.JPEG_BASE64, context="output")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert isinstance(result, attachment.Attachment)
+    assert result.content_type == "image/jpeg"
+    assert result.file_name.startswith("output-attachment-")
+    assert result.file_name.endswith(".jpg") or result.file_name.endswith(".jpeg")
+    assert os.path.exists(result.data)
+
+
+def test_decode_pdf_success(decoder, files_to_remove):
+    """Test successful decoding of PDF document."""
+    result = decoder.decode(constants.PDF_BASE64, context="metadata")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert isinstance(result, attachment.Attachment)
+    assert result.content_type == "application/pdf"
+    assert result.file_name.startswith("metadata-attachment-")
+    assert result.file_name.endswith(".pdf")
+    assert os.path.exists(result.data)
+
+
+def test_decode_gif_success(decoder, files_to_remove):
+    """Test successful decoding of GIF image."""
+    result = decoder.decode(constants.GIF89_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert isinstance(result, attachment.Attachment)
+    assert result.content_type == "image/gif"
+    assert result.file_name.endswith(".gif")
+    assert os.path.exists(result.data)
+
+
+def test_decode_json_success(decoder, files_to_remove):
+    """Test successful decoding of JSON data."""
+    result = decoder.decode(constants.JSON_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert isinstance(result, attachment.Attachment)
+    assert result.content_type == "application/json"
+    assert result.file_name.endswith(".json")
+    assert os.path.exists(result.data)
+
+
+def test_decode_invalid_base64(decoder):
+    """Test that an invalid base64 string returns None."""
+    invalid_base64 = "this is not valid base64!@#$%"
+
+    result = decoder.decode(invalid_base64, context="input")
+
+    assert result is None
+
+
+def test_decode_non_string_input(decoder):
+    """Test that non-string input returns None."""
+    result = decoder.decode(12345, context="input")  # type: ignore
+
+    assert result is None
+
+
+def test_decode_dict_input(decoder):
+    """Test that dict input returns None."""
+    result = decoder.decode({"key": "value"}, context="input")  # type: ignore
+
+    assert result is None
+
+
+def test_decode_none_input(decoder):
+    """Test that None input returns None."""
+    result = decoder.decode(None, context="input")  # type: ignore
+
+    assert result is None
+
+
+def test_decode_empty_string(decoder):
+    """Test that empty string returns None."""
+    result = decoder.decode("", context="input")
+
+    assert result is None
+
+
+def test_decode_octet_stream_returns_none(decoder):
+    """Test that unrecognizable binary data (octet-stream) returns None."""
+    # Random binary data that won't match any known format
+    result = decoder.decode(constants.RANDOM_BINARY_BASE64, context="input")
+
+    assert result is None
+
+
+def test_decode_plain_text_returns_none(decoder):
+    """Test that plain text returns None."""
+    result = decoder.decode(constants.PLAIN_TEXT_BASE64, context="input")
+
+    assert result is None
+
+
+def test_decode_creates_temp_file_with_correct_extension(decoder, files_to_remove):
+    """Test that temporary file is created with the correct extension."""
+    result = decoder.decode(constants.PNG_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    # Check that the temp file path ends with .png (suffix parameter)
+    assert result.data.endswith("png")
+    assert os.path.exists(result.data)
+
+    # Verify file content
+    with open(result.data, "rb") as f:
+        content = f.read()
+        assert content[:8] == b"\x89PNG\r\n\x1a\n"
+
+
+def test_decode_attachment_properties(decoder, files_to_remove):
+    """Test that Attachment object has all required properties."""
+    result = decoder.decode(constants.PNG_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert hasattr(result, "data")
+    assert hasattr(result, "file_name")
+    assert hasattr(result, "content_type")
+    assert result.data is not None
+    assert result.file_name is not None
+    assert result.content_type is not None
+
+
+def test_decode_filename_format(decoder, files_to_remove):
+    """Test that the filename follows the expected format."""
+    result = decoder.decode(constants.PNG_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    # check that the filename matches a backend pattern
+    pattern = re.compile(decoder_helpers.ATTACHMENT_FILE_NAME_REGEX)
+    assert bool(pattern.fullmatch(result.file_name)) is True
+
+
+def test_decode_different_contexts(decoder, files_to_remove):
+    """Test that different contexts are reflected in the filename."""
+    contexts = ["input", "output", "metadata"]
+    results = []
+
+    for ctx in contexts:
+        result = decoder.decode(constants.PNG_BASE64, context=ctx)
+        assert result is not None
+        # Register for cleanup
+        files_to_remove.append(result.data)
+
+        assert result.file_name.startswith(f"{ctx}-attachment-")
+        results.append(result)
+
+
+def test_decode_multiple_calls_create_unique_files(decoder, files_to_remove):
+    """Test that multiple decode calls create unique filenames."""
+    result1 = decoder.decode(constants.PNG_BASE64, context="input")
+    result2 = decoder.decode(constants.PNG_BASE64, context="input")
+
+    assert result1 is not None
+    assert result2 is not None
+    # Register for cleanup
+    files_to_remove.append(result1.data)
+    files_to_remove.append(result2.data)
+
+    assert result1.file_name != result2.file_name
+    assert result1.data != result2.data
+    assert os.path.exists(result1.data)
+    assert os.path.exists(result2.data)
+
+
+def test_decode_webp_success(decoder, files_to_remove):
+    """Test successful decoding of WebP image."""
+    result = decoder.decode(constants.WEBP_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert result.content_type == "image/webp"
+    assert result.file_name.endswith(".webp")
+
+
+def test_decode_svg_success(decoder, files_to_remove):
+    """Test successful decoding of SVG image."""
+    result = decoder.decode(constants.SVG_BASE64, context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert result.content_type == "image/svg+xml"
+    assert result.file_name.endswith(".svg")
+
+
+def test_decode_short_base64_data(decoder):
+    """Test that very short base64 (< 4 bytes decoded) returns None."""
+    # Encode just 2 bytes
+    short_data = b"ab"
+    short_base64 = base64.b64encode(short_data).decode("utf-8")
+
+    result = decoder.decode(short_base64, context="input")
+
+    assert result is None
+
+
+def test_decode_base64_with_whitespace(decoder, files_to_remove):
+    """Test that base64 with whitespace is handled correctly."""
+    png_base64 = constants.PNG_BASE64
+    # Add whitespace
+    png_base64_with_whitespace = f"  {png_base64}  \n"
+
+    result = decoder.decode(png_base64_with_whitespace.strip(), context="input")
+
+    assert result is not None
+    # Register for cleanup
+    files_to_remove.append(result.data)
+
+    assert result.content_type == "image/png"
+
+
+def test_decode_corrupted_base64_padding(decoder, files_to_remove):
+    """Test base64 with incorrect padding.
+
+    Note: Python's base64 decoder is lenient and may still decode
+    strings with missing padding, so this test verifies proper handling.
+    """
+    # Valid base64 but with missing padding
+    corrupted_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJ"  # Missing padding
+
+    result = decoder.decode(corrupted_base64, context="input")
+
+    # Python's base64 decoder is lenient, so this might still succeed
+    # If it succeeds, verify the result is valid and clean up
+    if result is not None:
+        files_to_remove.append(result.data)

sdks/python/tests/unit/api_objects/attachment/test_decoder_helpers.py

@@ -0,0 +1,221 @@
+import base64
+
+from opik.api_objects.attachment import decoder_helpers
+
+from . import constants
+
+
+class TestDetectMimeType:
+    """Test suite for detect_mime_type() function."""
+
+    def test_detect_png(self):
+        """Test PNG image detection using magic bytes."""
+        # PNG magic bytes: 89 50 4E 47 0D 0A 1A 0A
+        png_data = b"\x89PNG\r\n\x1a\n" + b"\x00" * 100
+        assert decoder_helpers.detect_mime_type(png_data) == "image/png"
+
+    def test_detect_jpeg(self):
+        """Test JPEG image detection using magic bytes."""
+        # Real JPEG SOI (FFD8) followed by an APP0/JFIF marker (FFE0), then
+        # body bytes, then FFD9 end marker.
+        jpeg_data = b"\xff\xd8\xff\xe0" + b"\x00" * 100 + b"\xff\xd9"
+        assert decoder_helpers.detect_mime_type(jpeg_data) == "image/jpeg"
+
+    def test_detect_jpeg_without_end_marker(self):
+        """Test JPEG without proper end marker should not be detected as JPEG."""
+        # JPEG start but no proper end
+        jpeg_data = b"\xff\xd8" + b"\x00" * 100
+        assert decoder_helpers.detect_mime_type(jpeg_data) != "image/jpeg"
+
+    def test_detect_gif87a(self):
+        """Test GIF87a format detection."""
+        gif_data = b"GIF87a" + b"\x00" * 100
+        assert decoder_helpers.detect_mime_type(gif_data) == "image/gif"
+
+    def test_detect_gif89a(self):
+        """Test GIF89a format detection."""
+        assert decoder_helpers.detect_mime_type(constants.GIF89_BYTES) == "image/gif"
+
+    def test_detect_pdf(self):
+        """Test PDF document detection."""
+        assert (
+            decoder_helpers.detect_mime_type(constants.PDF_BYTES) == "application/pdf"
+        )
+
+    def test_detect_webp(self):
+        """Test WebP image detection."""
+        assert decoder_helpers.detect_mime_type(constants.WEBP_BYTES) == "image/webp"
+
+    def test_detect_svg(self):
+        """Test SVG image detection."""
+        assert decoder_helpers.detect_mime_type(constants.SVG_BYTES) == "image/svg+xml"
+
+    def test_detect_svg_with_uppercase(self):
+        """Test SVG detection with an uppercase tag."""
+        svg_data = b"<SVG>content</SVG>"
+        assert decoder_helpers.detect_mime_type(svg_data) == "image/svg+xml"
+
+    def test_detect_svg_with_mixed_case(self):
+        """Test SVG detection with mixed case tag."""
+        svg_data = b"<SvG>content</SvG>"
+        assert decoder_helpers.detect_mime_type(svg_data) == "image/svg+xml"
+
+    def test_detect_mp4(self):
+        """Test MP4 video detection."""
+        # MP4: ftyp at offset 4
+        mp4_data = b"\x00\x00\x00\x20ftypisom" + b"\x00" * 100
+        assert decoder_helpers.detect_mime_type(mp4_data) == "video/mp4"
+
+    def test_detect_json_object(self):
+        """Test JSON object detection."""
+        json_data = b'{"key": "value", "number": 123}'
+        assert decoder_helpers.detect_mime_type(json_data) == "application/json"
+
+    def test_detect_json_array(self):
+        """Test JSON array detection."""
+        json_data = b'["item1", "item2", "item3"]'
+        assert decoder_helpers.detect_mime_type(json_data) == "application/json"
+
+    def test_detect_json_with_whitespace(self):
+        """Test JSON detection with leading whitespace."""
+        json_data = b'   \n\t{"key": "value"}'
+        assert decoder_helpers.detect_mime_type(json_data) == "application/json"
+
+    def test_detect_invalid_json_like(self):
+        """Test that invalid JSON-like content is not detected as JSON."""
+        # UTF-8 decoding should fail
+        invalid_data = b"\xff\xfe{invalid}"
+        assert (
+            decoder_helpers.detect_mime_type(invalid_data) == "application/octet-stream"
+        )
+
+    def test_detect_short_data(self):
+        """Test that data shorter than 4 bytes returns octet-stream."""
+        short_data = b"abc"
+        assert (
+            decoder_helpers.detect_mime_type(short_data) == "application/octet-stream"
+        )
+
+    def test_detect_empty_data(self):
+        """Test that empty data returns octet-stream."""
+        empty_data = b""
+        assert (
+            decoder_helpers.detect_mime_type(empty_data) == "application/octet-stream"
+        )
+
+    def test_detect_unknown_binary(self):
+        """Test that unknown binary data returns octet-stream."""
+        assert (
+            decoder_helpers.detect_mime_type(constants.RANDOM_BINARY_BYTES)
+            == "application/octet-stream"
+        )
+
+    def test_detect_plain_text_not_json(self):
+        """Test that plain text (non-JSON) returns octet-stream."""
+        assert (
+            decoder_helpers.detect_mime_type(constants.PLAIN_TEXT_BYTES)
+            == "application/octet-stream"
+        )
+
+    def test_real_png_base64(self):
+        """Test with a real minimal PNG image (1x1 transparent pixel)."""
+        # 1x1 transparent PNG
+        png_data = base64.b64decode(constants.PNG_BASE64)
+        assert decoder_helpers.detect_mime_type(png_data) == "image/png"
+
+    def test_real_jpeg_base64(self):
+        """Test with a real minimal JPEG image."""
+        # Minimal valid JPEG (1x1 red pixel)
+        jpeg_data = base64.b64decode(constants.JPEG_BASE64)
+        assert decoder_helpers.detect_mime_type(jpeg_data) == "image/jpeg"
+
+    def test_detect_svg_with_long_content(self):
+        """Test SVG detection with content longer than 1024 bytes."""
+        # Create SVG content longer than 1024 bytes
+        long_content = "x" * 2000
+        svg_data = (
+            f'<svg xmlns="http://www.w3.org/2000/svg">{long_content}</svg>'.encode(
+                "utf-8"
+            )
+        )
+        assert decoder_helpers.detect_mime_type(svg_data) == "image/svg+xml"
+
+    def test_detect_webp_invalid(self):
+        """Test that RIFF without WEBP marker is not detected as WebP."""
+        # RIFF but not WEBP
+        riff_data = b"RIFF\x00\x00\x00\x00XXXX" + b"\x00" * 100
+        assert decoder_helpers.detect_mime_type(riff_data) != "image/webp"
+
+    def test_detect_mp4_short_data(self):
+        """Test that data too short for MP4 detection is not detected as MP4."""
+        short_mp4 = b"\x00\x00\x00\x20ftyp"  # Only 9 bytes
+        assert decoder_helpers.detect_mime_type(short_mp4) != "video/mp4"
+
+
+class TestGetFileExtension:
+    """Test suite for get_file_extension() function."""
+
+    def test_get_extension_png(self):
+        """Test PNG extension extraction."""
+        assert decoder_helpers.get_file_extension("image/png") == "png"
+
+    def test_get_extension_jpeg(self):
+        """Test JPEG extension extraction (should return 'jpg' not 'jpe')."""
+        result = decoder_helpers.get_file_extension("image/jpeg")
+        # The function should return 'jpg' for image/jpeg
+        assert result in ("jpg", "jpeg")
+
+    def test_get_extension_pdf(self):
+        """Test PDF extension extraction."""
+        assert decoder_helpers.get_file_extension("application/pdf") == "pdf"
+
+    def test_get_extension_gif(self):
+        """Test GIF extension extraction."""
+        assert decoder_helpers.get_file_extension("image/gif") == "gif"
+
+    def test_get_extension_svg(self):
+        """Test SVG extension extraction from 'image/svg+xml'."""
+        assert decoder_helpers.get_file_extension("image/svg+xml") == "svg"
+
+    def test_get_extension_webp(self):
+        """Test WebP extension extraction."""
+        assert decoder_helpers.get_file_extension("image/webp") == "webp"
+
+    def test_get_extension_mp4(self):
+        """Test MP4 extension extraction."""
+        assert decoder_helpers.get_file_extension("video/mp4") == "mp4"
+
+    def test_get_extension_json(self):
+        """Test JSON extension extraction."""
+        assert decoder_helpers.get_file_extension("application/json") == "json"
+
+    def test_get_extension_with_parameters(self):
+        """Test extension extraction with MIME type parameters."""
+        assert decoder_helpers.get_file_extension("image/jpeg; charset=utf-8") == "jpeg"
+
+    def test_get_extension_empty_string(self):
+        """Test that an empty MIME type returns 'bin'."""
+        assert decoder_helpers.get_file_extension("") == "bin"
+
+    def test_get_extension_none(self):
+        """Test that a None MIME type returns 'bin'."""
+        # This will fail the "if not mime_type:" check
+        result = decoder_helpers.get_file_extension("")
+        assert result == "bin"
+
+    def test_get_extension_invalid_format(self):
+        """Test that an invalid MIME type format returns 'bin'."""
+        assert decoder_helpers.get_file_extension("invalidmimetype") == "bin"
+
+    def test_get_extension_unknown_type(self):
+        """Test that an unknown MIME type extracts subtype."""
+        assert (
+            decoder_helpers.get_file_extension("application/x-custom-type")
+            == "x-custom-type"
+        )
+
+    def test_get_extension_removes_leading_dot(self):
+        """Test that leading dots are removed from extensions."""
+        # mimetypes.guess_extension can return extensions with leading dots
+        result = decoder_helpers.get_file_extension("text/plain")
+        assert not result.startswith(".")

sdks/python/tests/unit/api_objects/conversation/test_conversation_factory.py

@@ -0,0 +1,87 @@
+import datetime
+
+import pytest
+
+from opik.rest_api import TracePublic
+
+from opik.api_objects.conversation import conversation_factory
+
+
+@pytest.mark.parametrize(
+    "traces,expected_discussion",
+    [
+        (
+            [
+                TracePublic(
+                    input={"x": "test input 997"},
+                    output={"output": "test output"},
+                    start_time=datetime.datetime.now(),
+                )
+            ],
+            [
+                {"role": "user", "content": "test input 997"},
+                {"role": "assistant", "content": "test output"},
+            ],
+        ),
+        (  # test that traces are sorted by time - the first trace should be first
+            [
+                TracePublic(
+                    input={"x": "test input 3"},
+                    output={"output": "test output"},
+                    start_time=datetime.datetime.now(),
+                ),
+                TracePublic(
+                    input={"x": "test input 1"},
+                    output={"output": "test output"},
+                    start_time=datetime.datetime.now() - datetime.timedelta(seconds=10),
+                ),
+                TracePublic(
+                    input={"x": "test input 2"},
+                    output={"output": "test output"},
+                    start_time=datetime.datetime.now() - datetime.timedelta(seconds=5),
+                ),
+            ],
+            [
+                {"role": "user", "content": "test input 1"},
+                {"role": "assistant", "content": "test output"},
+                {"role": "user", "content": "test input 2"},
+                {"role": "assistant", "content": "test output"},
+                {"role": "user", "content": "test input 3"},
+                {"role": "assistant", "content": "test output"},
+            ],
+        ),
+        (  # test that trace's input or output are filtered out if it isn't in the expected format
+            [
+                TracePublic(
+                    input={"y": "test input 1"},  # wrong input
+                    output={"output": "test output"},
+                    start_time=datetime.datetime.now(),
+                ),
+                TracePublic(
+                    input={"x": "test input 2"},
+                    output={"result": "test output"},  # wrong output
+                    start_time=datetime.datetime.now(),
+                ),
+            ],
+            [
+                {"role": "assistant", "content": "test output"},
+                {"role": "user", "content": "test input 2"},
+            ],
+        ),
+    ],
+)
+def test_create_conversation_from_traces(traces, expected_discussion):
+    def input_transform(input):
+        if "x" not in input:
+            return None
+        return input["x"]
+
+    def output_transform(output):
+        if "output" not in output:
+            return None
+        return output["output"]
+
+    discussion = conversation_factory.create_conversation_from_traces(
+        traces, input_transform, output_transform
+    )
+    assert discussion.as_json_list() == expected_discussion

sdks/python/tests/unit/api_objects/dashboard/test_dashboard_client.py

@@ -0,0 +1,668 @@
+import copy
+
+import pytest
+from unittest.mock import Mock, patch
+
+from opik import exceptions
+from opik.api_objects import opik_client
+from opik.api_objects.dashboard import types
+from opik.api_objects.dashboard.dashboard import Dashboard
+from opik.rest_api.types import dashboard_public as dp
+
+
+class FakeDashboardsApi:
+    def __init__(self, store):
+        self._store = store
+        self.update_calls = []
+        self.deleted = False
+
+    def update_dashboard(
+        self, dashboard_id, *, name=None, type=None, description=None, config=None
+    ):
+        self.update_calls.append(
+            {"name": name, "type": type, "description": description, "config": config}
+        )
+        if config is not None:
+            self._store["config"] = config
+        if name is not None:
+            self._store["name"] = name
+        if description is not None:
+            self._store["description"] = description
+        return self._make()
+
+    def get_dashboard_by_id(self, dashboard_id):
+        return self._make()
+
+    def delete_dashboard(self, dashboard_id):
+        self.deleted = True
+
+    def _make(self):
+        return dp.DashboardPublic(
+            id=self._store["id"],
+            name=self._store["name"],
+            type=self._store.get("type"),
+            description=self._store.get("description"),
+            config=self._store["config"],
+        )
+
+
+class FakeRest:
+    def __init__(self, store):
+        self.dashboards = FakeDashboardsApi(store)
+
+
+def _make_dashboard(config, dashboard_type="multi_project", project_id=None):
+    store = {"id": "d1", "name": "My dash", "type": dashboard_type, "config": config}
+    rest = FakeRest(store)
+    public = dp.DashboardPublic(
+        id="d1",
+        name="My dash",
+        type=dashboard_type,
+        config=copy.deepcopy(config),
+        project_id=project_id,
+    )
+    client = Mock(spec=opik_client.Opik)
+    return (
+        Dashboard(dashboard_public=public, rest_client=rest, client=client),
+        rest,
+        store,
+    )
+
+
+def _config_with_legacy_widget():
+    return {
+        "version": 4,
+        "mysteryTopLevel": "keep-me",
+        "sections": [
+            {
+                "id": "s1",
+                "title": "Overview",
+                "widgets": [
+                    {
+                        "id": "old",
+                        "type": "future_widget",
+                        "config": {"unknownField": 123},
+                        "futureWidgetProp": "x",
+                    }
+                ],
+                "layout": [{"i": "old", "x": 0, "y": 0, "w": 2, "h": 2}],
+            }
+        ],
+        "lastModified": 111,
+    }
+
+
+def test_add_widget__preserves_unknown_fields_round_trip():
+    dashboard, rest, _ = _make_dashboard(_config_with_legacy_widget())
+
+    widget_id = dashboard.add_widget(
+        types.DashboardWidget(
+            type=types.WidgetType.TEXT_MARKDOWN,
+            title="Notes",
+            config=types.TextMarkdownConfig(content="hello"),
+        ),
+    )
+
+    config = dashboard.config
+    assert config["mysteryTopLevel"] == "keep-me"
+    old = next(w for w in config["sections"][0]["widgets"] if w["id"] == "old")
+    assert old["futureWidgetProp"] == "x"
+    assert old["config"]["unknownField"] == 123
+    assert widget_id in [w["id"] for w in config["sections"][0]["widgets"]]
+    assert widget_id in [i["i"] for i in config["sections"][0]["layout"]]
+    assert config["version"] == 4
+    assert config["lastModified"] != 111
+
+
+def test_add_widget__project_scoped_injects_project_id():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config, project_id="proj-xyz")
+
+    widget_id = dashboard.add_widget(
+        types.DashboardWidget(
+            type=types.WidgetType.PROJECT_STATS_CARD,
+            config=types.ProjectStatsCardConfig(
+                metric=types.StatsCardMetric.TRACE_COUNT
+            ),
+        ),
+    )
+
+    widget = dashboard.config["sections"][0]["widgets"][0]
+    assert widget["id"] == widget_id
+    assert widget["config"]["projectId"] == "proj-xyz"
+
+
+def test_add_widget__project_scoped_without_dashboard_project_raises():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config, project_id=None)
+
+    with pytest.raises(exceptions.DashboardValidationError, match="project-scoped"):
+        dashboard.add_widget(
+            types.DashboardWidget(type=types.WidgetType.PROJECT_STATS_CARD),
+        )
+
+    assert rest.dashboards.update_calls == []
+
+
+def test_add_widget__incompatible_type_raises_before_write():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config, dashboard_type="multi_project")
+
+    with pytest.raises(exceptions.DashboardValidationError, match="not supported"):
+        dashboard.add_widget(
+            types.DashboardWidget(type=types.WidgetType.EXPERIMENT_LEADERBOARD)
+        )
+
+    assert rest.dashboards.update_calls == []
+
+
+def test_add_widget__unknown_section_raises():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config)
+
+    with pytest.raises(exceptions.DashboardValidationError, match="not found"):
+        dashboard.add_widget(
+            types.DashboardWidget(type=types.WidgetType.TEXT_MARKDOWN),
+            section_id="missing",
+        )
+
+
+def test_mutation__unknown_version_refused():
+    config = {
+        "version": 5,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config)
+
+    with pytest.raises(exceptions.DashboardValidationError, match="schema version 5"):
+        dashboard.add_widget(types.DashboardWidget(type=types.WidgetType.TEXT_MARKDOWN))
+    assert rest.dashboards.update_calls == []
+
+
+def test_reads_never_fail__on_unknown_widget_type():
+    dashboard, _, _ = _make_dashboard(_config_with_legacy_widget())
+
+    state = dashboard.state
+    assert len(state.sections) == 1
+    assert state.sections[0].widgets[0].type == "future_widget"
+
+
+def test_update_widget__merges_config():
+    config = {
+        "version": 4,
+        "sections": [
+            {
+                "id": "s1",
+                "title": "t",
+                "widgets": [
+                    {"id": "w1", "type": "text_markdown", "config": {"content": "a"}}
+                ],
+                "layout": [{"i": "w1", "x": 0, "y": 0, "w": 2, "h": 4}],
+            }
+        ],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config)
+
+    dashboard.update_widget("w1", title="New", config={"content": "b"})
+
+    widget = dashboard.config["sections"][0]["widgets"][0]
+    assert widget["title"] == "New"
+    assert widget["config"]["content"] == "b"
+    assert len(rest.dashboards.update_calls) == 1
+    assert rest.dashboards.update_calls[0]["config"] is not None
+
+
+def test_remove_widget__removes_widget_and_layout():
+    config = {
+        "version": 4,
+        "sections": [
+            {
+                "id": "s1",
+                "title": "t",
+                "widgets": [
+                    {"id": "w1", "type": "text_markdown", "config": {}},
+                    {"id": "w2", "type": "text_markdown", "config": {}},
+                ],
+                "layout": [
+                    {"i": "w1", "x": 0, "y": 0, "w": 2, "h": 4},
+                    {"i": "w2", "x": 2, "y": 0, "w": 2, "h": 4},
+                ],
+            }
+        ],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config)
+
+    dashboard.remove_widget("w1")
+
+    section = dashboard.config["sections"][0]
+    assert [w["id"] for w in section["widgets"]] == ["w2"]
+    assert [i["i"] for i in section["layout"]] == ["w2"]
+
+
+def test_remove_widget__missing_raises():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config)
+    with pytest.raises(exceptions.DashboardValidationError, match="not found"):
+        dashboard.remove_widget("ghost")
+
+
+def test_rename__sends_name_only():
+    dashboard, rest, _ = _make_dashboard(_config_with_legacy_widget())
+    dashboard.rename("Renamed")
+    assert dashboard.name == "Renamed"
+    assert rest.dashboards.update_calls[-1]["name"] == "Renamed"
+    assert rest.dashboards.update_calls[-1]["config"] is None
+
+
+def test_set_description__sends_description_only():
+    dashboard, rest, _ = _make_dashboard(_config_with_legacy_widget())
+    dashboard.set_description("desc")
+    assert dashboard.description == "desc"
+    assert rest.dashboards.update_calls[-1]["description"] == "desc"
+    assert rest.dashboards.update_calls[-1]["config"] is None
+
+
+def test_add_section__appends_and_returns_id():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config)
+    new_id = dashboard.add_section("Second")
+    section_ids = [s["id"] for s in dashboard.config["sections"]]
+    assert new_id in section_ids
+    assert len(section_ids) == 2
+
+
+def test_delete__calls_rest():
+    dashboard, rest, _ = _make_dashboard(_config_with_legacy_widget())
+    dashboard.delete()
+    assert rest.dashboards.deleted is True
+
+
+# --- Opik client-level methods ---
+
+
+def test_create_dashboard__default_section_and_config():
+    client = opik_client.Opik()
+    captured = {}
+
+    def fake_create(*, name, config, type, description, project_id, project_name):
+        captured.update(
+            name=name,
+            config=config,
+            type=type,
+            description=description,
+            project_id=project_id,
+            project_name=project_name,
+        )
+        return dp.DashboardPublic(id="d1", name=name, type=type, config=config)
+
+    with patch.object(
+        client._rest_client.dashboards, "create_dashboard", side_effect=fake_create
+    ):
+        dashboard = client.create_dashboard(
+            name="Prod", type=types.DashboardType.MULTI_PROJECT
+        )
+
+    assert dashboard.name == "Prod"
+    assert captured["type"] == "multi_project"
+    assert captured["config"]["version"] == types.DASHBOARD_VERSION
+    assert len(captured["config"]["sections"]) == 1
+    assert captured["config"]["sections"][0]["title"] == "Overview"
+    assert "lastModified" in captured["config"]
+
+
+def test_create_dashboard__with_provided_sections():
+    client = opik_client.Opik()
+    captured = {}
+
+    def fake_create(*, name, config, type, description, project_id, project_name):
+        captured["config"] = config
+        return dp.DashboardPublic(id="d1", name=name, type=type, config=config)
+
+    section = types.DashboardSection(
+        title="Custom",
+        widgets=[types.DashboardWidget(id="w1", type=types.WidgetType.TEXT_MARKDOWN)],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=2, h=4)],
+    )
+
+    with patch.object(
+        client._rest_client.dashboards, "create_dashboard", side_effect=fake_create
+    ):
+        client.create_dashboard(name="Prod", sections=[section])
+
+    assert captured["config"]["sections"][0]["title"] == "Custom"
+    assert captured["config"]["sections"][0]["widgets"][0]["id"] == "w1"
+
+
+def test_create_dashboard__project_scoped_widget_with_project_id_injects():
+    client = opik_client.Opik()
+
+    def fake_create(*, name, config, type, description, project_id, project_name):
+        return dp.DashboardPublic(id="d1", name=name, type=type, config=config)
+
+    section = types.DashboardSection(
+        title="S",
+        widgets=[
+            types.DashboardWidget(
+                id="w1",
+                type=types.WidgetType.PROJECT_STATS_CARD,
+                config=types.ProjectStatsCardConfig(
+                    metric=types.StatsCardMetric.TRACE_COUNT
+                ),
+            )
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=1, h=2)],
+    )
+
+    with patch.object(
+        client._rest_client.dashboards, "create_dashboard", side_effect=fake_create
+    ):
+        dashboard = client.create_dashboard(
+            name="Prod",
+            type=types.DashboardType.MULTI_PROJECT,
+            project_id="proj-123",
+            sections=[section],
+        )
+
+    widget = dashboard.config["sections"][0]["widgets"][0]
+    assert widget["config"]["projectId"] == "proj-123"
+
+
+def test_create_dashboard__project_scoped_widget_with_project_name_only_does_not_raise():
+    client = opik_client.Opik()
+
+    def fake_create(*, name, config, type, description, project_id, project_name):
+        return dp.DashboardPublic(id="d1", name=name, type=type, config=config)
+
+    section = types.DashboardSection(
+        title="S",
+        widgets=[
+            types.DashboardWidget(
+                id="w1",
+                type=types.WidgetType.PROJECT_STATS_CARD,
+                config=types.ProjectStatsCardConfig(
+                    metric=types.StatsCardMetric.TRACE_COUNT
+                ),
+            )
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=1, h=2)],
+    )
+
+    # project_name provided but not project_id — should not raise
+    with patch.object(
+        client._rest_client.dashboards, "create_dashboard", side_effect=fake_create
+    ):
+        client.create_dashboard(
+            name="Prod",
+            type=types.DashboardType.MULTI_PROJECT,
+            project_name="Default Project",
+            sections=[section],
+        )
+
+
+def test_create_dashboard__project_scoped_widget_without_any_project_raises():
+    client = opik_client.Opik()
+
+    section = types.DashboardSection(
+        title="S",
+        widgets=[
+            types.DashboardWidget(id="w1", type=types.WidgetType.PROJECT_STATS_CARD)
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=1, h=2)],
+    )
+
+    with patch.object(
+        client._rest_client.dashboards, "create_dashboard"
+    ) as mock_create:
+        with pytest.raises(exceptions.DashboardValidationError, match="project-scoped"):
+            client.create_dashboard(
+                name="Prod",
+                type=types.DashboardType.MULTI_PROJECT,
+                sections=[section],
+            )
+
+    mock_create.assert_not_called()
+
+
+def test_create_dashboard__incompatible_widget_raises_before_create():
+    client = opik_client.Opik()
+
+    section = types.DashboardSection(
+        title="Custom",
+        widgets=[
+            types.DashboardWidget(id="w1", type=types.WidgetType.EXPERIMENT_LEADERBOARD)
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=6, h=6)],
+    )
+
+    with patch.object(
+        client._rest_client.dashboards, "create_dashboard"
+    ) as mock_create:
+        with pytest.raises(exceptions.DashboardValidationError, match="not supported"):
+            client.create_dashboard(
+                name="Prod",
+                type=types.DashboardType.MULTI_PROJECT,
+                sections=[section],
+            )
+
+    mock_create.assert_not_called()
+
+
+def test_replace_sections__injects_project_id_for_project_scoped_widgets():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config, project_id="proj-abc")
+
+    new_section = types.DashboardSection(
+        title="t",
+        id="s1",
+        widgets=[
+            types.DashboardWidget(
+                id="w1",
+                type=types.WidgetType.PROJECT_STATS_CARD,
+                config=types.ProjectStatsCardConfig(
+                    metric=types.StatsCardMetric.TRACE_COUNT
+                ),
+            )
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=1, h=2)],
+    )
+    dashboard.replace_sections([new_section])
+
+    widget = dashboard.config["sections"][0]["widgets"][0]
+    assert widget["config"]["projectId"] == "proj-abc"
+
+
+def test_replace_sections__project_scoped_widget_without_dashboard_project_raises():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config, project_id=None)
+
+    new_section = types.DashboardSection(
+        title="t",
+        id="s1",
+        widgets=[
+            types.DashboardWidget(id="w1", type=types.WidgetType.PROJECT_STATS_CARD)
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=1, h=2)],
+    )
+
+    with pytest.raises(exceptions.DashboardValidationError, match="project-scoped"):
+        dashboard.replace_sections([new_section])
+
+    assert rest.dashboards.update_calls == []
+
+
+def test_replace_sections__incompatible_widget_raises_before_write():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config, dashboard_type="multi_project")
+
+    section = types.DashboardSection(
+        title="Custom",
+        widgets=[
+            types.DashboardWidget(id="w1", type=types.WidgetType.EXPERIMENT_LEADERBOARD)
+        ],
+        layout=[types.DashboardLayoutItem(id="w1", x=0, y=0, w=6, h=6)],
+    )
+
+    with pytest.raises(exceptions.DashboardValidationError, match="not supported"):
+        dashboard.replace_sections([section])
+
+    assert rest.dashboards.update_calls == []
+
+
+def test_get_dashboard__returns_wrapper():
+    client = opik_client.Opik()
+    public = dp.DashboardPublic(
+        id="d1",
+        name="My dash",
+        type="multi_project",
+        config={"version": 4, "sections": []},
+    )
+    with patch.object(
+        client._rest_client.dashboards, "get_dashboard_by_id", return_value=public
+    ):
+        dashboard = client.get_dashboard("d1")
+    assert dashboard.id == "d1"
+    assert dashboard.name == "My dash"
+
+
+def test_delete_dashboard__calls_rest():
+    client = opik_client.Opik()
+    with patch.object(
+        client._rest_client.dashboards, "delete_dashboard"
+    ) as mock_delete:
+        client.delete_dashboard("d1")
+        mock_delete.assert_called_once_with("d1")
+
+
+def test_add_widget__null_id_is_replaced():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config)
+
+    widget_id = dashboard.add_widget(
+        {"id": None, "type": "text_markdown", "config": {"content": "hi"}}
+    )
+
+    assert widget_id != "None"
+    assert widget_id is not None
+    layout_ids = [i["i"] for i in dashboard.config["sections"][0]["layout"]]
+    assert widget_id in layout_ids
+
+
+def test_add_widget__partial_size_raises():
+    config = {
+        "version": 4,
+        "sections": [{"id": "s1", "title": "t", "widgets": [], "layout": []}],
+        "lastModified": 1,
+    }
+    dashboard, rest, _ = _make_dashboard(config)
+
+    with pytest.raises(exceptions.DashboardValidationError, match="'w' and 'h'"):
+        dashboard.add_widget(
+            types.DashboardWidget(type=types.WidgetType.TEXT_MARKDOWN),
+            size={"w": 2},
+        )
+
+    assert rest.dashboards.update_calls == []
+
+
+def test_update_widget__null_config_from_api_does_not_raise():
+    config = {
+        "version": 4,
+        "sections": [
+            {
+                "id": "s1",
+                "title": "t",
+                "widgets": [{"id": "w1", "type": "text_markdown", "config": None}],
+                "layout": [{"i": "w1", "x": 0, "y": 0, "w": 2, "h": 4}],
+            }
+        ],
+        "lastModified": 1,
+    }
+    dashboard, _, _ = _make_dashboard(config)
+
+    dashboard.update_widget("w1", config={"content": "new"})
+
+    widget = dashboard.config["sections"][0]["widgets"][0]
+    assert widget["config"]["content"] == "new"
+
+
+def test_find_dashboards__paginates_and_respects_max_results():
+    from opik.api_objects.dashboard import rest_operations
+    from opik.rest_api.types import dashboard_page_public as dpp
+
+    pages = {
+        1: [
+            dp.DashboardPublic(
+                id=f"d{i}", name=f"d{i}", config={"version": 4, "sections": []}
+            )
+            for i in range(100)
+        ],
+        2: [
+            dp.DashboardPublic(
+                id="d100", name="d100", config={"version": 4, "sections": []}
+            )
+        ],
+    }
+
+    class PaginatedApi:
+        def __init__(self):
+            self.calls = []
+
+        def find_dashboards(self, *, page, size, name, project_id, sorting, filters):
+            self.calls.append(page)
+            return dpp.DashboardPagePublic(content=pages.get(page, []))
+
+    class Rest:
+        def __init__(self):
+            self.dashboards = PaginatedApi()
+
+    rest = Rest()
+    result = rest_operations.find_dashboards(
+        rest_client=rest, client=Mock(spec=opik_client.Opik), max_results=100
+    )
+    assert len(result) == 100
+    # only the first page is needed to satisfy max_results
+    assert rest.dashboards.calls == [1]
+    assert all(isinstance(d, Dashboard) for d in result)

sdks/python/tests/unit/api_objects/dashboard/test_layout.py

@@ -0,0 +1,191 @@
+"""Ported from apps/opik-frontend/src/lib/dashboard/layout.test.ts.
+
+Keeps the SDK auto-layout behavior identical to the frontend so widgets added
+via the SDK are positioned the same way the UI would position them.
+"""
+
+import pytest
+
+from opik.api_objects.dashboard import layout
+from opik.api_objects.dashboard.types import DashboardLayoutItem, WidgetType
+
+
+def _item(id: str, x: int, y: int, w: int, h: int) -> DashboardLayoutItem:
+    return DashboardLayoutItem(id=id, x=x, y=y, w=w, h=h)
+
+
+def test_grid_constants__match_frontend():
+    assert layout.GRID_COLUMNS == 6
+    assert layout.MAX_WIDGET_HEIGHT == 12
+    assert layout.MIN_WIDGET_WIDTH == 1
+    assert layout.MIN_WIDGET_HEIGHT == 1
+
+
+@pytest.mark.parametrize(
+    "widget_type,expected",
+    [
+        (WidgetType.PROJECT_METRICS.value, {"w": 2, "h": 4, "minW": 2, "minH": 4}),
+        (WidgetType.PROJECT_STATS_CARD.value, {"w": 1, "h": 2, "minW": 1, "minH": 2}),
+        (WidgetType.TEXT_MARKDOWN.value, {"w": 2, "h": 4, "minW": 1, "minH": 4}),
+        ("unknown-type", {"w": 2, "h": 2, "minW": 1, "minH": 1}),
+    ],
+)
+def test_get_widget_size_config__per_type(widget_type, expected):
+    assert layout.get_widget_size_config(widget_type) == expected
+
+
+@pytest.mark.parametrize(
+    "layout_items,expected",
+    [
+        ([], [0, 0, 0, 0, 0, 0]),
+        ([_item("w1", 0, 0, 2, 3)], [3, 3, 0, 0, 0, 0]),
+        (
+            [_item("w1", 0, 0, 2, 3), _item("w2", 2, 0, 2, 5)],
+            [3, 3, 5, 5, 0, 0],
+        ),
+        (
+            [_item("w1", 0, 0, 2, 3), _item("w2", 1, 3, 2, 2)],
+            [3, 5, 5, 0, 0, 0],
+        ),
+        ([_item("w1", 0, 0, 6, 4)], [4, 4, 4, 4, 4, 4]),
+        (
+            [_item("w1", 0, 2, 2, 3), _item("w2", 2, 0, 2, 2)],
+            [5, 5, 2, 2, 0, 0],
+        ),
+    ],
+)
+def test_get_column_heights(layout_items, expected):
+    assert layout.get_column_heights(layout_items) == expected
+
+
+@pytest.mark.parametrize(
+    "w,h,heights,expected",
+    [
+        (2, 3, [0, 0, 0, 0, 0, 0], {"x": 0, "y": 0}),
+        (2, 2, [3, 3, 0, 0, 0, 0], {"x": 2, "y": 0}),
+        (2, 2, [3, 3, 5, 5, 4, 4], {"x": 0, "y": 3}),
+        (6, 1, [2, 2, 3, 3, 2, 2], {"x": 0, "y": 3}),
+        (2, 3, [5, 5, 2, 2, 8, 8], {"x": 2, "y": 2}),
+        (1, 2, [3, 0, 2, 4, 1, 5], {"x": 1, "y": 0}),
+    ],
+)
+def test_find_first_available_position(w, h, heights, expected):
+    assert layout.find_first_available_position(w, h, heights) == expected
+
+
+def test_calculate_layout_for_adding_widget__first_widget_to_empty_layout():
+    result = layout.calculate_layout_for_adding_widget(
+        [], WidgetType.TEXT_MARKDOWN.value, "widget-1"
+    )
+    assert len(result) == 1
+    assert result[0].id == "widget-1"
+    assert (result[0].x, result[0].y, result[0].w, result[0].h) == (0, 0, 2, 4)
+
+
+def test_calculate_layout_for_adding_widget__custom_size():
+    result = layout.calculate_layout_for_adding_widget(
+        [], WidgetType.TEXT_MARKDOWN.value, "widget-1", size={"w": 3, "h": 5}
+    )
+    assert result[0].w == 3
+    assert result[0].h == 5
+
+
+def test_calculate_layout_for_adding_widget__clamps_oversized_custom_size():
+    result = layout.calculate_layout_for_adding_widget(
+        [],
+        WidgetType.TEXT_MARKDOWN.value,
+        "widget-1",
+        size={"w": 100, "h": 100},
+    )
+    assert result[0].w == layout.GRID_COLUMNS
+    assert result[0].h == layout.MAX_WIDGET_HEIGHT
+
+
+def test_calculate_layout_for_adding_widget__clamps_undersized_custom_size():
+    result = layout.calculate_layout_for_adding_widget(
+        [],
+        WidgetType.PROJECT_METRICS.value,
+        "widget-1",
+        size={"w": 1, "h": 1},
+    )
+    size_config = layout.get_widget_size_config(WidgetType.PROJECT_METRICS.value)
+    assert result[0].w == size_config["minW"]
+    assert result[0].h == size_config["minH"]
+
+
+def test_calculate_layout_for_adding_widget__second_widget_placed_next_to_first():
+    existing = [_item("widget-1", 0, 0, 2, 4)]
+    result = layout.calculate_layout_for_adding_widget(
+        existing, WidgetType.TEXT_MARKDOWN.value, "widget-2"
+    )
+    assert len(result) == 2
+    assert result[1].id == "widget-2"
+    assert (result[1].x, result[1].y) == (2, 0)
+
+
+def test_calculate_layout_for_adding_widget__sets_min_and_max_constraints():
+    result = layout.calculate_layout_for_adding_widget(
+        [], WidgetType.PROJECT_METRICS.value, "widget-1"
+    )
+    assert result[0].min_w == 2
+    assert result[0].min_h == 4
+    assert result[0].max_w == layout.GRID_COLUMNS
+    assert result[0].max_h == layout.MAX_WIDGET_HEIGHT
+
+
+def test_calculate_layout_for_adding_widget__optimal_position_with_complex_layout():
+    existing = [_item("widget-1", 0, 0, 3, 4), _item("widget-2", 3, 0, 3, 2)]
+    result = layout.calculate_layout_for_adding_widget(
+        existing, WidgetType.PROJECT_STATS_CARD.value, "widget-3"
+    )
+    assert result[2].y == 2
+
+
+def test_normalize_layout__empty():
+    assert layout.normalize_layout([]) == []
+
+
+def test_normalize_layout__clamps_position_to_grid_bounds():
+    normalized = layout.normalize_layout([_item("w1", 7, -1, 2, 3)])
+    assert normalized[0].x == 4
+    assert normalized[0].y == 0
+
+
+def test_normalize_layout__clamps_width_to_grid_columns():
+    normalized = layout.normalize_layout([_item("w1", 0, 0, 10, 3)])
+    assert normalized[0].w == layout.GRID_COLUMNS
+
+
+def test_normalize_layout__clamps_height_to_max():
+    normalized = layout.normalize_layout([_item("w1", 0, 0, 2, 20)])
+    assert normalized[0].h == layout.MAX_WIDGET_HEIGHT
+
+
+def test_normalize_layout__enforces_minimum_dimensions():
+    normalized = layout.normalize_layout([_item("w1", 0, 0, 0, 0)])
+    assert normalized[0].w == layout.MIN_WIDGET_WIDTH
+    assert normalized[0].h == layout.MIN_WIDGET_HEIGHT
+
+
+def test_normalize_layout__applies_widget_specific_constraints():
+    widgets = [{"id": "w1", "type": WidgetType.PROJECT_METRICS.value}]
+    normalized = layout.normalize_layout([_item("w1", 0, 0, 1, 2)], widgets)
+    assert normalized[0].min_w == 2
+    assert normalized[0].min_h == 4
+    assert normalized[0].w == 2
+    assert normalized[0].h == 4
+
+
+def test_remove_widget_from_layout():
+    items = [_item("w1", 0, 0, 2, 3), _item("w2", 2, 0, 2, 3), _item("w3", 4, 0, 2, 3)]
+    result = layout.remove_widget_from_layout(items, "w2")
+    assert [item.id for item in result] == ["w1", "w3"]
+    # original not mutated
+    assert len(items) == 3
+
+
+def test_remove_widget_from_layout__non_existent_id_is_noop():
+    items = [_item("w1", 0, 0, 2, 3)]
+    result = layout.remove_widget_from_layout(items, "ghost")
+    assert len(result) == 1
+    assert result[0].id == "w1"

sdks/python/tests/unit/api_objects/dashboard/test_types.py

@@ -0,0 +1,104 @@
+import pytest
+
+from opik import exceptions
+from opik.api_objects.dashboard import types
+from ....testlib import assert_equal
+
+
+def test_widget_serialization__uses_camelcase_and_enum_values():
+    widget = types.DashboardWidget(
+        type=types.WidgetType.PROJECT_STATS_CARD,
+        title="Traces",
+        config=types.ProjectStatsCardConfig(metric=types.StatsCardMetric.TRACE_COUNT),
+    )
+    result = widget.to_jsonable()
+
+    assert result["type"] == "project_stats_card"
+    assert result["title"] == "Traces"
+    assert "id" in result
+    assert_equal(
+        {
+            "source": "traces",
+            "metric": "trace_count",
+        },
+        result["config"],
+    )
+
+
+def test_project_metrics_config__defaults_and_breakdown_camelcase():
+    config = types.ProjectMetricsConfig(
+        metric_type=types.ProjectMetricType.DURATION,
+        breakdown=types.BreakdownConfig(
+            field=types.BreakdownField.METADATA, metadata_key="provider"
+        ),
+    )
+    result = config.to_jsonable()
+
+    assert result["metricType"] == "DURATION"
+    assert result["chartType"] == "line"
+    assert_equal({"field": "metadata", "metadataKey": "provider"}, result["breakdown"])
+
+
+def test_breakdown_config__metadata_field_requires_metadata_key():
+    with pytest.raises(exceptions.DashboardValidationError):
+        types.BreakdownConfig(field=types.BreakdownField.METADATA)
+
+
+def test_breakdown_config__non_metadata_field_does_not_require_key():
+    config = types.BreakdownConfig(field=types.BreakdownField.TAGS)
+    assert config.to_jsonable() == {"field": "tags"}
+
+
+def test_leaderboard_config__enable_ranking_requires_ranking_metric():
+    with pytest.raises(exceptions.DashboardValidationError):
+        types.ExperimentLeaderboardConfig(enable_ranking=True)
+
+
+def test_leaderboard_config__enable_ranking_with_metric_ok():
+    config = types.ExperimentLeaderboardConfig(
+        enable_ranking=True, ranking_metric="pass_rate"
+    )
+    assert config.to_jsonable()["rankingMetric"] == "pass_rate"
+
+
+@pytest.mark.parametrize("value", [0, 101, 500, "0", "200"])
+def test_leaderboard_config__max_rows_out_of_range_raises(value):
+    with pytest.raises(exceptions.DashboardValidationError):
+        types.ExperimentLeaderboardConfig(max_rows=value)
+
+
+@pytest.mark.parametrize("value", [1, 100, "10", 50])
+def test_leaderboard_config__max_rows_in_range_ok(value):
+    config = types.ExperimentLeaderboardConfig(max_rows=value)
+    assert config.to_jsonable()["maxRows"] == value
+
+
+@pytest.mark.parametrize("value", [0, 101])
+def test_feedback_scores_config__max_experiments_out_of_range_raises(value):
+    with pytest.raises(exceptions.DashboardValidationError):
+        types.ExperimentsFeedbackScoresConfig(max_experiments_count=value)
+
+
+def test_dashboard_state__defaults_to_known_version():
+    state = types.DashboardState()
+    result = state.to_jsonable()
+    assert result["version"] == types.DASHBOARD_VERSION
+    assert "lastModified" in result
+    assert result["sections"] == []
+
+
+def test_widget_accepts_raw_dict_config():
+    widget = types.DashboardWidget(type="future_widget", config={"someCamelField": 1})
+    assert widget.to_jsonable()["config"] == {"someCamelField": 1}
+
+
+def test_models_preserve_unknown_fields_on_parse():
+    state = types.DashboardState.model_validate(
+        {
+            "version": 4,
+            "lastModified": 1,
+            "sections": [],
+            "mysteryField": "keep",
+        }
+    )
+    assert state.to_jsonable()["mysteryField"] == "keep"

sdks/python/tests/unit/api_objects/dashboard/test_validation.py

@@ -0,0 +1,169 @@
+from unittest import mock
+
+import pytest
+
+from opik import exceptions
+from opik.api_objects.dashboard import types, validation
+
+
+def _section(section_id, widgets, layout_ids):
+    return {
+        "id": section_id,
+        "title": "t",
+        "widgets": [{"id": w, "type": "text_markdown", "config": {}} for w in widgets],
+        "layout": [{"i": i, "x": 0, "y": 0, "w": 1, "h": 1} for i in layout_ids],
+    }
+
+
+def test_validate_structure__valid_passes():
+    state = {"version": 4, "sections": [_section("s1", ["w1"], ["w1"])]}
+    validation.validate_structure(state)
+
+
+def test_validate_structure__widget_without_layout_raises():
+    state = {"version": 4, "sections": [_section("s1", ["w1"], [])]}
+    with pytest.raises(exceptions.DashboardValidationError, match="without a layout"):
+        validation.validate_structure(state)
+
+
+def test_validate_structure__orphan_layout_item_raises():
+    state = {"version": 4, "sections": [_section("s1", [], ["ghost"])]}
+    with pytest.raises(exceptions.DashboardValidationError, match="missing widgets"):
+        validation.validate_structure(state)
+
+
+def test_validate_structure__duplicate_widget_ids_across_sections_raises():
+    state = {
+        "version": 4,
+        "sections": [
+            _section("s1", ["dup"], ["dup"]),
+            _section("s2", ["dup"], ["dup"]),
+        ],
+    }
+    with pytest.raises(exceptions.DashboardValidationError, match="Duplicate widget"):
+        validation.validate_structure(state)
+
+
+def test_validate_structure__duplicate_section_ids_raises():
+    state = {
+        "version": 4,
+        "sections": [_section("s1", ["w1"], ["w1"]), _section("s1", ["w2"], ["w2"])],
+    }
+    with pytest.raises(exceptions.DashboardValidationError, match="Duplicate section"):
+        validation.validate_structure(state)
+
+
+def test_validate_widget_for_dashboard__incompatible_type_raises():
+    widget = {"type": "experiment_leaderboard", "config": {}}
+    with pytest.raises(exceptions.DashboardValidationError, match="not supported"):
+        validation.validate_widget_for_dashboard(widget, "multi_project")
+
+
+def test_validate_widget_for_dashboard__compatible_type_ok():
+    widget = {"type": "project_metrics", "config": {"metricType": "TRACE_COUNT"}}
+    validation.validate_widget_for_dashboard(widget, "multi_project")
+
+
+def test_validate_widget_for_dashboard__no_dashboard_type_skips_compatibility():
+    widget = {"type": "experiment_leaderboard", "config": {}}
+    validation.validate_widget_for_dashboard(widget, None)
+
+
+def test_validate_widget_for_dashboard__unknown_metric_warns_not_raises():
+    widget = {"type": "project_metrics", "config": {"metricType": "not_a_metric"}}
+    with mock.patch.object(validation.LOGGER, "warning") as mock_warning:
+        validation.validate_widget_for_dashboard(widget, "multi_project")
+    mock_warning.assert_called_once()
+    assert "Unknown project_metrics metricType" in mock_warning.call_args[0][0]
+
+
+def test_validate_widget_for_dashboard__dynamic_feedback_metric_no_warning():
+    widget = {
+        "type": "project_stats_card",
+        "config": {"metric": "feedback_scores.helpfulness"},
+    }
+    with mock.patch.object(validation.LOGGER, "warning") as mock_warning:
+        validation.validate_widget_for_dashboard(widget, "multi_project")
+    mock_warning.assert_not_called()
+
+
+def test_validate_writable_version__known_version_ok():
+    validation.validate_writable_version(4)
+    validation.validate_writable_version(None)
+
+
+def test_validate_writable_version__unknown_version_raises():
+    with pytest.raises(exceptions.DashboardValidationError, match="schema version 5"):
+        validation.validate_writable_version(5)
+
+
+def test_as_widget_dict__from_model__returns_jsonable():
+    widget = types.DashboardWidget(
+        id="w1", type=types.WidgetType.TEXT_MARKDOWN, config=types.TextMarkdownConfig()
+    )
+    result = validation.as_widget_dict(widget)
+    assert result["id"] == "w1"
+    assert result["type"] == "text_markdown"
+
+
+def test_as_widget_dict__from_dict__returns_same_dict():
+    d = {"id": "w1", "type": "text_markdown", "config": {}}
+    result = validation.as_widget_dict(d)
+    assert result is d
+
+
+def test_as_widget_dict__invalid_type__raises():
+    with pytest.raises(
+        exceptions.DashboardValidationError, match="Expected a DashboardWidget"
+    ):
+        validation.as_widget_dict(42)
+
+
+def test_as_section_dicts__from_models__returns_jsonable():
+    section = types.DashboardSection(title="S1")
+    result = validation.as_section_dicts([section])
+    assert len(result) == 1
+    assert result[0]["title"] == "S1"
+    assert "id" in result[0]
+
+
+def test_as_section_dicts__from_dicts__returns_same_dicts():
+    d = {"id": "s1", "title": "S1", "widgets": [], "layout": []}
+    result = validation.as_section_dicts([d])
+    assert result[0] is d
+
+
+def test_as_section_dicts__invalid_type__raises():
+    with pytest.raises(
+        exceptions.DashboardValidationError, match="Expected a DashboardSection"
+    ):
+        validation.as_section_dicts(["not-a-section"])
+
+
+def test_inject_project_id__project_scoped_widget_injects():
+    widget = {"type": "project_stats_card", "config": {"metric": "trace_count"}}
+    validation.inject_project_id(widget, "proj-123")
+    assert widget["config"]["projectId"] == "proj-123"
+
+
+def test_inject_project_id__project_metrics_injects():
+    widget = {"type": "project_metrics", "config": {"metricType": "DURATION"}}
+    validation.inject_project_id(widget, "proj-abc")
+    assert widget["config"]["projectId"] == "proj-abc"
+
+
+def test_inject_project_id__non_project_widget_untouched():
+    widget = {"type": "text_markdown", "config": {"content": "hi"}}
+    validation.inject_project_id(widget, "proj-123")
+    assert "projectId" not in widget["config"]
+
+
+def test_inject_project_id__no_project_id_raises():
+    widget = {"type": "project_stats_card", "config": {}}
+    with pytest.raises(exceptions.DashboardValidationError, match="project-scoped"):
+        validation.inject_project_id(widget, None)
+
+
+def test_inject_project_id__no_project_id_for_non_project_widget_ok():
+    widget = {"type": "text_markdown", "config": {}}
+    validation.inject_project_id(widget, None)

sdks/python/tests/unit/api_objects/dataset/test_converters.py

@@ -0,0 +1,432 @@
+import pandas as pd
+import pandas.testing
+import json
+import tempfile
+import os
+
+from opik.api_objects.dataset import converters
+from opik.api_objects.dataset.dataset_item import DatasetItem
+from ....testlib import ANY_BUT_NONE
+
+
+def test_from_pandas__all_columns_from_dataframe_represent_all_dataset_item_fields():
+    data_for_dataframe = {
+        "id": ["id-1", "id-2"],
+        "input": [{"input-key-1": "input-1"}, {"input-key-2": "input-2"}],
+        "expected_output": [
+            {"expected-output-key-1": "expected-output-1"},
+            {"expected-output-key-2": "expected-output-2"},
+        ],
+        "metadata": [{"metadata-key-1": "v1"}, {"metadata-key-2": "v2"}],
+        "span_id": ["span-id-1", "span-id-2"],
+        "trace_id": ["trace-id-1", "trace-id-2"],
+        "source": ["some-source-1", "some-source-2"],
+    }
+
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id="id-1",
+            input={"input-key-1": "input-1"},
+            expected_output={"expected-output-key-1": "expected-output-1"},
+            metadata={"metadata-key-1": "v1"},
+            span_id="span-id-1",
+            trace_id="trace-id-1",
+            source="some-source-1",
+        ),
+        DatasetItem(
+            id="id-2",
+            input={"input-key-2": "input-2"},
+            expected_output={"expected-output-key-2": "expected-output-2"},
+            metadata={"metadata-key-2": "v2"},
+            span_id="span-id-2",
+            trace_id="trace-id-2",
+            source="some-source-2",
+        ),
+    ]
+
+    dataframe = pd.DataFrame(data_for_dataframe)
+
+    actual_items = converters.from_pandas(
+        dataframe=dataframe, keys_mapping={}, ignore_keys=[]
+    )
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_from_pandas__only_input_presented_in_dataframe__items_are_constructed_with_default_values_for_missing_fields():
+    data_for_dataframe = {
+        "input": [{"input-key-1": "input-1"}, {"input-key-2": "input-2"}],
+    }
+
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-1": "input-1"},
+            source="sdk",
+        ),
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-2": "input-2"},
+            source="sdk",
+        ),
+    ]
+
+    dataframe = pd.DataFrame(data_for_dataframe)
+
+    actual_items = converters.from_pandas(
+        dataframe=dataframe, keys_mapping={}, ignore_keys=[]
+    )
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_from_pandas__dataframe_column_does_not_have_the_same_name_as_dataset_item_field__keys_mapping_is_used():
+    data_for_dataframe = {
+        "Input column name": [{"input-key-1": "input-1"}, {"input-key-2": "input-2"}],
+    }
+
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-1": "input-1"},
+            source="sdk",
+        ),
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-2": "input-2"},
+            source="sdk",
+        ),
+    ]
+
+    dataframe = pd.DataFrame(data_for_dataframe)
+
+    actual_items = converters.from_pandas(
+        dataframe=dataframe, keys_mapping={"Input column name": "input"}, ignore_keys=[]
+    )
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_from_pandas__dataframe_contains_extra_column_not_needed_for_dataset_item__ignore_keys_is_used():
+    data_for_dataframe = {
+        "input": [{"input-key-1": "input-1"}, {"input-key-2": "input-2"}],
+        "some-extra-column": [1, 2],
+    }
+
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-1": "input-1"},
+            source="sdk",
+        ),
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-2": "input-2"},
+            source="sdk",
+        ),
+    ]
+
+    dataframe = pd.DataFrame(data_for_dataframe)
+
+    actual_items = converters.from_pandas(
+        dataframe=dataframe, keys_mapping={}, ignore_keys=["some-extra-column"]
+    )
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_to_pandas__with_keys_mapping__span_id_trace_id_and_source_ignored__happyflow():
+    EXPECTED_DATAFRAME = pd.DataFrame(
+        {
+            "id": ["id-1", "id-2"],
+            "input": [{"input-key-1": "input-1"}, {"input-key-2": "input-2"}],
+            "Customized expected output": [
+                {"expected-output-key-1": "expected-output-1"},
+                {"expected-output-key-2": "expected-output-2"},
+            ],
+            "metadata": [{"metadata-key-1": "v1"}, {"metadata-key-2": "v2"}],
+        }
+    )
+
+    input_items = [
+        DatasetItem(
+            id="id-1",
+            input={"input-key-1": "input-1"},
+            expected_output={"expected-output-key-1": "expected-output-1"},
+            metadata={"metadata-key-1": "v1"},
+            span_id="span-id-1",
+            trace_id="trace-id-1",
+            source="some-source-1",
+        ),
+        DatasetItem(
+            id="id-2",
+            input={"input-key-2": "input-2"},
+            expected_output={"expected-output-key-2": "expected-output-2"},
+            metadata={"metadata-key-2": "v2"},
+            span_id="span-id-2",
+            trace_id="trace-id-2",
+            source="some-source-2",
+        ),
+    ]
+
+    actual_dataframe = converters.to_pandas(
+        input_items, keys_mapping={"expected_output": "Customized expected output"}
+    )
+
+    # check_like ignores columns and rows order
+    pandas.testing.assert_frame_equal(
+        actual_dataframe, EXPECTED_DATAFRAME, check_like=True
+    )
+
+
+def test_from_json__all_columns_from_dataframe_represent_all_dataset_item_fields():
+    input_json = """
+    [
+        {
+            "id": "id-1",
+            "input": {"input-key-1": "input-1"},
+            "expected_output": {"expected-output-key-1": "expected-output-1"},
+            "metadata": {"metadata-key-1": "v1"}
+        },
+        {
+            "id": "id-2",
+            "input": {"input-key-2": "input-2"},
+            "expected_output": {"expected-output-key-2": "expected-output-2"},
+            "metadata": {"metadata-key-2": "v2"}
+        }
+    ]
+"""
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id="id-1",
+            input={"input-key-1": "input-1"},
+            expected_output={"expected-output-key-1": "expected-output-1"},
+            metadata={"metadata-key-1": "v1"},
+        ),
+        DatasetItem(
+            id="id-2",
+            input={"input-key-2": "input-2"},
+            expected_output={"expected-output-key-2": "expected-output-2"},
+            metadata={"metadata-key-2": "v2"},
+        ),
+    ]
+
+    actual_items = converters.from_json(input_json, keys_mapping={}, ignore_keys=[])
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_from_json__only_input_presented_in_json__items_are_constructed_with_default_values_for_missing_fields():
+    input_json = """
+    [
+        {
+            "input": {"input-key-1": "input-1"}
+        },
+        {
+            "input": {"input-key-2": "input-2"}
+        }
+    ]
+    """
+
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-1": "input-1"},
+            source="sdk",
+        ),
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-2": "input-2"},
+            source="sdk",
+        ),
+    ]
+
+    actual_items = converters.from_json(input_json, keys_mapping={}, ignore_keys=[])
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_from_json__json_objects_contain_extra_key_not_needed_for_dataset_item__ignore_keys_is_used():
+    input_json = """
+    [
+        {
+            "input": {"input-key-1": "input-1"},
+            "extra_key": 42
+        },
+        {
+            "input": {"input-key-2": "input-2"},
+            "extra_key": 4242
+        }
+    ]
+    """
+
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-1": "input-1"},
+            source="sdk",
+        ),
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-2": "input-2"},
+            source="sdk",
+        ),
+    ]
+
+    actual_items = converters.from_json(
+        input_json, keys_mapping={}, ignore_keys=["extra_key"]
+    )
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_from_json__json_objects_dont_have_the_same_name_as_dataset_item_field__keys_mapping_is_used():
+    input_json = """
+    [
+        {
+            "JSON input key": {"input-key-1": "input-1"}
+        },
+        {
+            "JSON input key": {"input-key-2": "input-2"}
+        }
+    ]
+    """
+    EXPECTED_ITEMS = [
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-1": "input-1"},
+            source="sdk",
+        ),
+        DatasetItem(
+            id=ANY_BUT_NONE,
+            input={"input-key-2": "input-2"},
+            source="sdk",
+        ),
+    ]
+
+    actual_items = converters.from_json(
+        input_json, keys_mapping={"JSON input key": "input"}, ignore_keys=[]
+    )
+
+    assert actual_items == EXPECTED_ITEMS
+
+
+def test_to_json__with_keys_mapping__span_id_trace_id_and_source_ignored__happyflow():
+    EXPECTED_JSON = """
+    [
+        {
+            "id": "id-1",
+            "input": {"input-key-1": "input-1"},
+            "Customized expected output": {"expected-output-key-1": "expected-output-1"},
+            "metadata": {"metadata-key-1": "v1"}
+        },
+        {
+            "id": "id-2",
+            "input": {"input-key-2": "input-2"},
+            "Customized expected output": {"expected-output-key-2": "expected-output-2"},
+            "metadata": {"metadata-key-2": "v2"}
+        }
+    ]
+    """
+    input_items = [
+        DatasetItem(
+            id="id-1",
+            input={"input-key-1": "input-1"},
+            expected_output={"expected-output-key-1": "expected-output-1"},
+            metadata={"metadata-key-1": "v1"},
+            span_id="span-id-1",
+            trace_id="trace-id-1",
+            source="some-source-1",
+        ),
+        DatasetItem(
+            id="id-2",
+            input={"input-key-2": "input-2"},
+            expected_output={"expected-output-key-2": "expected-output-2"},
+            metadata={"metadata-key-2": "v2"},
+            span_id="span-id-2",
+            trace_id="trace-id-2",
+            source="some-source-2",
+        ),
+    ]
+
+    actual_json = converters.to_json(
+        input_items, keys_mapping={"expected_output": "Customized expected output"}
+    )
+
+    assert json.loads(actual_json) == json.loads(EXPECTED_JSON)
+
+
+def test_from_jsonl_file__happyflow():
+    jsonl_content = """
+{"input": {"user_question": "What is the capital of France?"}, "expected_output": {"assistant_answer": "The capital of France is Paris."}}
+{"input": {"user_question": "How many planets are in our solar system?"}, "expected_output": {"assistant_answer": "There are 8 planets in our solar system: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune."}}
+"""
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
+        temp_file.write(jsonl_content)
+        temp_file_path = temp_file.name
+
+    try:
+        result = converters.from_jsonl_file(
+            temp_file_path, keys_mapping={}, ignore_keys=[]
+        )
+
+        assert result[0].input == {"user_question": "What is the capital of France?"}
+        assert result[0].expected_output == {
+            "assistant_answer": "The capital of France is Paris."
+        }
+
+        assert result[1].input == {
+            "user_question": "How many planets are in our solar system?"
+        }
+        assert result[1].expected_output == {
+            "assistant_answer": "There are 8 planets in our solar system: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune."
+        }
+    finally:
+        os.unlink(temp_file_path)
+
+
+def test_from_jsonl_file__empty_file():
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
+        temp_file_path = temp_file.name
+
+    try:
+        result = converters.from_jsonl_file(
+            temp_file_path, keys_mapping={}, ignore_keys=[]
+        )
+        assert isinstance(result, list)
+        assert len(result) == 0
+    finally:
+        os.unlink(temp_file_path)
+
+
+def test_from_jsonl_file__file_with_empty_lines():
+    jsonl_content = """
+{"input": {"user_question": "What is the capital of France?"}, "expected_output": {"assistant_answer": "The capital of France is Paris."}}
+
+{"input": {"user_question": "How many planets are in our solar system?"}, "expected_output": {"assistant_answer": "There are 8 planets in our solar system: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune."}}
+
+"""
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as temp_file:
+        temp_file.write(jsonl_content)
+        temp_file_path = temp_file.name
+
+    try:
+        result = converters.from_jsonl_file(
+            temp_file_path, keys_mapping={}, ignore_keys=[]
+        )
+
+        assert len(result) == 2
+
+        assert result[0].input == {"user_question": "What is the capital of France?"}
+        assert result[0].expected_output == {
+            "assistant_answer": "The capital of France is Paris."
+        }
+
+        assert result[1].input == {
+            "user_question": "How many planets are in our solar system?"
+        }
+        assert result[1].expected_output == {
+            "assistant_answer": "There are 8 planets in our solar system: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune."
+        }
+    finally:
+        os.unlink(temp_file_path)

sdks/python/tests/unit/api_objects/dataset/test_dataset_client.py

@@ -0,0 +1,165 @@
+from unittest.mock import Mock
+from opik.api_objects.dataset.dataset import Dataset
+
+
+def test_insert_deduplication__two_dicts_passed_with_the_same_content__only_one_is_inserted():
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+    )
+
+    item_dict = {
+        "input": {"key": "value", "key2": "value2"},
+        "expected_output": {"key": "value", "key2": "value2"},
+        "metadata": {"key": "value", "key2": "value2"},
+    }
+
+    # Insert the identical items
+    dataset.insert([item_dict, item_dict])
+
+    assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 1, (
+        "create_or_update_dataset_items should be called only once"
+    )
+
+    call_args = mock_rest_client.datasets.create_or_update_dataset_items.call_args
+    inserted_items = call_args[1]["items"]
+
+    assert len(inserted_items) == 1, "Only one item should be inserted"
+
+
+def test_insert_deduplication__two_dicts_passed_with_the_different_content__both_are_inserted():
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+    )
+
+    item_dict1 = {
+        "input": {"key": "value1"},
+        "expected_output": {"key": "output1"},
+        "metadata": {"key": "meta1"},
+    }
+    item_dict2 = {
+        "input": {"key": "value2"},
+        "expected_output": {"key": "output2"},
+        "metadata": {"key": "meta2"},
+    }
+
+    # Insert the different items
+    dataset.insert([item_dict1, item_dict2])
+
+    assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 1, (
+        "create_or_update_dataset_items should be called only once"
+    )
+
+    call_args = mock_rest_client.datasets.create_or_update_dataset_items.call_args
+    inserted_items = call_args[1]["items"]
+
+    assert len(inserted_items) == 2, "Two items should be inserted"
+
+
+def test_insert_deduplication__three_dicts_passed__one_unique__two_duplicates__two_different_items_are_inserted():
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+    )
+
+    item_dict1 = {
+        "input": {"key": "value1"},
+        "expected_output": {"key": "output1"},
+        "metadata": {"key": "meta1"},
+    }
+    item_dict2 = {
+        "input": {"key": "value2"},
+        "expected_output": {"key": "output2"},
+        "metadata": {"key": "meta2"},
+    }
+
+    # Insert 3 items: one unique and two duplicates
+    dataset.insert([item_dict1, item_dict2, item_dict1])
+
+    assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 1, (
+        "create_or_update_dataset_items should be called only once"
+    )
+
+    call_args = mock_rest_client.datasets.create_or_update_dataset_items.call_args
+    inserted_rest_items = call_args[1]["items"]
+
+    assert len(inserted_rest_items) == 2, "Two items should be inserted"
+
+
+def test_update__happyflow():
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+    )
+
+    initial_item = {
+        "input": {"key": "initial_value"},
+        "expected_output": {"key": "initial_output"},
+        "metadata": {"key": "initial_metadata"},
+    }
+
+    dataset.insert([initial_item])
+
+    assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 1, (
+        "create_or_update_dataset_items should be called once for insertion"
+    )
+
+    insert_call_args = (
+        mock_rest_client.datasets.create_or_update_dataset_items.call_args
+    )
+    inserted_items = insert_call_args[1]["items"]
+
+    assert len(inserted_items) == 1, "One item should be inserted"
+
+    # Create an updated version of the item
+    updated_item = {
+        "id": inserted_items[0].id,
+        "input": {"key": "updated_value"},
+        "expected_output": {"key": "updated_output"},
+        "metadata": {"key": "updated_metadata"},
+    }
+
+    # Update the item
+    dataset.update([updated_item])
+
+    # Check that create_or_update_dataset_items was called twice in total (once for insertion, once for update)
+    assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 2, (
+        "create_or_update_dataset_items should be called twice in total"
+    )
+
+    # Get the arguments passed to create_or_update_dataset_items for update
+    update_call_args = (
+        mock_rest_client.datasets.create_or_update_dataset_items.call_args
+    )
+    updated_rest_items = update_call_args[1]["items"]
+
+    # Check that one item was updated
+    assert len(updated_rest_items) == 1, "One item should be updated"
+
+    # Verify the content of the updated item
+    assert updated_rest_items[0].data["input"] == {"key": "updated_value"}, (
+        "Input should be updated"
+    )
+    assert updated_rest_items[0].data["expected_output"] == {"key": "updated_output"}, (
+        "Expected output should be updated"
+    )
+    assert updated_rest_items[0].data["metadata"] == {"key": "updated_metadata"}, (
+        "Metadata should be updated"
+    )

sdks/python/tests/unit/api_objects/dataset/test_dataset_items_count.py

@@ -0,0 +1,178 @@
+from unittest.mock import Mock
+from opik.api_objects.dataset.dataset import Dataset
+from opik.rest_api.types.dataset_public import DatasetPublic
+
+
+def test_dataset_items_count__cached_value__returns_cached_count():
+    """Test that dataset_items_count returns cached value when available."""
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=5,
+    )
+
+    count = dataset.dataset_items_count
+
+    assert count == 5
+    mock_rest_client.datasets.get_dataset_by_id.assert_not_called()
+
+
+def test_dataset_items_count__no_cached_value__fetches_from_backend():
+    """Test that dataset_items_count fetches from backend when cache is None."""
+    mock_rest_client = Mock()
+
+    mock_dataset_public = DatasetPublic(name="test_dataset", dataset_items_count=10)
+    mock_rest_client.datasets.get_dataset_by_id.return_value = mock_dataset_public
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=None,
+    )
+
+    count = dataset.dataset_items_count
+
+    assert count == 10
+    mock_rest_client.datasets.get_dataset_by_id.assert_called_once_with(id=dataset.id)
+
+
+def test_dataset_items_count__fetched_once__cached_for_subsequent_calls():
+    """Test that dataset_items_count is fetched once and then cached."""
+    mock_rest_client = Mock()
+
+    mock_dataset_public = DatasetPublic(name="test_dataset", dataset_items_count=10)
+    mock_rest_client.datasets.get_dataset_by_id.return_value = mock_dataset_public
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=None,
+    )
+
+    count1 = dataset.dataset_items_count
+    count2 = dataset.dataset_items_count
+    count3 = dataset.dataset_items_count
+
+    assert count1 == 10
+    assert count2 == 10
+    assert count3 == 10
+    mock_rest_client.datasets.get_dataset_by_id.assert_called_once()
+
+
+def test_delete__invalidates_cached_count():
+    """Test that delete() invalidates the cached count."""
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=5,
+    )
+
+    assert dataset.dataset_items_count == 5
+
+    dataset.delete(["item1", "item2"])
+
+    mock_dataset_public = DatasetPublic(name="test_dataset", dataset_items_count=3)
+    mock_rest_client.datasets.get_dataset_by_id.return_value = mock_dataset_public
+
+    count = dataset.dataset_items_count
+
+    assert count == 3
+    mock_rest_client.datasets.get_dataset_by_id.assert_called_once()
+
+
+def test_update__invalidates_cached_count():
+    """Test that update() invalidates the cached count."""
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=5,
+    )
+
+    assert dataset.dataset_items_count == 5
+
+    updated_item = {
+        "id": "item1",
+        "input": {"key": "updated_value"},
+        "expected_output": {"key": "updated_output"},
+    }
+    dataset.update([updated_item])
+
+    mock_dataset_public = DatasetPublic(name="test_dataset", dataset_items_count=5)
+    mock_rest_client.datasets.get_dataset_by_id.return_value = mock_dataset_public
+
+    count = dataset.dataset_items_count
+
+    assert count == 5
+    mock_rest_client.datasets.get_dataset_by_id.assert_called_once()
+
+
+def test_insert__invalidates_cached_count():
+    """Test that insert() invalidates the cached count."""
+    mock_rest_client = Mock()
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=5,
+    )
+
+    assert dataset.dataset_items_count == 5
+
+    new_items = [
+        {
+            "input": {"key": "value1"},
+            "expected_output": {"key": "output1"},
+        },
+        {
+            "input": {"key": "value2"},
+            "expected_output": {"key": "output2"},
+        },
+    ]
+    dataset.insert(new_items)
+
+    mock_dataset_public = DatasetPublic(name="test_dataset", dataset_items_count=7)
+    mock_rest_client.datasets.get_dataset_by_id.return_value = mock_dataset_public
+
+    count = dataset.dataset_items_count
+
+    assert count == 7
+    mock_rest_client.datasets.get_dataset_by_id.assert_called_once()
+
+
+def test_backend_returns_none_count__property_returns_none():
+    """Test that if backend returns None for count, property returns None."""
+    mock_rest_client = Mock()
+
+    mock_dataset_public = DatasetPublic(name="test_dataset", dataset_items_count=None)
+    mock_rest_client.datasets.get_dataset_by_id.return_value = mock_dataset_public
+
+    dataset = Dataset(
+        name="test_dataset",
+        description="Test description",
+        project_name="Test project",
+        rest_client=mock_rest_client,
+        dataset_items_count=None,
+    )
+
+    count = dataset.dataset_items_count
+
+    assert count is None
+    mock_rest_client.datasets.get_dataset_by_id.assert_called_once()

sdks/python/tests/unit/api_objects/dataset/test_rate_limit_retry.py

@@ -0,0 +1,122 @@
+import unittest
+from unittest.mock import Mock, patch
+
+from opik.api_objects.dataset import dataset
+from opik.rest_api.core.api_error import ApiError
+
+
+class TestDatasetRateLimitRetry(unittest.TestCase):
+    """Test rate limit retry behavior for dataset operations using public API."""
+
+    def _create_dataset_with_mock_client(self) -> tuple[dataset.Dataset, Mock]:
+        """Create a Dataset instance with a mocked REST client."""
+        mock_rest_client = Mock()
+        dataset_obj = dataset.Dataset(
+            name="test_dataset",
+            description="test",
+            project_name="Test project",
+            rest_client=mock_rest_client,
+        )
+        return dataset_obj, mock_rest_client
+
+    @patch("opik.api_objects.rest_helpers._sleep")
+    def test_insert__429_with_retry_after_header__retries_with_correct_delay(
+        self, mock_sleep: Mock
+    ) -> None:
+        """Test that 429 errors with RateLimit-Reset header are retried with correct delay."""
+        dataset_obj, mock_rest_client = self._create_dataset_with_mock_client()
+
+        # First call raises 429 with rate limit headers, second call succeeds
+        rate_limit_error = ApiError(
+            status_code=429,
+            headers={
+                "RateLimit-Reset": "5",  # 5 seconds retry after
+            },
+            body="Rate limit exceeded",
+        )
+        mock_rest_client.datasets.create_or_update_dataset_items.side_effect = [
+            rate_limit_error,
+            None,  # Success on second attempt
+        ]
+
+        # Execute using public API
+        dataset_obj.insert([{"input": "test"}])
+
+        # Verify retry behavior
+        assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 2
+        mock_sleep.assert_called_once_with(5.0)
+
+    @patch("opik.api_objects.rest_helpers._sleep")
+    def test_insert__429_without_header__uses_fallback_delay(
+        self, mock_sleep: Mock
+    ) -> None:
+        """Test that 429 errors without headers use fallback 1 second delay."""
+        dataset_obj, mock_rest_client = self._create_dataset_with_mock_client()
+
+        # First two calls raise 429 without headers, third succeeds
+        rate_limit_error = ApiError(
+            status_code=429,
+            headers={},
+            body="Rate limit exceeded",
+        )
+        mock_rest_client.datasets.create_or_update_dataset_items.side_effect = [
+            rate_limit_error,
+            rate_limit_error,
+            None,  # Success on third attempt
+        ]
+
+        # Execute using public API
+        dataset_obj.insert([{"input": "test"}])
+
+        # Verify fallback delay: always 1 second when no header
+        assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 3
+        assert mock_sleep.call_count == 2
+        # Both retries should use 1 second delay
+        assert all(call[0][0] == 1 for call in mock_sleep.call_args_list)
+
+    def test_insert__non_429_error__raises_immediately(self) -> None:
+        """Test that non-429 errors are raised immediately without retry."""
+        dataset_obj, mock_rest_client = self._create_dataset_with_mock_client()
+
+        # Simulate a 500 error
+        error = ApiError(
+            status_code=500,
+            headers={},
+            body="Internal server error",
+        )
+        mock_rest_client.datasets.create_or_update_dataset_items.side_effect = error
+
+        # Execute & Verify
+        with self.assertRaises(ApiError) as context:
+            dataset_obj.insert([{"input": "test"}])
+
+        assert context.exception.status_code == 500
+        # Should only try once for non-429 errors
+        assert mock_rest_client.datasets.create_or_update_dataset_items.call_count == 1
+
+    @patch("opik.api_objects.rest_helpers._sleep")
+    def test_delete__429_with_retry_after_header__retries_with_correct_delay(
+        self, mock_sleep: Mock
+    ) -> None:
+        """Test that delete operation also handles 429 errors correctly."""
+        dataset_obj, mock_rest_client = self._create_dataset_with_mock_client()
+
+        # First call raises 429 with rate limit headers, second call succeeds
+        rate_limit_error = ApiError(
+            status_code=429,
+            headers={
+                "RateLimit-Reset": "3",
+            },
+            body="Rate limit exceeded",
+        )
+        mock_rest_client.datasets.delete_dataset_items.side_effect = [
+            rate_limit_error,
+            None,  # Success on second attempt
+        ]
+
+        # Execute using public API
+        dataset_obj.delete(["item-id-1"])
+
+        # Verify retry behavior
+        assert mock_rest_client.datasets.delete_dataset_items.call_count == 2
+        mock_sleep.assert_called_once_with(3.0)

sdks/python/tests/unit/api_objects/dataset/test_stream_dataset_items.py

@@ -0,0 +1,74 @@
+"""Unit tests for stream_dataset_items() in rest_operations."""
+
+from unittest.mock import Mock, patch
+
+from opik.api_objects.dataset import rest_operations
+from opik.rest_api.types import dataset_item as rest_dataset_item
+
+_SHADOW_WARNING = (
+    "Dataset item data contains keys that shadow DatasetItem fields and will be ignored: %s. "
+    "Rename these keys in your dataset to preserve them."
+)
+
+
+def _make_rest_item(item_id: str, data: dict) -> rest_dataset_item.DatasetItem:
+    return rest_dataset_item.DatasetItem(
+        id=item_id,
+        source="sdk",
+        data=data,
+    )
+
+
+def test_stream_dataset_items__colliding_id_key__uses_real_id_and_warns():
+    real_id = "real-uuid-1234"
+    rest_item = _make_rest_item(real_id, {"id": "COLLISION", "question": "What?"})
+
+    mock_rest_client = Mock()
+
+    with (
+        patch(
+            "opik.api_objects.dataset.rest_operations.rest_stream_parser.read_and_parse_stream",
+            side_effect=[[rest_item], []],
+        ),
+        patch.object(rest_operations.LOGGER, "warning") as mock_warn,
+    ):
+        items = list(
+            rest_operations.stream_dataset_items(
+                rest_client=mock_rest_client,
+                dataset_name="test-dataset",
+                project_name=None,
+            )
+        )
+
+    assert len(items) == 1
+    assert items[0].id == real_id
+
+    mock_warn.assert_called_once_with(_SHADOW_WARNING, ["id"])
+
+
+def test_stream_dataset_items__colliding_id_key__warning_emitted_only_once():
+    """Warning is logged once per stream even when multiple items have the collision."""
+    items_data = [
+        _make_rest_item(f"uuid-{i}", {"id": f"hotpot-{i}", "question": "Q?"})
+        for i in range(3)
+    ]
+
+    mock_rest_client = Mock()
+
+    with (
+        patch(
+            "opik.api_objects.dataset.rest_operations.rest_stream_parser.read_and_parse_stream",
+            side_effect=[items_data, []],
+        ),
+        patch.object(rest_operations.LOGGER, "warning") as mock_warn,
+    ):
+        result = list(
+            rest_operations.stream_dataset_items(
+                rest_client=mock_rest_client,
+                dataset_name="test-dataset",
+                project_name=None,
+            )
+        )
+
+    assert len(result) == 3
+    mock_warn.assert_called_once_with(_SHADOW_WARNING, ["id"])

sdks/python/tests/unit/api_objects/dataset/test_suite/test_converters.py

@@ -0,0 +1,513 @@
+"""Unit tests for test_suite.converters module."""
+
+import json
+import os
+import tempfile
+
+import pandas as pd
+import pandas.testing
+import pytest
+from unittest import mock
+
+from opik.api_objects.dataset import dataset_item
+from opik.api_objects.dataset.test_suite import converters
+
+
+# ---------------------------------------------------------------------------
+# evaluators_to_assertions
+# ---------------------------------------------------------------------------
+
+
+def test_evaluators_to_assertions__multiple_evaluators__concatenates():
+    e1 = mock.MagicMock()
+    e1.assertions = ["A1"]
+    e2 = mock.MagicMock()
+    e2.assertions = ["A2", "A3"]
+
+    assert converters.evaluators_to_assertions([e1, e2]) == ["A1", "A2", "A3"]
+
+
+def test_evaluators_to_assertions__empty_list__returns_empty():
+    assert converters.evaluators_to_assertions([]) == []
+
+
+# ---------------------------------------------------------------------------
+# version_evaluators_to_assertions
+# ---------------------------------------------------------------------------
+
+
+def test_version_evaluators_to_assertions__llm_judge__extracts_assertions():
+    from opik.evaluation.suite_evaluators import LLMJudge
+
+    judge = LLMJudge(assertions=["Response is accurate"], track=False)
+    config = judge.to_config().model_dump(by_alias=True)
+
+    evaluator_item = mock.MagicMock()
+    evaluator_item.type = "llm_judge"
+    evaluator_item.config = config
+
+    assert converters.version_evaluators_to_assertions([evaluator_item]) == [
+        "Response is accurate"
+    ]
+
+
+def test_version_evaluators_to_assertions__none__returns_empty():
+    assert converters.version_evaluators_to_assertions(None) == []
+
+
+def test_version_evaluators_to_assertions__non_llm_judge__skipped():
+    evaluator_item = mock.MagicMock()
+    evaluator_item.type = "custom_scorer"
+
+    assert converters.version_evaluators_to_assertions([evaluator_item]) == []
+
+
+# ---------------------------------------------------------------------------
+# version_policy_to_execution_policy
+# ---------------------------------------------------------------------------
+
+
+def test_version_policy_to_execution_policy__converts():
+    policy = mock.MagicMock()
+    policy.runs_per_item = 5
+    policy.pass_threshold = 3
+
+    assert converters.version_policy_to_execution_policy(policy) == {
+        "runs_per_item": 5,
+        "pass_threshold": 3,
+    }
+
+
+def test_version_policy_to_execution_policy__none__returns_default():
+    assert converters.version_policy_to_execution_policy(None) == {
+        "runs_per_item": 1,
+        "pass_threshold": 1,
+    }
+
+
+# ---------------------------------------------------------------------------
+# dataset_item_to_suite_item_dict (DatasetItem → TestSuiteItem)
+# ---------------------------------------------------------------------------
+
+
+def test_dataset_item_to_suite_item_dict__all_fields():
+    item = dataset_item.DatasetItem(
+        id="item-1",
+        description="Test item",
+        question="What is 2+2?",
+        execution_policy=dataset_item.ExecutionPolicyItem(
+            runs_per_item=3, pass_threshold=2
+        ),
+    )
+
+    assert converters.dataset_item_to_suite_item_dict(item) == {
+        "id": "item-1",
+        "data": {"question": "What is 2+2?"},
+        "assertions": [],
+        "description": "Test item",
+        "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+    }
+
+
+def test_dataset_item_to_suite_item_dict__minimal():
+    item = dataset_item.DatasetItem(id="item-1", question="Hello")
+
+    result = converters.dataset_item_to_suite_item_dict(item)
+
+    assert result == {"id": "item-1", "data": {"question": "Hello"}, "assertions": []}
+    assert "description" not in result
+    assert "execution_policy" not in result
+
+
+def test_dataset_item_to_suite_item_dict__with_evaluators__extracts_assertions():
+    from opik.evaluation.suite_evaluators import LLMJudge
+
+    judge = LLMJudge(assertions=["Is correct"], track=False)
+    config = judge.to_config().model_dump(by_alias=True)
+
+    item = dataset_item.DatasetItem(
+        id="item-2",
+        evaluators=[
+            dataset_item.EvaluatorItem(
+                name="llm_judge", type="llm_judge", config=config
+            ),
+        ],
+        question="Hello",
+    )
+
+    assert converters.dataset_item_to_suite_item_dict(item)["assertions"] == [
+        "Is correct"
+    ]
+
+
+# ---------------------------------------------------------------------------
+# suite_item_dict_to_dataset_item (TestSuiteItem → DatasetItem)
+# ---------------------------------------------------------------------------
+
+
+def test_suite_item_dict_to_dataset_item__all_fields():
+    item = {
+        "id": "item-1",
+        "data": {"question": "Hello", "context": "test"},
+        "assertions": ["Is polite"],
+        "description": "A test case",
+        "execution_policy": {"runs_per_item": 5, "pass_threshold": 3},
+    }
+
+    ds_item = converters.suite_item_dict_to_dataset_item(item)
+
+    assert ds_item.id == "item-1"
+    assert ds_item.get_content() == {"question": "Hello", "context": "test"}
+    assert ds_item.description == "A test case"
+    assert ds_item.execution_policy is not None
+    assert ds_item.execution_policy.runs_per_item == 5
+    assert ds_item.execution_policy.pass_threshold == 3
+    assert ds_item.evaluators is not None
+    assert len(ds_item.evaluators) == 1
+    assert ds_item.evaluators[0].type == "llm_judge"
+
+
+def test_suite_item_dict_to_dataset_item__minimal__generates_id():
+    ds_item = converters.suite_item_dict_to_dataset_item(
+        {"data": {"question": "Hello"}}
+    )
+
+    assert ds_item.id is not None
+    assert len(ds_item.id) > 0
+    assert ds_item.get_content() == {"question": "Hello"}
+    assert ds_item.evaluators is None
+    assert ds_item.execution_policy is None
+
+
+# ---------------------------------------------------------------------------
+# to_json / from_json
+# ---------------------------------------------------------------------------
+
+
+SAMPLE_ITEMS = [
+    {
+        "id": "item-1",
+        "data": {"question": "How do I get a refund?", "context": "Premium user"},
+        "assertions": ["Response is polite"],
+        "description": "Refund scenario",
+        "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+    },
+    {
+        "id": "item-2",
+        "data": {"question": "Is my account hacked?"},
+        "assertions": [],
+    },
+]
+
+
+def test_to_json__happyflow():
+    EXPECTED = [
+        {
+            "id": "item-1",
+            "data": {"question": "How do I get a refund?", "context": "Premium user"},
+            "assertions": ["Response is polite"],
+            "description": "Refund scenario",
+            "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+        },
+        {
+            "id": "item-2",
+            "data": {"question": "Is my account hacked?"},
+            "assertions": [],
+        },
+    ]
+
+    assert json.loads(converters.to_json(SAMPLE_ITEMS)) == EXPECTED
+
+
+def test_to_json__empty_list():
+    assert json.loads(converters.to_json([])) == []
+
+
+def test_from_json__happyflow():
+    json_str = json.dumps(
+        [
+            {"data": {"question": "Hello"}, "assertions": ["Is polite"]},
+            {"data": {"question": "Bye"}},
+        ]
+    )
+
+    EXPECTED = [
+        {"data": {"question": "Hello"}, "assertions": ["Is polite"]},
+        {"data": {"question": "Bye"}},
+    ]
+
+    assert converters.from_json(json_str, {}, []) == EXPECTED
+
+
+def test_from_json__with_keys_mapping():
+    json_str = json.dumps(
+        [
+            {"test_data": {"question": "Hello"}, "checks": ["Is polite"]},
+        ]
+    )
+
+    EXPECTED = [{"data": {"question": "Hello"}, "assertions": ["Is polite"]}]
+
+    assert (
+        converters.from_json(
+            json_str, {"test_data": "data", "checks": "assertions"}, []
+        )
+        == EXPECTED
+    )
+
+
+def test_from_json__with_ignore_keys():
+    json_str = json.dumps(
+        [
+            {"data": {"question": "Hello"}, "internal_note": "skip this"},
+        ]
+    )
+
+    EXPECTED = [{"data": {"question": "Hello"}}]
+
+    assert converters.from_json(json_str, {}, ["internal_note"]) == EXPECTED
+
+
+def test_from_json__non_array__raises_value_error():
+    json_str = json.dumps({"data": {"question": "Hello"}})
+
+    with pytest.raises(ValueError, match="must be an array"):
+        converters.from_json(json_str, {}, [])
+
+
+# ---------------------------------------------------------------------------
+# to_pandas / from_pandas
+# ---------------------------------------------------------------------------
+
+
+def test_to_pandas__happyflow():
+    EXPECTED = pd.DataFrame(
+        [
+            {
+                "id": "item-1",
+                "data": {
+                    "question": "How do I get a refund?",
+                    "context": "Premium user",
+                },
+                "assertions": ["Response is polite"],
+                "description": "Refund scenario",
+                "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+            },
+            {
+                "id": "item-2",
+                "data": {"question": "Is my account hacked?"},
+                "assertions": [],
+            },
+        ]
+    )
+
+    pandas.testing.assert_frame_equal(converters.to_pandas(SAMPLE_ITEMS), EXPECTED)
+
+
+def test_to_pandas__empty_list():
+    assert len(converters.to_pandas([])) == 0
+
+
+def test_from_pandas__happyflow():
+    dataframe = pd.DataFrame(
+        [
+            {"data": {"question": "Hello"}, "assertions": ["Is polite"]},
+            {"data": {"question": "Bye"}},
+        ]
+    )
+
+    EXPECTED = [
+        {"data": {"question": "Hello"}, "assertions": ["Is polite"]},
+        {"data": {"question": "Bye"}},
+    ]
+
+    assert converters.from_pandas(dataframe, {}, []) == EXPECTED
+
+
+def test_from_pandas__with_keys_mapping():
+    dataframe = pd.DataFrame([{"test_data": {"question": "Hello"}}])
+
+    EXPECTED = [{"data": {"question": "Hello"}}]
+
+    assert converters.from_pandas(dataframe, {"test_data": "data"}, []) == EXPECTED
+
+
+def test_from_pandas__nan_values__skipped():
+    dataframe = pd.DataFrame(
+        [
+            {"data": {"question": "Hello"}, "assertions": ["Is polite"]},
+            {"data": {"question": "Bye"}, "assertions": float("nan")},
+        ]
+    )
+
+    result = converters.from_pandas(dataframe, {}, [])
+
+    assert result[0] == {"data": {"question": "Hello"}, "assertions": ["Is polite"]}
+    assert result[1] == {"data": {"question": "Bye"}}
+
+
+# ---------------------------------------------------------------------------
+# from_jsonl_file
+# ---------------------------------------------------------------------------
+
+
+def test_from_jsonl_file__happyflow():
+    jsonl_content = (
+        '{"data": {"question": "What is 2+2?"}, "assertions": ["Is correct"]}\n'
+        '{"data": {"question": "Capital of France?"}}\n'
+    )
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+        f.write(jsonl_content)
+        path = f.name
+
+    try:
+        EXPECTED = [
+            {"data": {"question": "What is 2+2?"}, "assertions": ["Is correct"]},
+            {"data": {"question": "Capital of France?"}},
+        ]
+
+        assert converters.from_jsonl_file(path, {}, []) == EXPECTED
+    finally:
+        os.unlink(path)
+
+
+def test_from_jsonl_file__empty_lines__skipped():
+    jsonl_content = '{"data": {"question": "Q1"}}\n\n{"data": {"question": "Q2"}}\n\n'
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+        f.write(jsonl_content)
+        path = f.name
+
+    try:
+        assert len(converters.from_jsonl_file(path, {}, [])) == 2
+    finally:
+        os.unlink(path)
+
+
+def test_from_jsonl_file__with_keys_mapping():
+    jsonl_content = '{"test_data": {"question": "Hello"}}\n'
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
+        f.write(jsonl_content)
+        path = f.name
+
+    try:
+        EXPECTED = [{"data": {"question": "Hello"}}]
+
+        assert converters.from_jsonl_file(path, {"test_data": "data"}, []) == EXPECTED
+    finally:
+        os.unlink(path)
+
+
+# ---------------------------------------------------------------------------
+# Round-trip tests
+# ---------------------------------------------------------------------------
+
+
+def test_adapter_roundtrip__suite_to_dataset_to_suite__preserves_all_fields():
+    original = {
+        "id": "item-1",
+        "data": {"question": "Hello", "context": "Premium"},
+        "assertions": ["Is polite", "Is helpful"],
+        "description": "Test case",
+        "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+    }
+
+    ds_item = converters.suite_item_dict_to_dataset_item(original)
+    recovered = converters.dataset_item_to_suite_item_dict(ds_item)
+
+    assert recovered["id"] == original["id"]
+    assert recovered["data"] == original["data"]
+    assert sorted(recovered["assertions"]) == sorted(original["assertions"])
+    assert recovered["description"] == original["description"]
+    assert recovered["execution_policy"] == original["execution_policy"]
+
+
+def test_adapter_roundtrip__two_cycles__stable():
+    original = {
+        "id": "item-1",
+        "data": {"question": "Hello"},
+        "assertions": ["Is polite"],
+        "execution_policy": {"runs_per_item": 5, "pass_threshold": 3},
+    }
+
+    suite_item = original
+    for _ in range(2):
+        ds_item = converters.suite_item_dict_to_dataset_item(suite_item)
+        suite_item = converters.dataset_item_to_suite_item_dict(ds_item)
+
+    assert suite_item["data"] == original["data"]
+    assert suite_item["assertions"] == original["assertions"]
+    assert suite_item["execution_policy"] == original["execution_policy"]
+
+
+def test_json_roundtrip__export_import_export__stable():
+    json_str_1 = converters.to_json(SAMPLE_ITEMS)
+    imported = converters.from_json(json_str_1, {}, [])
+    json_str_2 = converters.to_json(imported)
+
+    assert json.loads(json_str_1) == json.loads(json_str_2)
+
+
+def test_json_roundtrip__import_export_import__stable():
+    json_str = json.dumps(
+        [
+            {
+                "data": {"question": "Hello"},
+                "assertions": ["Is polite"],
+                "description": "Test",
+                "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+            },
+        ]
+    )
+
+    items_1 = converters.from_json(json_str, {}, [])
+    exported = converters.to_json(items_1)
+    items_2 = converters.from_json(exported, {}, [])
+
+    assert items_1 == items_2
+
+
+def test_pandas_roundtrip__export_import_export__stable():
+    df_1 = converters.to_pandas(SAMPLE_ITEMS)
+    imported = converters.from_pandas(df_1, {}, [])
+    df_2 = converters.to_pandas(imported)
+
+    pandas.testing.assert_frame_equal(df_1, df_2)
+
+
+def test_full_roundtrip__json_through_adapters__two_cycles():
+    """JSON → from_json → adapter → adapter → to_json, repeated twice."""
+    original_json = json.dumps(
+        [
+            {
+                "id": "item-1",
+                "data": {"question": "Refund?", "tier": "premium"},
+                "assertions": ["Is polite", "No hallucination"],
+                "description": "Refund scenario",
+                "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+            },
+        ]
+    )
+
+    current_json = original_json
+    for _ in range(2):
+        suite_items = converters.from_json(current_json, {}, [])
+        ds_items = [converters.suite_item_dict_to_dataset_item(i) for i in suite_items]
+        suite_items = [converters.dataset_item_to_suite_item_dict(i) for i in ds_items]
+        current_json = converters.to_json(suite_items)
+
+    EXPECTED = {
+        "id": "item-1",
+        "data": {"question": "Refund?", "tier": "premium"},
+        "assertions": ["Is polite", "No hallucination"],
+        "description": "Refund scenario",
+        "execution_policy": {"runs_per_item": 3, "pass_threshold": 2},
+    }
+
+    final = json.loads(current_json)
+    assert len(final) == 1
+    result = final[0]
+    assert result["id"] == EXPECTED["id"]
+    assert result["data"] == EXPECTED["data"]
+    assert sorted(result["assertions"]) == sorted(EXPECTED["assertions"])
+    assert result["description"] == EXPECTED["description"]
+    assert result["execution_policy"] == EXPECTED["execution_policy"]

sdks/python/tests/unit/api_objects/dataset/test_suite/test_report_file.py

@@ -0,0 +1,317 @@
+"""Unit tests for test suite result file generation."""
+
+import json
+import os
+
+
+from opik.api_objects.dataset.test_suite import (
+    suite_result_constructor,
+    types as suite_types,
+)
+from opik.api_objects.dataset.test_suite.report_processors import file_writer
+from opik.api_objects.dataset import dataset_item
+from opik.evaluation import evaluation_result, test_result, test_case
+from opik.evaluation.metrics import score_result
+
+
+def _make_test_result(
+    dataset_item_id: str,
+    trial_id: int,
+    scores: list[tuple[str, float]],
+    task_output: dict | None = None,
+    dataset_item_content: dict | None = None,
+    execution_policy: dict | None = None,
+    task_execution_time: float | None = None,
+    scoring_time: float | None = None,
+) -> test_result.TestResult:
+    ds_item = None
+    if execution_policy is not None:
+        ds_item = dataset_item.DatasetItem(
+            id=dataset_item_id,
+            execution_policy=dataset_item.ExecutionPolicyItem(
+                runs_per_item=execution_policy.get("runs_per_item"),
+                pass_threshold=execution_policy.get("pass_threshold"),
+            ),
+        )
+
+    return test_result.TestResult(
+        test_case=test_case.TestCase(
+            trace_id=f"trace-{dataset_item_id}-{trial_id}",
+            dataset_item_id=dataset_item_id,
+            task_output=task_output or {"input": "test", "output": "result"},
+            dataset_item_content=dataset_item_content or {"question": "What?"},
+            dataset_item=ds_item,
+        ),
+        score_results=[
+            score_result.ScoreResult(name=name, value=value) for name, value in scores
+        ],
+        trial_id=trial_id,
+        task_execution_time=task_execution_time,
+        scoring_time=scoring_time,
+    )
+
+
+def _make_suite_result(
+    test_results_list: list[test_result.TestResult],
+    suite_name: str | None = None,
+    total_time: float | None = None,
+) -> suite_types.TestSuiteResult:
+    eval_result = evaluation_result.EvaluationResult(
+        experiment_id="exp-123",
+        dataset_id="dataset-456",
+        experiment_name="my-experiment",
+        test_results=test_results_list,
+        experiment_url="http://example.com/experiment/exp-123",
+        trial_count=1,
+    )
+    return suite_result_constructor.build_suite_result(
+        eval_result,
+        suite_name=suite_name,
+        total_time=total_time,
+    )
+
+
+class TestToReportDict:
+    def test_to_report_dict__single_item_with_mixed_scores__returns_correct_structure(
+        self,
+    ):
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("Is polite", True), ("Is helpful", False)],
+                task_output={"input": "hi", "output": "hello"},
+                dataset_item_content={"question": "hi"},
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+                task_execution_time=1.234,
+                scoring_time=0.567,
+            )
+        ]
+        suite_result = _make_suite_result(test_results_list, suite_name="My Suite")
+
+        result = suite_result.to_report_dict()
+
+        assert result["suite_passed"] is False
+        assert result["items_passed"] == 0
+        assert result["items_total"] == 1
+        assert result["pass_rate"] == 0.0
+        assert result["experiment_id"] == "exp-123"
+        assert result["experiment_name"] == "my-experiment"
+        assert result["experiment_url"] == "http://example.com/experiment/exp-123"
+        assert result["suite_name"] == "My Suite"
+        assert "generated_at" in result
+
+        assert len(result["items"]) == 1
+        item = result["items"][0]
+        assert item["dataset_item_id"] == "item-1"
+        assert item["passed"] is False
+        assert item["runs_passed"] == 0
+        assert item["execution_policy"] == {"runs_per_item": 1, "pass_threshold": 1}
+
+        assert len(item["runs"]) == 1
+        run = item["runs"][0]
+        assert run["trial_id"] == 0
+        assert run["passed"] is False
+        assert run["input"] == "hi"
+        assert run["output"] == "hello"
+        assert run["trace_id"] == "trace-item-1-0"
+        assert run["task_execution_time_seconds"] == 1.234
+        assert run["scoring_time_seconds"] == 0.567
+
+        assert len(run["assertions"]) == 2
+        assert run["assertions"][0]["name"] == "Is polite"
+        assert run["assertions"][0]["passed"] is True
+        assert run["assertions"][1]["name"] == "Is helpful"
+        assert run["assertions"][1]["passed"] is False
+
+    def test_to_report_dict__all_items_pass__suite_passed_true(self):
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+            _make_test_result(
+                dataset_item_id="item-2",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(test_results_list)
+
+        result = suite_result.to_report_dict()
+
+        assert result["suite_passed"] is True
+        assert result["items_passed"] == 2
+        assert result["pass_rate"] == 1.0
+
+    def test_to_report_dict__with_total_time__includes_rounded_value(self):
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(test_results_list, total_time=12.3456)
+
+        result = suite_result.to_report_dict()
+
+        assert result["total_time_seconds"] == 12.346
+
+    def test_to_report_dict__scoring_failed__marks_assertion_failed_with_reason(self):
+        test_results_list = [
+            test_result.TestResult(
+                test_case=test_case.TestCase(
+                    trace_id="trace-1",
+                    dataset_item_id="item-1",
+                    task_output={"output": "test"},
+                    dataset_item_content={},
+                    dataset_item=dataset_item.DatasetItem(
+                        id="item-1",
+                        execution_policy=dataset_item.ExecutionPolicyItem(
+                            runs_per_item=1,
+                            pass_threshold=1,
+                        ),
+                    ),
+                ),
+                score_results=[
+                    score_result.ScoreResult(
+                        name="A1",
+                        value=0,
+                        scoring_failed=True,
+                        reason="Model error",
+                    ),
+                ],
+                trial_id=0,
+            )
+        ]
+        suite_result = _make_suite_result(test_results_list)
+
+        result = suite_result.to_report_dict()
+
+        assertion = result["items"][0]["runs"][0]["assertions"][0]
+        assert assertion["passed"] is False
+        assert assertion["scoring_failed"] is True
+        assert assertion["reason"] == "Model error"
+
+
+class TestSaveReport:
+    def test_save_report__valid_input__writes_json_file(self, tmp_path):
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(
+            test_results_list, suite_name="Test Suite", total_time=5.0
+        )
+        output_path = str(tmp_path / "report.json")
+
+        result_path = file_writer.save_report(suite_result, output_path)
+
+        assert result_path == output_path
+        assert os.path.exists(output_path)
+
+        with open(output_path) as f:
+            data = json.load(f)
+
+        assert data["suite_name"] == "Test Suite"
+        assert data["suite_passed"] is True
+        assert len(data["items"]) == 1
+
+    def test_save_report__no_path__uses_experiment_name(self, tmp_path, monkeypatch):
+        monkeypatch.chdir(tmp_path)
+
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(test_results_list)
+
+        result_path = file_writer.save_report(suite_result)
+
+        assert "my-experiment" in os.path.basename(result_path)
+        assert result_path.endswith(".json")
+        assert os.path.exists(result_path)
+
+    def test_save_report__nested_path__creates_parent_directories(self, tmp_path):
+        output_path = str(tmp_path / "nested" / "dir" / "report.json")
+
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(test_results_list)
+
+        result_path = file_writer.save_report(suite_result, output_path)
+
+        assert os.path.exists(result_path)
+
+
+class TestBuildDefaultReportPath:
+    def test_build_default_report_path__unsafe_characters__replaces_with_underscore(
+        self,
+    ):
+        path = file_writer.build_default_report_path("my/suite:name")
+        assert os.path.basename(path) == "my_suite_name.json"
+
+    def test_build_default_report_path__safe_characters__keeps_unchanged(self):
+        path = file_writer.build_default_report_path("my-suite_v1.0")
+        assert os.path.basename(path) == "my-suite_v1.0.json"
+
+    def test_build_default_report_path__spaces__replaces_with_underscore(self):
+        path = file_writer.build_default_report_path("my suite")
+        assert os.path.basename(path) == "my_suite.json"
+
+
+class TestTestSuiteResultMethods:
+    def test_to_dict__passing_suite__returns_dict_with_suite_passed_true(self):
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(test_results_list)
+
+        result = suite_result.to_dict()
+
+        assert isinstance(result, dict)
+        assert result["suite_passed"] is True
+        assert "items" in result
+
+    def test_to_report_dict__save_to_file__produces_valid_json(self, tmp_path):
+        test_results_list = [
+            _make_test_result(
+                dataset_item_id="item-1",
+                trial_id=0,
+                scores=[("A1", True)],
+                execution_policy={"runs_per_item": 1, "pass_threshold": 1},
+            ),
+        ]
+        suite_result = _make_suite_result(test_results_list)
+        output_path = str(tmp_path / "result.json")
+
+        path = file_writer.save_report(suite_result, output_path)
+
+        assert os.path.exists(path)
+        with open(path) as f:
+            data = json.load(f)
+        assert data["suite_passed"] is True

sdks/python/tests/unit/api_objects/experiment/test_helpers.py

@@ -0,0 +1,119 @@
+import types
+
+import pytest
+
+from opik.api_objects import experiment
+from tests.conftest import random_chars
+
+
+def fake_prompt(with_postfix: bool = False):
+    postfix = random_chars()
+
+    def __internal_api__to_info_dict__():
+        return {
+            "name": fake_prompt_obj.name,
+            "version": {
+                "template": fake_prompt_obj.prompt,
+            },
+        }
+
+    fake_prompt_obj = types.SimpleNamespace(
+        __internal_api__version_id__="some-prompt-version-id",
+        prompt="some-prompt-value",
+        name="some-prompt-name",
+        __internal_api__to_info_dict__=__internal_api__to_info_dict__,
+    )
+
+    if with_postfix:
+        fake_prompt_obj.prompt += postfix
+        fake_prompt_obj.__internal_api__version_id__ += postfix
+        fake_prompt_obj.name += postfix
+
+    return fake_prompt_obj
+
+
+@pytest.mark.parametrize(
+    argnames="input_kwargs,expected",
+    argvalues=[
+        (
+            {"experiment_config": None, "prompts": None},
+            {"metadata": None, "prompt_versions": None},
+        ),
+        (
+            {"experiment_config": {}, "prompts": None},
+            {"metadata": None, "prompt_versions": None},
+        ),
+        (
+            {"experiment_config": None, "prompts": [fake_prompt()]},
+            {
+                "metadata": {"prompts": {"some-prompt-name": "some-prompt-value"}},
+                "prompt_versions": [{"id": "some-prompt-version-id"}],
+            },
+        ),
+        (
+            {"experiment_config": {}, "prompts": [fake_prompt()]},
+            {
+                "metadata": {"prompts": {"some-prompt-name": "some-prompt-value"}},
+                "prompt_versions": [{"id": "some-prompt-version-id"}],
+            },
+        ),
+        (
+            {"experiment_config": {"some-key": "some-value"}, "prompts": None},
+            {"metadata": {"some-key": "some-value"}, "prompt_versions": None},
+        ),
+        (
+            {
+                "experiment_config": "NOT-DICT-VALUE-THAT-WILL-BE-IGNORED-AND-REPLACED-WITH-DICT-WITH-PROMPT",
+                "prompts": [fake_prompt()],
+            },
+            {
+                "metadata": {"prompts": {"some-prompt-name": "some-prompt-value"}},
+                "prompt_versions": [{"id": "some-prompt-version-id"}],
+            },
+        ),
+    ],
+)
+def test_experiment_build_metadata_from_prompt_versions(input_kwargs, expected):
+    metadata, prompt_versions = experiment.build_metadata_and_prompt_versions(
+        **input_kwargs
+    )
+
+    assert metadata == expected["metadata"]
+    assert prompt_versions == expected["prompt_versions"]
+
+
+def test_check_prompt_args_with_none_arguments():
+    result = experiment.handle_prompt_args(prompt=None, prompts=None)
+    assert result is None
+
+
+def test_check_prompt_args_with_none_and_empty_list():
+    result = experiment.handle_prompt_args(prompt=None, prompts=[])
+    assert result is None
+
+
+def test_check_prompt_args_with_single_prompt():
+    mock_prompt = fake_prompt(with_postfix=True)
+    result = experiment.handle_prompt_args(prompt=mock_prompt, prompts=None)
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert result[0] == mock_prompt
+
+
+def test_check_prompt_args_with_prompts_list():
+    mock_prompt_1 = fake_prompt(with_postfix=True)
+    mock_prompt_2 = fake_prompt(with_postfix=True)
+    prompts = [mock_prompt_1, mock_prompt_2]
+    result = experiment.handle_prompt_args(prompt=None, prompts=prompts)
+    assert result == prompts
+
+
+def test_check_prompt_args_with_both_prompt_and_prompts():
+    mock_prompt = fake_prompt(with_postfix=True)
+    mock_prompt_list = [
+        fake_prompt(with_postfix=True),
+        fake_prompt(with_postfix=True),
+    ]
+    result = experiment.handle_prompt_args(prompt=mock_prompt, prompts=mock_prompt_list)
+    assert isinstance(result, list)
+    assert result == mock_prompt_list

sdks/python/tests/unit/api_objects/prompt/test_chat_prompt_template.py

@@ -0,0 +1,587 @@
+import pytest
+from opik.api_objects.prompt import ChatPromptTemplate, PromptType
+from opik import exceptions
+
+
+def test_chat_prompt_template__format__simple_text_message__happyflow():
+    """Test basic formatting of a simple text message."""
+    messages = [
+        {
+            "role": "user",
+            "content": "Hi, my name is {{name}} and I live in {{city}}.",
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format({"name": "Harry", "city": "London"})
+
+    assert result == [
+        {"role": "user", "content": "Hi, my name is Harry and I live in London."}
+    ]
+
+
+def test_chat_prompt_template__format__multiple_messages__happyflow():
+    """Test formatting multiple messages with different roles."""
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant in {{location}}."},
+        {"role": "user", "content": "What is the capital of {{country}}?"},
+        {"role": "assistant", "content": "The capital is {{capital}}."},
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format(
+        {
+            "location": "London",
+            "country": "France",
+            "capital": "Paris",
+        }
+    )
+
+    assert result == [
+        {"role": "system", "content": "You are a helpful assistant in London."},
+        {"role": "user", "content": "What is the capital of France?"},
+        {"role": "assistant", "content": "The capital is Paris."},
+    ]
+
+
+def test_chat_prompt_template__format__multimodal_content__happyflow():
+    """Test formatting messages with multimodal content (text + image) when vision is supported."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this {{object}}:"},
+                {"type": "image_url", "image_url": {"url": "{{image_url}}"}},
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    # Format with vision supported to get structured content
+    result = tested.format(
+        {"object": "painting", "image_url": "https://example.com/image.jpg"},
+        supported_modalities={"vision": True},
+    )
+
+    assert result == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this painting:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/image.jpg"},
+                },
+            ],
+        }
+    ]
+
+
+def test_chat_prompt_template__format__jinja2_template__happyflow():
+    """Test formatting with Jinja2 template type."""
+    messages = [
+        {
+            "role": "user",
+            "content": "Hi, my name is {{ name }} and I live in {{ city }}.",
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages, template_type=PromptType.JINJA2)
+
+    result = tested.format({"name": "Harry", "city": "London"})
+
+    assert result == [
+        {"role": "user", "content": "Hi, my name is Harry and I live in London."}
+    ]
+
+
+def test_chat_prompt_template__format__jinja2_with_control_flow():
+    """Test Jinja2 formatting with control flow."""
+    messages = [
+        {
+            "role": "user",
+            "content": """
+            {% if is_wizard %}
+            {{ name }} is a wizard who lives in {{ city }}.
+            {% else %}
+            {{ name }} is a muggle who lives in {{ city }}.
+            {% endif %}
+            """,
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages, template_type=PromptType.JINJA2)
+
+    wizard_result = tested.format(
+        {"name": "Harry", "city": "London", "is_wizard": True}
+    )
+    assert (
+        "Harry is a wizard who lives in London." in wizard_result[0]["content"].strip()
+    )
+
+    muggle_result = tested.format(
+        {"name": "Dudley", "city": "Surrey", "is_wizard": False}
+    )
+    assert (
+        "Dudley is a muggle who lives in Surrey." in muggle_result[0]["content"].strip()
+    )
+
+
+def test_chat_prompt_template__format__jinja2_with_loops():
+    """Test Jinja2 formatting with loops."""
+    messages = [
+        {
+            "role": "user",
+            "content": """
+            {{ name }}'s friends are:
+            {% for friend in friends %}
+            - {{ friend }}
+            {% endfor %}
+            """,
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages, template_type=PromptType.JINJA2)
+
+    result = tested.format({"name": "Harry", "friends": ["Ron", "Hermione", "Neville"]})
+
+    content = result[0]["content"]
+    assert "Harry's friends are:" in content
+    assert "- Ron" in content
+    assert "- Hermione" in content
+    assert "- Neville" in content
+
+
+def test_chat_prompt_template__format__empty_content():
+    """Test formatting with empty content."""
+    messages = [
+        {"role": "user", "content": ""},
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format({})
+
+    assert result == [{"role": "user", "content": ""}]
+
+
+def test_chat_prompt_template__format__message_without_role__skipped():
+    """Test that messages without a role are skipped."""
+    messages = [
+        {"role": "user", "content": "Hello {{name}}"},
+        {"content": "This message has no role"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format({"name": "Harry"})
+
+    # Only messages with roles should be included
+    assert result == [
+        {"role": "user", "content": "Hello Harry"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+
+
+def test_chat_prompt_template__required_modalities__text_only():
+    """Test required_modalities returns empty set for text-only messages."""
+    messages = [
+        {"role": "user", "content": "Simple text message"},
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.required_modalities()
+
+    assert result == set()
+
+
+def test_chat_prompt_template__required_modalities__with_vision():
+    """Test required_modalities detects vision modality."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this image:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/img.jpg"},
+                },
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.required_modalities()
+
+    assert "vision" in result
+
+
+def test_chat_prompt_template__required_modalities__multiple_messages():
+    """Test required_modalities across multiple messages."""
+    messages = [
+        {"role": "user", "content": "Text only"},
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "With image:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/img.jpg"},
+                },
+            ],
+        },
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.required_modalities()
+
+    assert "vision" in result
+
+
+def test_chat_prompt_template__format__unsupported_modality_replaced_with_placeholder():
+    """Test that unsupported modalities are replaced with placeholders."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this image:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/img.jpg"},
+                },
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    # Format with vision not supported
+    result = tested.format({}, supported_modalities={"vision": False})
+
+    assert len(result) == 1
+    # When vision is not supported, image should be replaced with placeholder
+    content = result[0]["content"]
+    assert isinstance(content, str)
+    assert "<<<image>>>" in content
+    assert "<<</image>>>" in content
+    assert "Describe this image:" in content
+
+
+def test_chat_prompt_template__format__supported_modality_preserved():
+    """Test that supported modalities are preserved as structured content."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this image:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/img.jpg"},
+                },
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    # Format with vision supported
+    result = tested.format({}, supported_modalities={"vision": True})
+
+    assert result == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this image:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/img.jpg"},
+                },
+            ],
+        }
+    ]
+
+
+def test_chat_prompt_template__format__multimodal_with_template_variables():
+    """Test multimodal content with template variables in both text and image URL."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Analyze this {{object}}:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/{{filename}}"},
+                },
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format(
+        {"object": "diagram", "filename": "diagram.png"},
+        supported_modalities={"vision": True},
+    )
+
+    assert result == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Analyze this diagram:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "https://example.com/diagram.png"},
+                },
+            ],
+        }
+    ]
+
+
+def test_chat_prompt_template__format__image_with_detail_parameter():
+    """Test that image detail parameter is preserved during formatting."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Analyze:"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "{{url}}", "detail": "high"},
+                },
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format(
+        {"url": "https://example.com/img.jpg"}, supported_modalities={"vision": True}
+    )
+
+    assert result == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Analyze:"},
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": "https://example.com/img.jpg",
+                        "detail": "high",
+                    },
+                },
+            ],
+        }
+    ]
+
+
+def test_chat_prompt_template__messages_property():
+    """Test that messages property returns the original messages."""
+    messages = [
+        {"role": "user", "content": "Hello {{name}}"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    assert tested.messages == messages
+
+
+def test_chat_prompt_template__format__override_template_type():
+    """Test that template_type can be overridden in format call."""
+    messages = [
+        {"role": "user", "content": "Name: {{name}}, City: {{city}}"},
+    ]
+
+    # Create with Mustache default
+    tested = ChatPromptTemplate(messages, template_type=PromptType.MUSTACHE)
+
+    # Override with Jinja2 in format
+    result = tested.format(
+        {"name": "Harry", "city": "London"},
+        template_type=PromptType.JINJA2,
+    )
+
+    assert result == [{"role": "user", "content": "Name: Harry, City: London"}]
+
+
+def test_chat_prompt_template__format__one_placeholder_used_multiple_times():
+    """Test formatting with the same placeholder used multiple times."""
+    messages = [
+        {
+            "role": "user",
+            "content": "My name is {{name}}. I repeat, my name is {{name}}.",
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format({"name": "Harry"})
+
+    assert result == [
+        {"role": "user", "content": "My name is Harry. I repeat, my name is Harry."}
+    ]
+
+
+def test_chat_prompt_template__format__empty_messages_list():
+    """Test formatting with empty messages list."""
+    messages = []
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format({})
+
+    assert result == []
+
+
+def test_chat_prompt_template__format__message_with_missing_content():
+    """Test formatting when message has no content field."""
+    messages = [
+        {"role": "user"},  # No content field
+    ]
+
+    tested = ChatPromptTemplate(messages)
+
+    result = tested.format({})
+
+    assert result == [{"role": "user", "content": ""}]
+
+
+def test_chat_prompt_template__format__passed_arguments_not_in_template__error_raised():
+    """Test that extra format arguments not in template raise an error."""
+    messages = [
+        {"role": "user", "content": "Hi, my name is {{name}}, I live in {{city}}."}
+    ]
+
+    tested = ChatPromptTemplate(messages, validate_placeholders=True)
+
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format({"name": "Harry", "city": "London", "nemesis_name": "Voldemort"})
+
+    assert exc_info.value.format_arguments == set(["name", "city", "nemesis_name"])
+    assert exc_info.value.prompt_placeholders == set(["name", "city"])
+    assert exc_info.value.symmetric_difference == set(["nemesis_name"])
+
+
+def test_chat_prompt_template__format__some_placeholders_missing__error_raised():
+    """Test that missing required placeholders raise an error."""
+    messages = [
+        {"role": "user", "content": "Hi, my name is {{name}}, I live in {{city}}."}
+    ]
+
+    tested = ChatPromptTemplate(messages, validate_placeholders=True)
+
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format({"name": "Harry"})
+
+    assert exc_info.value.format_arguments == set(["name"])
+    assert exc_info.value.prompt_placeholders == set(["name", "city"])
+    assert exc_info.value.symmetric_difference == set(["city"])
+
+
+def test_chat_prompt_template__format__placeholders_mismatch_both_ways__error_raised():
+    """Test error when some placeholders are missing AND extra arguments provided."""
+    messages = [
+        {"role": "user", "content": "Hi, my name is {{name}}, I live in {{city}}."}
+    ]
+
+    tested = ChatPromptTemplate(messages, validate_placeholders=True)
+
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format({"name": "Harry", "nemesis_name": "Voldemort"})
+
+    assert exc_info.value.format_arguments == set(["name", "nemesis_name"])
+    assert exc_info.value.prompt_placeholders == set(["name", "city"])
+    assert exc_info.value.symmetric_difference == set(["city", "nemesis_name"])
+
+
+def test_chat_prompt_template__format__multimodal_placeholders__validates_all():
+    """Test that placeholders in multimodal content are validated."""
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe this {{object}}:"},
+                {"type": "image_url", "image_url": {"url": "{{image_url}}"}},
+            ],
+        }
+    ]
+
+    tested = ChatPromptTemplate(messages, validate_placeholders=True)
+
+    # Missing image_url placeholder
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format({"object": "painting"})
+
+    assert exc_info.value.format_arguments == set(["object"])
+    assert exc_info.value.prompt_placeholders == set(["object", "image_url"])
+    assert exc_info.value.symmetric_difference == set(["image_url"])
+
+
+def test_chat_prompt_template__format__multiple_messages_placeholders__validates_all():
+    """Test that placeholders across multiple messages are validated."""
+    messages = [
+        {"role": "system", "content": "You are an assistant in {{location}}."},
+        {"role": "user", "content": "What is the capital of {{country}}?"},
+    ]
+
+    tested = ChatPromptTemplate(messages, validate_placeholders=True)
+
+    # Missing country placeholder
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format({"location": "London"})
+
+    assert exc_info.value.format_arguments == set(["location"])
+    assert exc_info.value.prompt_placeholders == set(["location", "country"])
+    assert exc_info.value.symmetric_difference == set(["country"])
+
+
+def test_chat_prompt_template__format__validation_disabled__no_error():
+    """Test that validation can be disabled."""
+    messages = [
+        {"role": "user", "content": "Hi, my name is {{name}}, I live in {{city}}."}
+    ]
+
+    tested = ChatPromptTemplate(messages, validate_placeholders=False)
+
+    # Should not raise an error even though city is missing
+    result = tested.format({"name": "Harry"})
+
+    # The template will leave unformatted placeholders
+    assert result == [
+        {"role": "user", "content": "Hi, my name is Harry, I live in {{city}}."}
+    ]
+
+
+def test_chat_prompt_template__format__jinja2_no_validation():
+    """Test that Jinja2 templates don't validate placeholders."""
+    messages = [
+        {"role": "user", "content": "Hi, my name is {{ name }}, I live in {{ city }}."}
+    ]
+
+    tested = ChatPromptTemplate(messages, template_type=PromptType.JINJA2)
+
+    # Jinja2 templates don't validate, so this should not raise an error
+    # Jinja2 will just render missing variables as empty
+    result = tested.format({"name": "Harry"})
+
+    # Should render successfully (Jinja2 handles missing variables gracefully)
+    assert len(result) == 1
+    assert result[0]["role"] == "user"

sdks/python/tests/unit/api_objects/prompt/test_mask_context.py

@@ -0,0 +1,59 @@
+from opik.api_objects.prompt.mask_context import (
+    get_active_prompt_masks,
+    get_mask_for_prompt,
+    prompt_mask_context,
+)
+
+
+class TestPromptMaskContext:
+    def test_no_context__returns_none(self):
+        assert get_active_prompt_masks() is None
+
+    def test_inside_context__returns_masks_dict(self):
+        masks = {"prompt-1": "mask-a", "prompt-2": "mask-b"}
+        with prompt_mask_context(masks):
+            assert get_active_prompt_masks() == masks
+
+    def test_get_mask_for_prompt__prompt_in_dict__returns_mask_id(self):
+        masks = {"prompt-uuid": "mask-uuid"}
+        with prompt_mask_context(masks):
+            assert get_mask_for_prompt("prompt-uuid") == "mask-uuid"
+
+    def test_get_mask_for_prompt__prompt_not_in_dict__returns_none(self):
+        masks = {"other-prompt": "mask-abc"}
+        with prompt_mask_context(masks):
+            assert get_mask_for_prompt("unknown-prompt") is None
+
+    def test_get_mask_for_prompt__no_context__returns_none(self):
+        assert get_mask_for_prompt("any-prompt-id") is None
+
+    def test_after_exit__resets_to_none(self):
+        with prompt_mask_context({"prompt-1": "mask-1"}):
+            pass
+        assert get_active_prompt_masks() is None
+
+    def test_context_with_none_masks__returns_none(self):
+        with prompt_mask_context(None):
+            assert get_active_prompt_masks() is None
+            assert get_mask_for_prompt("any-id") is None
+
+    def test_exception_inside__resets(self):
+        try:
+            with prompt_mask_context({"p": "m"}):
+                assert get_active_prompt_masks() == {"p": "m"}
+                raise ValueError("boom")
+        except ValueError:
+            pass
+        assert get_active_prompt_masks() is None
+
+    def test_nested_contexts__inner_wins_then_restores_outer(self):
+        outer = {"p1": "m1"}
+        inner = {"p2": "m2"}
+        with prompt_mask_context(outer):
+            assert get_active_prompt_masks() == outer
+            with prompt_mask_context(inner):
+                assert get_active_prompt_masks() == inner
+                assert get_mask_for_prompt("p2") == "m2"
+                assert get_mask_for_prompt("p1") is None
+            assert get_active_prompt_masks() == outer
+        assert get_active_prompt_masks() is None

sdks/python/tests/unit/api_objects/prompt/test_prompt_cache.py

@@ -0,0 +1,500 @@
+"""Unit tests for the prompt client-side cache."""
+
+import time
+from typing import Optional
+from unittest import mock
+
+import pytest
+
+from opik import config as opik_config
+from opik.api_objects.prompt import prompt_cache
+from opik.api_objects.prompt.prompt_cache import (
+    PromptCache,
+    get_global_cache,
+)
+
+
+def _get_ttl() -> int:
+    return opik_config.OpikConfig().prompt_cache_ttl_seconds
+
+
+@pytest.fixture(autouse=True)
+def clear_caches_after_test():
+    yield
+    get_global_cache().clear()
+
+
+@pytest.fixture
+def cache():
+    c = PromptCache()
+    yield c
+    c.clear()
+
+
+def _make_mock_prompt(
+    name: str = "my-prompt",
+    commit: str = "abc123",
+    prompt_id: Optional[str] = None,
+) -> mock.Mock:
+    p = mock.Mock()
+    p.name = name
+    p.commit = commit
+    p.__internal_api__prompt_id__ = prompt_id
+    return p
+
+
+class TestPromptCache:
+    def test_get__missing_key__returns_none(self, cache):
+        assert cache.get(("p", None, None, "text")) is None
+
+    def test_get_or_fetch__cache_miss__fetches_and_caches(self, cache):
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+        result = cache.get_or_fetch(
+            ("p", None, None, "text"), fetch_fn, ttl_seconds=300
+        )
+        assert result is p
+        fetch_fn.assert_called_once()
+
+    def test_get_or_fetch__cache_hit__does_not_call_fetch_fn(self, cache):
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+        cache.get_or_fetch(("p", None, None, "text"), fetch_fn, ttl_seconds=300)
+        fetch_fn.reset_mock()
+        result = cache.get_or_fetch(
+            ("p", None, None, "text"), fetch_fn, ttl_seconds=300
+        )
+        assert result is p
+        fetch_fn.assert_not_called()
+
+    def test_get_or_fetch__fetch_returns_none__returns_none(self, cache):
+        fetch_fn = mock.Mock(return_value=None)
+        result = cache.get_or_fetch(
+            ("missing", None, None, "text"), fetch_fn, ttl_seconds=300
+        )
+        assert result is None
+
+    def test_clear__populated_cache__removes_all_entries(self, cache):
+        p = _make_mock_prompt()
+        cache.get_or_fetch(
+            ("p", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=None
+        )
+        cache.clear()
+        assert cache.get(("p", None, None, "text")) is None
+
+    def test_clear__running_refresh_thread__stops_thread(self, cache):
+        p = _make_mock_prompt()
+        cache.get_or_fetch(
+            ("p", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=300
+        )
+        assert cache._thread is not None and cache._thread.is_alive()
+        cache.clear()
+        assert cache._thread is None
+
+    def test_refresh_thread__unpinned_entry__starts_thread(self, cache):
+        p = _make_mock_prompt()
+        cache.get_or_fetch(
+            ("p", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=300
+        )
+        assert cache._thread is not None
+        assert cache._thread.is_alive()
+
+    def test_refresh_thread__pinned_entry__does_not_start_thread(self, cache):
+        p = _make_mock_prompt()
+        cache.get_or_fetch(
+            ("p", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=None
+        )
+        assert cache._thread is None
+
+    def test_get_or_fetch__different_keys__returns_separate_entries(self, cache):
+        p1 = _make_mock_prompt(commit="c1")
+        p2 = _make_mock_prompt(commit="c2")
+        cache.get_or_fetch(
+            ("p", "c1", None, "text"), mock.Mock(return_value=p1), ttl_seconds=None
+        )
+        cache.get_or_fetch(
+            ("p", "c2", None, "text"), mock.Mock(return_value=p2), ttl_seconds=None
+        )
+        assert cache.get(("p", "c1", None, "text")) is p1
+        assert cache.get(("p", "c2", None, "text")) is p2
+
+
+class TestBackgroundRefresh:
+    def test_refresh__stale_entry__updates_prompt(self):
+        new_prompt = _make_mock_prompt(commit=None)
+        callback = mock.Mock(return_value=new_prompt)
+
+        cache = get_global_cache()
+        base = time.monotonic()
+        with mock.patch("opik.api_objects.prompt.prompt_cache.time") as mock_time:
+            mock_time.monotonic.return_value = base
+            prompt_cache.get_or_fetch("p", None, None, "text", callback)
+            callback.reset_mock()
+            callback.return_value = new_prompt
+
+            mock_time.monotonic.return_value = base + _get_ttl() + 1
+            mock_time.sleep = mock.Mock()
+            cache._refresh_stale_entries()
+
+            assert callback.call_count >= 1
+            assert cache.get(("p", None, None, "text", None, None)) is new_prompt
+
+    def test_refresh__non_stale_entry__skips_callback(self):
+        p = _make_mock_prompt(commit=None)
+        callback = mock.Mock(return_value=p)
+
+        cache = get_global_cache()
+        base = time.monotonic()
+        with mock.patch("opik.api_objects.prompt.prompt_cache.time") as mock_time:
+            mock_time.monotonic.return_value = base
+            prompt_cache.get_or_fetch("p", None, None, "text", callback)
+            callback.reset_mock()
+
+            mock_time.monotonic.return_value = base + 0.2
+            cache._refresh_stale_entries()
+
+            callback.assert_not_called()
+
+    def test_refresh__callback_raises__thread_survives(self):
+        p = _make_mock_prompt(commit=None)
+        callback = mock.Mock(side_effect=[p, RuntimeError("boom")])
+
+        cache = get_global_cache()
+        base = time.monotonic()
+        with mock.patch("opik.api_objects.prompt.prompt_cache.time") as mock_time:
+            mock_time.monotonic.return_value = base
+            prompt_cache.get_or_fetch("p", None, None, "text", callback)
+
+            mock_time.monotonic.return_value = base + _get_ttl() + 1
+            cache._refresh_stale_entries()
+
+            assert cache._thread.is_alive()
+
+
+class TestGetOrFetch:
+    def test_get_or_fetch__cache_miss__calls_fetch_fn(self):
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+        result = prompt_cache.get_or_fetch("p", None, None, "text", fetch_fn)
+        assert result is p
+        fetch_fn.assert_called_once()
+
+    def test_get_or_fetch__fetch_returns_none__returns_none(self):
+        fetch_fn = mock.Mock(return_value=None)
+        result = prompt_cache.get_or_fetch("missing", None, None, "text", fetch_fn)
+        assert result is None
+
+    def test_get_or_fetch__cache_hit__does_not_call_fetch_fn(self):
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+        prompt_cache.get_or_fetch("p", None, None, "text", fetch_fn)
+        fetch_fn.reset_mock()
+
+        result = prompt_cache.get_or_fetch("p", None, None, "text", fetch_fn)
+        assert result is p
+        fetch_fn.assert_not_called()
+
+    def test_get_or_fetch__pinned_commit__no_refresh_thread(self):
+        p = _make_mock_prompt(commit="abc")
+        fetch_fn = mock.Mock(return_value=p)
+        prompt_cache.get_or_fetch("p", "abc", None, "text", fetch_fn)
+        cache = get_global_cache()
+        assert cache._thread is None
+
+    def test_get_or_fetch__unpinned_commit__starts_refresh_thread(self):
+        p = _make_mock_prompt(commit=None)
+        fetch_fn = mock.Mock(return_value=p)
+        prompt_cache.get_or_fetch("p", None, None, "text", fetch_fn)
+        cache = get_global_cache()
+        assert cache._thread is not None
+        assert cache._thread.is_alive()
+
+    def test_get_or_fetch__cache_hit__returns_same_object(self):
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+        prompt_cache.get_or_fetch("p", None, "proj", "text", fetch_fn)
+        first = prompt_cache.get_or_fetch("p", None, "proj", "text", fetch_fn)
+        second = prompt_cache.get_or_fetch("p", None, "proj", "text", fetch_fn)
+        assert first is second is p
+
+    def test_get_or_fetch__different_keys__returns_separate_entries(self):
+        p1 = _make_mock_prompt(commit="c1")
+        p2 = _make_mock_prompt(commit="c2")
+        prompt_cache.get_or_fetch("p", "c1", None, "text", mock.Mock(return_value=p1))
+        prompt_cache.get_or_fetch("p", "c2", None, "text", mock.Mock(return_value=p2))
+        r1 = prompt_cache.get_or_fetch("p", "c1", None, "text", mock.Mock())
+        r2 = prompt_cache.get_or_fetch("p", "c2", None, "text", mock.Mock())
+        assert r1 is p1
+        assert r2 is p2
+
+
+class TestGetOrFetchVersionSelector:
+    """Tests for the ``version`` parameter on the module-level ``get_or_fetch``."""
+
+    def test_get_or_fetch__different_versions__return_separate_entries(self):
+        p1 = _make_mock_prompt(commit="commitA")
+        p2 = _make_mock_prompt(commit="commitB")
+        fetch1 = mock.Mock(return_value=p1)
+        fetch2 = mock.Mock(return_value=p2)
+
+        first = prompt_cache.get_or_fetch("p", None, None, "text", fetch1, version="v1")
+        second = prompt_cache.get_or_fetch(
+            "p", None, None, "text", fetch2, version="v2"
+        )
+
+        assert first is p1
+        assert second is p2
+        fetch1.assert_called_once()
+        fetch2.assert_called_once()
+
+    def test_get_or_fetch__same_version__second_call_hits_cache(self):
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+
+        first = prompt_cache.get_or_fetch(
+            "p", None, None, "text", fetch_fn, version="v2"
+        )
+        second = prompt_cache.get_or_fetch(
+            "p", None, None, "text", fetch_fn, version="v2"
+        )
+
+        assert first is second is p
+        fetch_fn.assert_called_once()
+
+    def test_get_or_fetch__commit_pin_and_version_pin__do_not_collide(self):
+        p_commit = _make_mock_prompt(commit="abc12345")
+        p_version = _make_mock_prompt(commit="def67890")
+        fetch_commit = mock.Mock(return_value=p_commit)
+        fetch_version = mock.Mock(return_value=p_version)
+
+        from_commit = prompt_cache.get_or_fetch(
+            "p", "abc12345", None, "text", fetch_commit
+        )
+        from_version = prompt_cache.get_or_fetch(
+            "p", None, None, "text", fetch_version, version="v3"
+        )
+
+        assert from_commit is p_commit
+        assert from_version is p_version
+        fetch_commit.assert_called_once()
+        fetch_version.assert_called_once()
+
+    def test_get_or_fetch__version_selector__starts_refresh_thread(self):
+        # Sequential versions can be reassigned by the backend if the underlying
+        # version is deleted and recreated, so they must follow the normal TTL
+        # refresh (not pinned indefinitely).
+        p = _make_mock_prompt()
+        fetch_fn = mock.Mock(return_value=p)
+        prompt_cache.get_or_fetch("p", None, None, "text", fetch_fn, version="v1")
+        cache = get_global_cache()
+        assert cache._thread is not None
+        assert cache._thread.is_alive()
+
+
+class TestPromptCacheEdgeCases:
+    def test_get_or_fetch__multiple_unpinned_inserts__reuses_single_refresh_thread(
+        self, cache
+    ):
+        p = _make_mock_prompt()
+        cache.get_or_fetch(
+            ("a", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=300
+        )
+        thread1 = cache._thread
+        cache.get_or_fetch(
+            ("b", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=300
+        )
+        assert cache._thread is thread1
+
+    def test_clear__on_empty_cache__is_noop(self, cache):
+        cache.clear()
+        assert cache.get(("any", None, None, "text")) is None
+
+    def test_refresh__pinned_entry__not_refreshed_even_after_ttl(self, cache):
+        p = _make_mock_prompt()
+        callback = mock.Mock(return_value=p)
+
+        base = time.monotonic()
+        with mock.patch("opik.api_objects.prompt.prompt_cache.time") as mock_time:
+            mock_time.monotonic.return_value = base
+            cache.get_or_fetch(("p", "v1", None, "text"), callback, ttl_seconds=None)
+            callback.reset_mock()
+
+            mock_time.monotonic.return_value = base + 0.2
+            cache._refresh_stale_entries()
+
+            callback.assert_not_called()
+        assert cache.get(("p", "v1", None, "text")) is p
+
+    def test_get_or_fetch__fetch_returns_none__not_cached(self, cache):
+        fetch_fn = mock.Mock(return_value=None)
+        cache.get_or_fetch(("p", None, None, "text"), fetch_fn, ttl_seconds=300)
+        assert cache.get(("p", None, None, "text")) is None
+
+    def test_refresh__callback_returns_none__preserves_original_prompt(self):
+        original = _make_mock_prompt(commit=None)
+        callback = mock.Mock(side_effect=[original, None])
+
+        cache = get_global_cache()
+        base = time.monotonic()
+        with mock.patch("opik.api_objects.prompt.prompt_cache.time") as mock_time:
+            mock_time.monotonic.return_value = base
+            prompt_cache.get_or_fetch("p", None, None, "text", callback)
+
+            mock_time.monotonic.return_value = base + _get_ttl() + 1
+            cache._refresh_stale_entries()
+
+        assert cache.get(("p", None, None, "text", None, None)) is original
+
+    def test_lru_eviction__oldest_entry_removed_when_max_size_exceeded(self):
+        cache = PromptCache(max_size=3)
+        prompts = [_make_mock_prompt(commit=f"c{i}") for i in range(4)]
+        for i, p in enumerate(prompts):
+            cache.get_or_fetch(
+                (f"p{i}", f"c{i}", None, "text"),
+                mock.Mock(return_value=p),
+                ttl_seconds=None,
+            )
+
+        assert cache.get(("p0", "c0", None, "text")) is None
+        assert cache.get(("p1", "c1", None, "text")) is prompts[1]
+        assert cache.get(("p2", "c2", None, "text")) is prompts[2]
+        assert cache.get(("p3", "c3", None, "text")) is prompts[3]
+        cache.clear()
+
+    def test_lru_eviction__access_refreshes_position(self):
+        cache = PromptCache(max_size=3)
+        prompts = [_make_mock_prompt(commit=f"c{i}") for i in range(3)]
+        for i, p in enumerate(prompts):
+            cache.get_or_fetch(
+                (f"p{i}", f"c{i}", None, "text"),
+                mock.Mock(return_value=p),
+                ttl_seconds=None,
+            )
+
+        # Access p0 so it becomes most-recently-used
+        cache.get(("p0", "c0", None, "text"))
+
+        # Insert a 4th entry — p1 (the actual LRU) should be evicted, not p0
+        p3 = _make_mock_prompt(commit="c3")
+        cache.get_or_fetch(
+            ("p3", "c3", None, "text"),
+            mock.Mock(return_value=p3),
+            ttl_seconds=None,
+        )
+
+        assert cache.get(("p0", "c0", None, "text")) is prompts[0]
+        assert cache.get(("p1", "c1", None, "text")) is None
+        assert cache.get(("p2", "c2", None, "text")) is prompts[2]
+        assert cache.get(("p3", "c3", None, "text")) is p3
+        cache.clear()
+
+    def test_clear__called_twice__is_safe(self, cache):
+        p = _make_mock_prompt()
+        cache.get_or_fetch(
+            ("p", None, None, "text"), mock.Mock(return_value=p), ttl_seconds=300
+        )
+        cache.clear()
+        cache.clear()
+        assert cache._thread is None
+
+
+class TestPromptAutoInjection:
+    """Test that get_prompt injects prompts into the active trace/span context via opik_prompts."""
+
+    @staticmethod
+    def _make_prompt_with_info_dict(name="my-prompt", commit="abc123", info_dict=None):
+        p = _make_mock_prompt(name=name, commit=commit)
+        p.__internal_api__to_info_dict__ = mock.Mock(
+            return_value=info_dict or {"name": name, "version": {"commit": commit}}
+        )
+        return p
+
+    def _call_get_prompt(self, cache_return_value):
+        from opik.api_objects import opik_client
+
+        client = opik_client.Opik()
+        with mock.patch(
+            "opik.api_objects.prompt.prompt_cache.get_or_fetch",
+            return_value=cache_return_value,
+        ):
+            return client.get_prompt(
+                name="my-prompt", commit="abc123", project_name=None
+            )
+
+    def test_get_prompt__in_track_context__injects_into_metadata(self):
+        info_dict = {"name": "my-prompt", "version": {"commit": "abc123"}}
+        mock_prompt = self._make_prompt_with_info_dict(info_dict=info_dict)
+
+        mock_trace_data = mock.Mock()
+        mock_trace_data.metadata = None
+        mock_span_data = mock.Mock()
+        mock_span_data.metadata = None
+
+        with (
+            mock.patch(
+                "opik.context_storage.get_trace_data", return_value=mock_trace_data
+            ),
+            mock.patch(
+                "opik.context_storage.top_span_data", return_value=mock_span_data
+            ),
+        ):
+            self._call_get_prompt(mock_prompt)
+
+        mock_trace_data.update.assert_called_once_with(
+            metadata={"opik_prompts": [info_dict]}
+        )
+        mock_span_data.update.assert_called_once_with(
+            metadata={"opik_prompts": [info_dict]}
+        )
+
+    def test_get_prompt__appends_to_existing_prompts(self):
+        existing_prompt_info = {"name": "old-prompt", "version": {"commit": "old123"}}
+        new_info_dict = {"name": "my-prompt", "version": {"commit": "abc123"}}
+        mock_prompt = self._make_prompt_with_info_dict(info_dict=new_info_dict)
+
+        mock_trace_data = mock.Mock()
+        mock_trace_data.metadata = {
+            "opik_prompts": [existing_prompt_info],
+            "other_key": "value",
+        }
+        mock_span_data = mock.Mock()
+        mock_span_data.metadata = {"opik_prompts": [existing_prompt_info]}
+
+        with (
+            mock.patch(
+                "opik.context_storage.get_trace_data", return_value=mock_trace_data
+            ),
+            mock.patch(
+                "opik.context_storage.top_span_data", return_value=mock_span_data
+            ),
+        ):
+            self._call_get_prompt(mock_prompt)
+
+        mock_trace_data.update.assert_called_once_with(
+            metadata={"opik_prompts": [existing_prompt_info, new_info_dict]}
+        )
+        mock_span_data.update.assert_called_once_with(
+            metadata={"opik_prompts": [existing_prompt_info, new_info_dict]}
+        )
+
+    def test_get_prompt__no_track_context__no_error(self):
+        """When there is no active trace context, injection silently does nothing."""
+        mock_prompt = self._make_prompt_with_info_dict()
+
+        with (
+            mock.patch("opik.context_storage.get_trace_data", return_value=None),
+            mock.patch("opik.context_storage.top_span_data", return_value=None),
+        ):
+            result = self._call_get_prompt(mock_prompt)
+            assert result is mock_prompt
+
+    def test_get_prompt__none_result__no_injection(self):
+        with (
+            mock.patch("opik.context_storage.get_trace_data") as mock_get_trace,
+            mock.patch("opik.context_storage.top_span_data") as mock_top_span,
+        ):
+            result = self._call_get_prompt(None)
+
+        assert result is None
+        mock_get_trace.assert_not_called()
+        mock_top_span.assert_not_called()

sdks/python/tests/unit/api_objects/prompt/test_prompt_client.py

@@ -0,0 +1,793 @@
+"""Unit tests for PromptClient to verify API endpoint selection logic."""
+
+from unittest import mock
+from typing import Optional
+
+import pytest
+
+from opik import exceptions
+from opik.api_objects import opik_client as opik_client_module
+from opik.api_objects.prompt import client as prompt_client
+from opik.api_objects.prompt import prompt_cache
+from opik.api_objects.prompt import types as prompt_types
+from opik.rest_api import core as rest_api_core
+from opik.rest_api.types import prompt_version_detail
+
+
+@pytest.fixture
+def mock_rest_client():
+    """Create a mock REST client."""
+    client = mock.Mock()
+    client.prompts = mock.Mock()
+    return client
+
+
+@pytest.fixture
+def client(mock_rest_client):
+    """Create a PromptClient with a mock REST client."""
+    return prompt_client.PromptClient(mock_rest_client)
+
+
+def _make_mock_version(
+    template: str = "test template",
+    version_type: str = "mustache",
+    metadata: Optional[dict] = None,
+    version_number: Optional[str] = "v1",
+    tags: Optional[list] = None,
+    environments: Optional[list] = None,
+) -> prompt_version_detail.PromptVersionDetail:
+    """Helper to create a mock PromptVersionDetail."""
+    return prompt_version_detail.PromptVersionDetail(
+        id="version-id",
+        prompt_id="prompt-id",
+        template=template,
+        type=version_type,
+        metadata=metadata,
+        commit="abc123",
+        version_number=version_number,
+        tags=tags,
+        environments=environments,
+        template_structure="text",
+    )
+
+
+def _make_404_error() -> rest_api_core.ApiError:
+    """Helper to create a 404 ApiError."""
+    error = rest_api_core.ApiError(status_code=404, body=None)
+    return error
+
+
+class TestPromptClientEndpointSelection:
+    """Tests to verify that PromptClient calls the correct REST endpoint based on parameters."""
+
+    def test_create_new_prompt_with_container_params_calls_create_prompt(
+        self, client, mock_rest_client
+    ):
+        """When creating a new prompt with id/description/tags, should call create_prompt endpoint.
+        Also asserts that version_number survives the tag-rebuild branch."""
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            _make_404_error(),
+            _make_mock_version(tags=None),
+        ]
+
+        result = client.create_prompt(
+            name="test-prompt",
+            prompt="test template",
+            metadata=None,
+            type=prompt_types.PromptType.MUSTACHE,
+            id="custom-id",
+            description="A test prompt",
+            tags=["test", "unit"],
+        )
+
+        mock_rest_client.prompts.create_prompt.assert_called_once()
+        mock_rest_client.prompts.create_prompt_version.assert_not_called()
+
+        call_kwargs = mock_rest_client.prompts.create_prompt.call_args[1]
+        assert call_kwargs["name"] == "test-prompt"
+        assert call_kwargs["id"] == "custom-id"
+        assert call_kwargs["description"] == "A test prompt"
+        assert call_kwargs["tags"] == ["test", "unit"]
+        assert result.version_number == "v1"
+        assert result.tags == ["test", "unit"]
+
+    def test_create_new_prompt_without_container_params_calls_create_version(
+        self, client, mock_rest_client
+    ):
+        """When creating a new prompt without id/description/tags, should call create_prompt_version.
+        Also asserts that version_number from the create response is surfaced."""
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = _make_404_error()
+        mock_rest_client.prompts.create_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        result = client.create_prompt(
+            name="test-prompt",
+            prompt="test template",
+            metadata=None,
+            type=prompt_types.PromptType.MUSTACHE,
+        )
+
+        mock_rest_client.prompts.create_prompt_version.assert_called_once()
+        mock_rest_client.prompts.create_prompt.assert_not_called()
+        assert result.version_number == "v1"
+
+    def test_update_existing_prompt_always_calls_create_version(
+        self, client, mock_rest_client
+    ):
+        """When updating an existing prompt, should always call create_prompt_version regardless of params."""
+        existing_version = _make_mock_version(template="old template")
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = existing_version
+        mock_rest_client.prompts.create_prompt_version.return_value = (
+            _make_mock_version(template="new template")
+        )
+
+        client.create_prompt(
+            name="test-prompt",
+            prompt="new template",
+            metadata=None,
+            type=prompt_types.PromptType.MUSTACHE,
+            id="custom-id",
+            description="Updated description",
+            tags=["updated"],
+        )
+
+        mock_rest_client.prompts.create_prompt_version.assert_called_once()
+        mock_rest_client.prompts.create_prompt.assert_not_called()
+
+    def test_create_new_prompt_with_only_id_calls_create_prompt(
+        self, client, mock_rest_client
+    ):
+        """When creating a new prompt with only id parameter, should call create_prompt."""
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            _make_404_error(),
+            _make_mock_version(),
+        ]
+
+        client.create_prompt(
+            name="test-prompt",
+            prompt="test template",
+            metadata=None,
+            type=prompt_types.PromptType.MUSTACHE,
+            id="custom-id",
+        )
+
+        mock_rest_client.prompts.create_prompt.assert_called_once()
+        mock_rest_client.prompts.create_prompt_version.assert_not_called()
+
+    def test_create_new_prompt_with_only_description_calls_create_prompt(
+        self, client, mock_rest_client
+    ):
+        """When creating a new prompt with only description parameter, should call create_prompt."""
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            _make_404_error(),
+            _make_mock_version(),
+        ]
+
+        client.create_prompt(
+            name="test-prompt",
+            prompt="test template",
+            metadata=None,
+            type=prompt_types.PromptType.MUSTACHE,
+            description="A test prompt",
+        )
+
+        mock_rest_client.prompts.create_prompt.assert_called_once()
+        mock_rest_client.prompts.create_prompt_version.assert_not_called()
+
+    def test_create_new_prompt_with_only_tags_calls_create_prompt(
+        self, client, mock_rest_client
+    ):
+        """When creating a new prompt with only tags parameter, should call create_prompt.
+        Tags-only takes the rebuild branch — assert version_number survives it
+        (regression for the bug where tag injection dropped version_number)."""
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            _make_404_error(),
+            _make_mock_version(tags=None),
+        ]
+
+        result = client.create_prompt(
+            name="test-prompt",
+            prompt="test template",
+            metadata=None,
+            type=prompt_types.PromptType.MUSTACHE,
+            tags=["test"],
+        )
+
+        mock_rest_client.prompts.create_prompt.assert_called_once()
+        mock_rest_client.prompts.create_prompt_version.assert_not_called()
+        assert result.version_number == "v1"
+        assert result.tags == ["test"]
+
+
+class TestInternalCreateMask:
+    """Tests for __internal__create_mask method."""
+
+    def test_create_mask_calls_create_prompt_version_with_mask_type(
+        self, client, mock_rest_client
+    ):
+        expected = _make_mock_version()
+        mock_rest_client.prompts.create_prompt_version.return_value = expected
+
+        result = client._PromptClient__internal_api__create_mask(
+            name="test-prompt",
+            prompt="masked template",
+        )
+
+        mock_rest_client.prompts.create_prompt_version.assert_called_once()
+        call_kwargs = mock_rest_client.prompts.create_prompt_version.call_args[1]
+        assert call_kwargs["name"] == "test-prompt"
+        assert call_kwargs["version"].template == "masked template"
+        assert call_kwargs["version"].version_type == "mask"
+        assert result == expected
+
+    def test_create_mask_passes_all_parameters(self, client, mock_rest_client):
+        mock_rest_client.prompts.create_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        client._PromptClient__internal_api__create_mask(
+            name="test-prompt",
+            prompt="masked template",
+            type=prompt_types.PromptType.JINJA2,
+            metadata={"key": "value"},
+            template_structure="chat",
+            project_name="my-project",
+            change_description="mask for testing",
+        )
+
+        call_kwargs = mock_rest_client.prompts.create_prompt_version.call_args[1]
+        assert call_kwargs["name"] == "test-prompt"
+        assert call_kwargs["template_structure"] == "chat"
+        assert call_kwargs["project_name"] == "my-project"
+        version = call_kwargs["version"]
+        assert version.template == "masked template"
+        assert version.version_type == "mask"
+        assert version.type == prompt_types.PromptType.JINJA2
+        assert version.metadata == {"key": "value"}
+        assert version.change_description == "mask for testing"
+
+    def test_create_mask_defaults(self, client, mock_rest_client):
+        mock_rest_client.prompts.create_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        client._PromptClient__internal_api__create_mask(
+            name="test-prompt",
+            prompt="template",
+        )
+
+        call_kwargs = mock_rest_client.prompts.create_prompt_version.call_args[1]
+        assert call_kwargs["template_structure"] == "text"
+        assert call_kwargs["project_name"] is None
+        version = call_kwargs["version"]
+        assert version.type == prompt_types.PromptType.MUSTACHE
+        assert version.metadata is None
+        assert version.change_description is None
+
+
+class TestGetPromptByVersionSelector:
+    """Tests for the new sequential ``version`` parameter (e.g. ``"v3"``)."""
+
+    @pytest.fixture(autouse=True)
+    def clear_global_cache(self):
+        yield
+        prompt_cache.get_global_cache().clear()
+
+    def test_get_prompt__version_arg__threads_through_as_version_number(
+        self, client, mock_rest_client
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        client.get_prompt(name="my-prompt", version="v3")
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_called_once()
+        call_kwargs = mock_rest_client.prompts.retrieve_prompt_version.call_args[1]
+        assert call_kwargs["name"] == "my-prompt"
+        assert call_kwargs["version_number"] == "v3"
+        assert call_kwargs["commit"] is None
+
+    def test_get_prompt__commit_and_version_both_set__raises_value_error(
+        self, client, mock_rest_client
+    ):
+        with pytest.raises(ValueError, match=r"Provide either `commit` or `version`"):
+            client.get_prompt(name="my-prompt", commit="abc12345", version="v1")
+        mock_rest_client.prompts.retrieve_prompt_version.assert_not_called()
+
+    def test_get_prompt_with_cache__commit_and_version_both_set__raises_value_error(
+        self, client, mock_rest_client
+    ):
+        from opik.api_objects.prompt.text import prompt as text_prompt_module
+
+        with pytest.raises(ValueError, match=r"Provide either `commit` or `version`"):
+            client.get_prompt_with_cache(
+                name="my-prompt",
+                commit="abc12345",
+                project_name=None,
+                template_structure="text",
+                prompt_cls=text_prompt_module.Prompt,
+                version="v1",
+            )
+
+    def test_get_prompt_with_cache__commit_and_environment__raises_value_error(
+        self, client, mock_rest_client
+    ):
+        from opik.api_objects.prompt.text import prompt as text_prompt_module
+
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            client.get_prompt_with_cache(
+                name="my-prompt",
+                commit="abc12345",
+                project_name=None,
+                template_structure="text",
+                prompt_cls=text_prompt_module.Prompt,
+                environment="staging",
+            )
+
+    def test_get_prompt_with_cache__version_and_environment__raises_value_error(
+        self, client, mock_rest_client
+    ):
+        from opik.api_objects.prompt.text import prompt as text_prompt_module
+
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            client.get_prompt_with_cache(
+                name="my-prompt",
+                commit=None,
+                project_name=None,
+                template_structure="text",
+                prompt_cls=text_prompt_module.Prompt,
+                version="v1",
+                environment="staging",
+            )
+
+    def test_get_prompt_with_cache__different_versions__do_not_collide_in_cache(
+        self, client, mock_rest_client
+    ):
+        from opik.api_objects.prompt.text import prompt as text_prompt_module
+
+        v1 = _make_mock_version(template="content for v1")
+        v2 = _make_mock_version(template="content for v2")
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [v1, v2]
+
+        first = client.get_prompt_with_cache(
+            name="my-prompt",
+            commit=None,
+            project_name=None,
+            template_structure="text",
+            prompt_cls=text_prompt_module.Prompt,
+            version="v1",
+        )
+        second = client.get_prompt_with_cache(
+            name="my-prompt",
+            commit=None,
+            project_name=None,
+            template_structure="text",
+            prompt_cls=text_prompt_module.Prompt,
+            version="v2",
+        )
+
+        assert first is not None and first.prompt == "content for v1"
+        assert second is not None and second.prompt == "content for v2"
+        assert mock_rest_client.prompts.retrieve_prompt_version.call_count == 2
+
+    def test_get_prompt_with_cache__commit_and_version_pin__do_not_collide(
+        self, client, mock_rest_client
+    ):
+        from opik.api_objects.prompt.text import prompt as text_prompt_module
+
+        by_commit = _make_mock_version(template="content by commit")
+        by_version = _make_mock_version(template="content by version")
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            by_commit,
+            by_version,
+        ]
+
+        commit_result = client.get_prompt_with_cache(
+            name="my-prompt",
+            commit="abc12345",
+            project_name=None,
+            template_structure="text",
+            prompt_cls=text_prompt_module.Prompt,
+        )
+        version_result = client.get_prompt_with_cache(
+            name="my-prompt",
+            commit=None,
+            project_name=None,
+            template_structure="text",
+            prompt_cls=text_prompt_module.Prompt,
+            version="v3",
+        )
+
+        assert commit_result is not None
+        assert commit_result.prompt == "content by commit"
+        assert version_result is not None
+        assert version_result.prompt == "content by version"
+        assert mock_rest_client.prompts.retrieve_prompt_version.call_count == 2
+
+    def test_get_prompt_with_cache__same_version_twice__second_hits_cache(
+        self, client, mock_rest_client
+    ):
+        from opik.api_objects.prompt.text import prompt as text_prompt_module
+
+        v2 = _make_mock_version(template="v2 content")
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = v2
+
+        first = client.get_prompt_with_cache(
+            name="my-prompt",
+            commit=None,
+            project_name=None,
+            template_structure="text",
+            prompt_cls=text_prompt_module.Prompt,
+            version="v2",
+        )
+        second = client.get_prompt_with_cache(
+            name="my-prompt",
+            commit=None,
+            project_name=None,
+            template_structure="text",
+            prompt_cls=text_prompt_module.Prompt,
+            version="v2",
+        )
+
+        assert first is second
+        assert mock_rest_client.prompts.retrieve_prompt_version.call_count == 1
+
+
+class TestGetPromptWithCacheBypass:
+    """Tests for no_cache parameter in Opik.get_prompt()."""
+
+    @pytest.fixture(autouse=True)
+    def clear_global_cache(self):
+        yield
+        prompt_cache.get_global_cache().clear()
+
+    @pytest.fixture
+    def opik_client(self, mock_rest_client):
+        client = opik_client_module.Opik()
+        client._rest_client = mock_rest_client
+        return client
+
+    @pytest.mark.parametrize(
+        "no_cache, expected_extra_calls",
+        [
+            (False, 0),
+            (True, 1),
+        ],
+        ids=["no_cache_false__uses_cache", "no_cache_true__hits_backend"],
+    )
+    def test_get_prompt__second_call_after_cache_warm__respects_no_cache(
+        self, opik_client, mock_rest_client, no_cache, expected_extra_calls
+    ):
+        version = _make_mock_version()
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = version
+
+        opik_client.get_prompt(name="my-prompt", commit=None, project_name=None)
+        call_count_after_warm = (
+            mock_rest_client.prompts.retrieve_prompt_version.call_count
+        )
+
+        opik_client.get_prompt(
+            name="my-prompt", commit=None, project_name=None, no_cache=no_cache
+        )
+
+        assert (
+            mock_rest_client.prompts.retrieve_prompt_version.call_count
+            == call_count_after_warm + expected_extra_calls
+        )
+
+    def test_get_prompt__no_cache_true__returns_fresh_value_from_backend(
+        self, opik_client, mock_rest_client
+    ):
+        old_version = _make_mock_version(template="old template")
+        new_version = _make_mock_version(template="new template")
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            old_version,
+            new_version,
+        ]
+
+        opik_client.get_prompt(name="my-prompt", commit=None, project_name=None)
+
+        result = opik_client.get_prompt(
+            name="my-prompt", commit=None, project_name=None, no_cache=True
+        )
+
+        assert result is not None
+        assert result.prompt == "new template"
+
+    def test_get_prompt__no_cache_true__backend_returns_none__returns_none(
+        self, opik_client, mock_rest_client
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = _make_404_error()
+
+        result = opik_client.get_prompt(
+            name="missing-prompt", commit=None, project_name=None, no_cache=True
+        )
+
+        assert result is None
+
+
+class TestPromptEnvironment:
+    """Unit tests for the prompt ``environment`` plumbing."""
+
+    @pytest.fixture(autouse=True)
+    def clear_global_cache(self):
+        yield
+        prompt_cache.get_global_cache().clear()
+
+    def test_get_prompt__forwards_environment_to_retrieve(
+        self, client, mock_rest_client
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        client.get_prompt(name="env-prompt", environment="staging")
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_called_once()
+        call_kwargs = mock_rest_client.prompts.retrieve_prompt_version.call_args[1]
+        assert call_kwargs["environment"] == "staging"
+
+    def test_get_prompt__commit_and_environment__raises_value_error(
+        self, client, mock_rest_client
+    ):
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            client.get_prompt(
+                name="env-prompt", commit="abc12345", environment="staging"
+            )
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_not_called()
+
+    def test_get_prompt__version_and_environment__raises_value_error(
+        self, client, mock_rest_client
+    ):
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            client.get_prompt(name="env-prompt", version="v3", environment="staging")
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_not_called()
+
+    def test_get_prompt_with_cache__different_environments__not_cached_together(
+        self, mock_rest_client
+    ):
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        staging_version = prompt_version_detail.PromptVersionDetail(
+            id="staging-id",
+            prompt_id="prompt-id",
+            template="staging template",
+            type="mustache",
+            metadata=None,
+            commit="aaa",
+            template_structure="text",
+            environments=["staging"],
+        )
+        production_version = prompt_version_detail.PromptVersionDetail(
+            id="prod-id",
+            prompt_id="prompt-id",
+            template="prod template",
+            type="mustache",
+            metadata=None,
+            commit="bbb",
+            template_structure="text",
+            environments=["production"],
+        )
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = [
+            staging_version,
+            production_version,
+        ]
+
+        staging = opik_client.get_prompt(name="env-prompt", environment="staging")
+        production = opik_client.get_prompt(name="env-prompt", environment="production")
+
+        assert staging is not None and production is not None
+        assert staging.commit == "aaa"
+        assert production.commit == "bbb"
+        assert mock_rest_client.prompts.retrieve_prompt_version.call_count == 2
+
+    def test_set_prompt_environments__sets_via_latest_version(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        opik_client.set_prompt_environments("env-prompt", ["staging"])
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_called_once()
+        retrieve_kwargs = mock_rest_client.prompts.retrieve_prompt_version.call_args[1]
+        assert retrieve_kwargs["name"] == "env-prompt"
+        mock_rest_client.prompts.set_prompt_version_environment.assert_called_once_with(
+            version_id="version-id",
+            environments=["staging"],
+        )
+
+    def test_set_prompt_environments__empty_list_clears(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        opik_client.set_prompt_environments("env-prompt", [])
+
+        mock_rest_client.prompts.set_prompt_version_environment.assert_called_once_with(
+            version_id="version-id",
+            environments=[],
+        )
+
+    def test_set_prompt_environments__multiple_envs_in_one_call(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        opik_client.set_prompt_environments("env-prompt", ["staging", "production"])
+
+        mock_rest_client.prompts.set_prompt_version_environment.assert_called_once_with(
+            version_id="version-id",
+            environments=["staging", "production"],
+        )
+
+    def test_set_prompt_environments__deduplicates_input(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        opik_client.set_prompt_environments(
+            "env-prompt", ["staging", "staging", "production"]
+        )
+
+        mock_rest_client.prompts.set_prompt_version_environment.assert_called_once_with(
+            version_id="version-id",
+            environments=["staging", "production"],
+        )
+
+    def test_set_prompt_environments__forwards_project_name(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        opik_client.set_prompt_environments(
+            "env-prompt", ["staging"], project_name="my-project"
+        )
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_called_once()
+        retrieve_kwargs = mock_rest_client.prompts.retrieve_prompt_version.call_args[1]
+        assert retrieve_kwargs["name"] == "env-prompt"
+        assert retrieve_kwargs["project_name"] == "my-project"
+
+    def test_set_prompt_environments__forwards_version(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        opik_client.set_prompt_environments("env-prompt", ["staging"], version="v3")
+
+        mock_rest_client.prompts.retrieve_prompt_version.assert_called_once()
+        retrieve_kwargs = mock_rest_client.prompts.retrieve_prompt_version.call_args[1]
+        assert retrieve_kwargs["name"] == "env-prompt"
+        # The SDK forwards ``version`` as the wire-level ``version_number`` field —
+        # ``commit`` is no longer part of this method's surface.
+        assert retrieve_kwargs["version_number"] == "v3"
+        assert "commit" not in retrieve_kwargs
+        mock_rest_client.prompts.set_prompt_version_environment.assert_called_once_with(
+            version_id="version-id",
+            environments=["staging"],
+        )
+
+    def test_set_prompt_environments__prompt_not_found__raises_prompt_not_found(
+        self, mock_rest_client
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = (
+            rest_api_core.ApiError(status_code=404, body=None)
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        with pytest.raises(exceptions.PromptNotFoundError, match="missing-prompt"):
+            opik_client.set_prompt_environments("missing-prompt", ["staging"])
+
+        mock_rest_client.prompts.set_prompt_version_environment.assert_not_called()
+
+    def test_set_prompt_environments__version_not_found__raises_prompt_not_found_with_version(
+        self, mock_rest_client
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.side_effect = (
+            rest_api_core.ApiError(status_code=404, body=None)
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        with pytest.raises(exceptions.PromptNotFoundError, match="v7"):
+            opik_client.set_prompt_environments("env-prompt", ["staging"], version="v7")
+
+        mock_rest_client.prompts.set_prompt_version_environment.assert_not_called()
+
+    # The backend reports an unknown environment as 404 or 409 from the
+    # workspace-registry check; both must surface as EnvironmentNotFoundError.
+    @pytest.mark.parametrize("status_code", [404, 409])
+    def test_set_prompt_environments__environment_not_found__raises_environment_not_found(
+        self, mock_rest_client, status_code
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+        mock_rest_client.prompts.set_prompt_version_environment.side_effect = (
+            rest_api_core.ApiError(status_code=status_code, body=None)
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        with pytest.raises(exceptions.EnvironmentNotFoundError, match="unknown-env"):
+            opik_client.set_prompt_environments("env-prompt", ["unknown-env"])
+
+    def test_set_prompt_environments__other_api_error__bubbles_up(
+        self, mock_rest_client
+    ):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+        mock_rest_client.prompts.set_prompt_version_environment.side_effect = (
+            rest_api_core.ApiError(status_code=500, body=None)
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+
+        with pytest.raises(rest_api_core.ApiError):
+            opik_client.set_prompt_environments("env-prompt", ["staging"])
+
+    def test_set_prompt_environments__invalidates_cache(self, mock_rest_client):
+        mock_rest_client.prompts.retrieve_prompt_version.return_value = (
+            _make_mock_version()
+        )
+
+        opik_client = opik_client_module.Opik()
+        opik_client._rest_client = mock_rest_client
+        resolved_project = opik_client._resolve_project_name(None)
+
+        cache = prompt_cache.get_global_cache()
+        cache.clear()
+        try:
+            sentinel = mock.MagicMock()
+            cache.get_or_fetch(
+                key=("env-prompt", None, resolved_project, "text", "staging"),
+                fetch_fn=lambda: sentinel,
+                ttl_seconds=None,
+            )
+            cache.get_or_fetch(
+                key=("other-prompt", None, resolved_project, "text", "staging"),
+                fetch_fn=lambda: sentinel,
+                ttl_seconds=None,
+            )
+
+            opik_client.set_prompt_environments("env-prompt", ["production"])
+
+            assert (
+                cache.get(("env-prompt", None, resolved_project, "text", "staging"))
+                is None
+            ), "stale entry for the updated prompt should be evicted"
+            assert (
+                cache.get(("other-prompt", None, resolved_project, "text", "staging"))
+                is sentinel
+            ), "entries for unrelated prompts must not be evicted"
+        finally:
+            cache.clear()

sdks/python/tests/unit/api_objects/prompt/test_prompt_sync.py

@@ -0,0 +1,180 @@
+"""Tests for Prompt and ChatPrompt resilience when backend sync fails."""
+
+from unittest.mock import patch, MagicMock
+
+import httpx
+import pytest
+
+from opik.rest_api.core import ApiError
+from opik.api_objects.prompt.text.prompt import Prompt
+from opik.api_objects.prompt.chat.chat_prompt import ChatPrompt
+
+
+_API_ERRORS = [
+    pytest.param(ApiError(status_code=500, body="Internal Server Error"), id="api_500"),
+    pytest.param(ApiError(status_code=503, body="Service Unavailable"), id="api_503"),
+    pytest.param(httpx.ConnectError("Connection refused"), id="connect_error"),
+    pytest.param(httpx.TimeoutException("Request timed out"), id="timeout"),
+]
+
+
+def _mock_opik_client():
+    """Patch the opik_client.get_client_cached() used inside sync_with_backend."""
+    return patch(
+        "opik.api_objects.opik_client.get_client_cached",
+        return_value=MagicMock(),
+    )
+
+
+class TestPromptSyncFailure:
+    """Prompt should be created locally even when sync_with_backend fails with API errors."""
+
+    @pytest.mark.parametrize("error", _API_ERRORS)
+    def test_prompt_created_despite_backend_failure(self, error):
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                side_effect=error,
+            ),
+        ):
+            prompt = Prompt(name="test-prompt", prompt="Hello {{name}}")
+
+        assert prompt.name == "test-prompt"
+        assert prompt.prompt == "Hello {{name}}"
+        assert prompt.commit is None
+        assert prompt.synced is False
+        assert prompt.format(name="World") == "Hello World"
+
+    def test_sync_returns_false_on_api_error(self):
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                side_effect=ApiError(status_code=500),
+            ),
+        ):
+            prompt = Prompt(name="test-prompt", prompt="Hello {{name}}")
+            result = prompt.sync_with_backend()
+
+        assert result is False
+        assert prompt.synced is False
+
+    def test_sync_returns_true_on_success(self):
+        mock_version = MagicMock()
+        mock_version.commit = "abc123"
+        mock_version.prompt_id = "pid"
+        mock_version.id = "vid"
+        mock_version.change_description = None
+        mock_version.tags = []
+
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                return_value=mock_version,
+            ),
+        ):
+            prompt = Prompt(name="test-prompt", prompt="Hello {{name}}")
+
+        assert prompt.commit == "abc123"
+        assert prompt.synced is True
+
+        # Re-sync also succeeds
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                return_value=mock_version,
+            ),
+        ):
+            assert prompt.sync_with_backend() is True
+            assert prompt.synced is True
+
+    def test_non_api_error_propagates(self):
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                side_effect=ValueError("bad value"),
+            ),
+        ):
+            with pytest.raises(ValueError, match="bad value"):
+                Prompt(name="test-prompt", prompt="Hello {{name}}")
+
+
+class TestChatPromptSyncFailure:
+    """ChatPrompt should be created locally even when sync_with_backend fails with API errors."""
+
+    MESSAGES = [{"role": "user", "content": "Hello {{name}}"}]
+
+    @pytest.mark.parametrize("error", _API_ERRORS)
+    def test_chat_prompt_created_despite_backend_failure(self, error):
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                side_effect=error,
+            ),
+        ):
+            prompt = ChatPrompt(name="test-chat", messages=self.MESSAGES)
+
+        assert prompt.name == "test-chat"
+        assert prompt.template == self.MESSAGES
+        assert prompt.commit is None
+        assert prompt.synced is False
+
+    def test_sync_returns_false_on_api_error(self):
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                side_effect=ApiError(status_code=500),
+            ),
+        ):
+            prompt = ChatPrompt(name="test-chat", messages=self.MESSAGES)
+            result = prompt.sync_with_backend()
+
+        assert result is False
+        assert prompt.synced is False
+
+    def test_sync_returns_true_on_success(self):
+        mock_version = MagicMock()
+        mock_version.commit = "abc123"
+        mock_version.prompt_id = "pid"
+        mock_version.id = "vid"
+        mock_version.change_description = None
+        mock_version.tags = []
+
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                return_value=mock_version,
+            ),
+        ):
+            prompt = ChatPrompt(name="test-chat", messages=self.MESSAGES)
+
+        assert prompt.commit == "abc123"
+        assert prompt.synced is True
+
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                return_value=mock_version,
+            ),
+        ):
+            assert prompt.sync_with_backend() is True
+            assert prompt.synced is True
+
+    def test_non_api_error_propagates(self):
+        with (
+            _mock_opik_client(),
+            patch(
+                "opik.api_objects.prompt.client.PromptClient.create_prompt",
+                side_effect=ValueError("bad value"),
+            ),
+        ):
+            with pytest.raises(ValueError, match="bad value"):
+                ChatPrompt(name="test-chat", messages=self.MESSAGES)

sdks/python/tests/unit/api_objects/prompt/test_prompt_template.py

@@ -0,0 +1,138 @@
+import pytest
+from opik.api_objects.prompt import PromptTemplate, PromptType
+from opik import exceptions
+
+
+def test_prompt__format__happyflow():
+    PROMPT_TEMPLATE = "Hi, my name is {{name}}, I live in {{city}}."
+
+    tested = PromptTemplate(PROMPT_TEMPLATE)
+
+    result = tested.format(name="Harry", city="London")
+    assert result == "Hi, my name is Harry, I live in London."
+
+
+def test_prompt__format__one_placeholder_used_multiple_times():
+    PROMPT_TEMPLATE = (
+        "Hi, my name is {{name}}, I live in {{city}}. I repeat, my name is {{name}}"
+    )
+
+    tested = PromptTemplate(PROMPT_TEMPLATE)
+
+    result = tested.format(name="Harry", city="London")
+    assert (
+        result == "Hi, my name is Harry, I live in London. I repeat, my name is Harry"
+    )
+
+
+def test_prompt__format__passed_arguments_that_are_not_in_template__error_raised_with_correct_report_info():
+    PROMPT_TEMPLATE = "Hi, my name is {{name}}, I live in {{city}}."
+
+    tested = PromptTemplate(PROMPT_TEMPLATE)
+
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format(name="Harry", city="London", nemesis_name="Voldemort")
+
+    assert exc_info.value.format_arguments == set(["name", "city", "nemesis_name"])
+    assert exc_info.value.prompt_placeholders == set(
+        [
+            "name",
+            "city",
+        ]
+    )
+    assert exc_info.value.symmetric_difference == set(["nemesis_name"])
+
+
+def test_prompt__format__some_placeholders_dont_have_corresponding_format_arguments__error_raised_with_correct_report_info():
+    PROMPT_TEMPLATE = "Hi, my name is {{name}}, I live in {{city}}."
+
+    tested = PromptTemplate(PROMPT_TEMPLATE)
+
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format(name="Harry")
+
+    assert exc_info.value.format_arguments == set(["name"])
+    assert exc_info.value.prompt_placeholders == set(["name", "city"])
+    assert exc_info.value.symmetric_difference == set(["city"])
+
+
+def test_prompt__format__some_placeholders_dont_have_corresponding_format_arguments_AND_there_are_format_arguments_that_are_not_in_the_template__error_raised_with_correct_report_info():
+    PROMPT_TEMPLATE = "Hi, my name is {{name}}, I live in {{city}}."
+
+    tested = PromptTemplate(PROMPT_TEMPLATE)
+
+    with pytest.raises(
+        exceptions.PromptPlaceholdersDontMatchFormatArguments
+    ) as exc_info:
+        tested.format(name="Harry", nemesis_name="Voldemort")
+
+    assert exc_info.value.format_arguments == set(["name", "nemesis_name"])
+    assert exc_info.value.prompt_placeholders == set(["name", "city"])
+    assert exc_info.value.symmetric_difference == set(["city", "nemesis_name"])
+
+
+def test_prompt__format_jinja2__happyflow():
+    PROMPT_TEMPLATE = "Hi, my name is {{ name }}, I live in {{ city }}."
+
+    tested = PromptTemplate(PROMPT_TEMPLATE, type=PromptType.JINJA2)
+
+    result = tested.format(name="Harry", city="London")
+    assert result == "Hi, my name is Harry, I live in London."
+
+
+def test_prompt__format_jinja2__with_control_flow():
+    PROMPT_TEMPLATE = """
+    {% if is_wizard %}
+    {{ name }} is a wizard who lives in {{ city }}.
+    {% else %}
+    {{ name }} is a muggle who lives in {{ city }}.
+    {% endif %}
+    """
+
+    tested = PromptTemplate(PROMPT_TEMPLATE, type=PromptType.JINJA2)
+
+    wizard_result = tested.format(name="Harry", city="London", is_wizard=True)
+    assert "Harry is a wizard who lives in London." in wizard_result.strip()
+
+    muggle_result = tested.format(name="Dudley", city="Surrey", is_wizard=False)
+    assert "Dudley is a muggle who lives in Surrey." in muggle_result.strip()
+
+
+def test_prompt__format_jinja2__with_loops():
+    PROMPT_TEMPLATE = """
+    {{ name }}'s friends are:
+    {% for friend in friends %}
+    - {{ friend }}
+    {% endfor %}
+    """
+
+    tested = PromptTemplate(PROMPT_TEMPLATE, type=PromptType.JINJA2)
+
+    result = tested.format(name="Harry", friends=["Ron", "Hermione", "Neville"])
+
+    assert "Harry's friends are:" in result
+    assert "- Ron" in result
+    assert "- Hermione" in result
+    assert "- Neville" in result
+
+
+def test_prompt__format_jinja2__with_filters():
+    PROMPT_TEMPLATE = "{{ name | upper }} lives in {{ city | lower }}."
+
+    tested = PromptTemplate(PROMPT_TEMPLATE, type=PromptType.JINJA2)
+
+    result = tested.format(name="Harry", city="LONDON")
+    assert result == "HARRY lives in london."
+
+
+def test_prompt__format__none_values_render_as_empty_strings() -> None:
+    PROMPT_TEMPLATE = "Primary: {{primary}} Secondary: {{secondary}}"
+
+    tested = PromptTemplate(PROMPT_TEMPLATE)
+
+    result = tested.format(primary="cat", secondary=None)
+    assert result == "Primary: cat Secondary: "

sdks/python/tests/unit/api_objects/span/test_span_data.py

@@ -0,0 +1,41 @@
+import datetime
+
+from opik.api_objects.span.span_data import SpanData
+from ....testlib import assert_equal, ANY_BUT_NONE
+
+
+def test_span_data__as_start_parameters__includes_parent_span_id():
+    span_data = SpanData(
+        trace_id="trace-1",
+        id="span-1",
+        parent_span_id="parent-1",
+        start_time=datetime.datetime.now(),
+        project_name="test",
+        name="span-name",
+        source="sdk",
+    )
+
+    expected_parameters = {
+        "id": "span-1",
+        "start_time": ANY_BUT_NONE,
+        "project_name": "test",
+        "trace_id": "trace-1",
+        "parent_span_id": "parent-1",
+        "name": "span-name",
+        "source": "sdk",
+    }
+
+    assert_equal(expected_parameters, span_data.as_start_parameters)
+
+
+def test_span_data__as_start_parameters__excludes_parent_span_id_when_none():
+    span_data = SpanData(
+        trace_id="trace-1",
+        id="span-1",
+        parent_span_id=None,
+        start_time=datetime.datetime.now(),
+        project_name="test",
+    )
+
+    params = span_data.as_start_parameters
+    assert "parent_span_id" not in params

sdks/python/tests/unit/api_objects/test_connection_resource_manager.py

@@ -0,0 +1,395 @@
+import threading
+import time
+from typing import List, Optional, Tuple
+from unittest import mock
+
+from opik import config as opik_config
+from opik.api_objects import connection_resources
+
+
+class FakeBundle:
+    """Stand-in for SharedConnectionResourcesBundle that records close/flush calls."""
+
+    def __init__(self) -> None:
+        self.flush_timeout: Optional[int] = None
+        self.close_calls: List[Tuple[Optional[int], bool]] = []
+        self.flush_calls: List[Optional[int]] = []
+
+    @property
+    def closed(self) -> bool:
+        return len(self.close_calls) > 0
+
+    def close(self, timeout: Optional[int], *, flush: bool) -> None:
+        self.close_calls.append((timeout, flush))
+
+    def flush(self, timeout: Optional[int]) -> None:
+        self.flush_calls.append(timeout)
+
+
+def _manager():
+    """Returns (manager, created) where ``created`` accumulates built bundles.
+
+    A fake builder keeps these tests focused on the ref-counting lifecycle
+    without spinning up real transport threads.
+    """
+    created: List[FakeBundle] = []
+
+    def builder(config, *, use_batching):
+        bundle = FakeBundle()
+        created.append(bundle)
+        return bundle
+
+    return connection_resources.ConnectionResourceManager(builder=builder), created
+
+
+def _config(workspace: str = "default") -> opik_config.OpikConfig:
+    return opik_config.OpikConfig(workspace=workspace)
+
+
+def test_acquire__same_config__reuses_single_bundle():
+    manager, created = _manager()
+    config = _config()
+
+    lease_a = manager.acquire(config, use_batching=True)
+    lease_b = manager.acquire(config, use_batching=True)
+
+    assert len(created) == 1
+    assert lease_a.resources is lease_b.resources is created[0]
+    assert manager.reference_count(config, use_batching=True) == 2
+    assert manager.active_connection_count() == 1
+
+
+def test_acquire__distinct_configs__builds_separate_bundles():
+    manager, created = _manager()
+    config_a = _config("ws-a")
+    config_b = _config("ws-b")
+
+    lease_a = manager.acquire(config_a, use_batching=True)
+    lease_b = manager.acquire(config_b, use_batching=True)
+
+    assert len(created) == 2
+    assert lease_a.resources is not lease_b.resources
+    assert manager.active_connection_count() == 2
+
+
+def test_acquire__differing_use_batching__builds_separate_bundles():
+    manager, created = _manager()
+    config = _config()
+
+    manager.acquire(config, use_batching=True)
+    manager.acquire(config, use_batching=False)
+
+    assert len(created) == 2
+    assert manager.active_connection_count() == 2
+
+
+def test_release__not_last_reference__keeps_bundle_open():
+    manager, created = _manager()
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)
+    manager.acquire(config, use_batching=True)
+
+    lease_a.release(timeout=None, close_on_zero=True)
+
+    assert not created[0].closed
+    # A durable release (flush=True, the default) still drains the shared queue
+    # so this handle's data is persisted while the bundle stays alive.
+    assert created[0].flush_calls == [None]
+    assert manager.reference_count(config, use_batching=True) == 1
+
+
+def test_release__not_last_reference_with_flush__drains_shared_queue_without_closing():
+    # Durability contract under sharing: end(flush=True) on a handle that shares
+    # its bundle must flush the shared queue now — otherwise a co-located
+    # handle's later flush=False teardown could discard this handle's data.
+    manager, created = _manager()
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)
+    manager.acquire(config, use_batching=True)
+
+    lease_a.release(timeout=3, flush=True, close_on_zero=True)
+
+    assert created[0].flush_calls == [3]
+    assert not created[0].closed
+    assert manager.reference_count(config, use_batching=True) == 1
+
+
+def test_release__not_last_reference_flush_false__neither_flushes_nor_closes():
+    manager, created = _manager()
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)
+    manager.acquire(config, use_batching=True)
+
+    lease_a.release(timeout=None, flush=False, close_on_zero=True)
+
+    assert created[0].flush_calls == []
+    assert not created[0].closed
+
+
+def test_release__not_last_reference_gc_path__never_flushes():
+    # A GC finalizer (close_on_zero=False) must never do network I/O, even though
+    # it releases with flush=True by default.
+    manager, created = _manager()
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)
+    manager.acquire(config, use_batching=True)
+
+    lease_a.release(timeout=None, close_on_zero=False)
+
+    assert created[0].flush_calls == []
+    assert not created[0].closed
+    assert manager.reference_count(config, use_batching=True) == 1
+
+
+def test_release__concurrent_durable_and_teardown__no_close_during_shared_flush():
+    # Regression: a shared end(flush=True) pre-flushes while still holding its
+    # reference, so a concurrent last-release (flush=False) cannot evict + close
+    # the bundle and clear the queue mid-flush. The event forces the teardown to
+    # race the in-flight shared flush; on the buggy (decrement-then-flush) order
+    # the close would run while in_flush is True.
+    flush_started = threading.Event()
+
+    class RaceDetectBundle(FakeBundle):
+        def __init__(self) -> None:
+            super().__init__()
+            self.in_flush = False
+            self.closed_during_flush = False
+
+        def flush(self, timeout: Optional[int]) -> None:
+            self.in_flush = True
+            flush_started.set()
+            time.sleep(0.1)  # window in which a racing close must not run
+            self.in_flush = False
+            super().flush(timeout)
+
+        def close(self, timeout: Optional[int], *, flush: bool) -> None:
+            if self.in_flush:
+                self.closed_during_flush = True
+            super().close(timeout, flush=flush)
+
+    bundle = RaceDetectBundle()
+    manager = connection_resources.ConnectionResourceManager(
+        builder=lambda config, *, use_batching: bundle
+    )
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)  # durable holder
+    lease_b = manager.acquire(config, use_batching=True)  # fire-and-forget holder
+
+    def durable_release() -> None:
+        lease_a.release(timeout=None, flush=True, close_on_zero=True)
+
+    def teardown_release() -> None:
+        flush_started.wait(timeout=2)  # release into the in-flight shared flush
+        lease_b.release(timeout=None, flush=False, close_on_zero=True)
+
+    threads = [
+        threading.Thread(target=durable_release),
+        threading.Thread(target=teardown_release),
+    ]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert not bundle.closed_during_flush  # close never ran mid-flush
+    assert bundle.closed  # bundle still torn down once, after the flush
+    assert bundle.flush_calls == [None]  # the durable holder drained the queue
+    assert manager.active_connection_count() == 0
+
+
+def test_release__last_reference__closes_and_evicts():
+    manager, created = _manager()
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)
+    lease_b = manager.acquire(config, use_batching=True)
+
+    lease_a.release(timeout=None, close_on_zero=True)
+    lease_b.release(timeout=None, close_on_zero=True)
+
+    assert created[0].close_calls == [(None, True)]
+    assert manager.reference_count(config, use_batching=True) == 0
+    assert manager.active_connection_count() == 0
+
+
+def test_release__last_reference__forwards_timeout_and_flush():
+    manager, created = _manager()
+    lease = manager.acquire(_config(), use_batching=True)
+
+    lease.release(timeout=7, flush=False, close_on_zero=True)
+
+    assert created[0].close_calls == [(7, False)]
+
+
+def test_release__not_close_on_zero__last_reference__decrements_without_closing():
+    # The GC-finalizer path: dropping the last reference must only decrement.
+    # Closing (thread joins, network flush) is never safe inside garbage
+    # collection, so the bundle is left cached instead.
+    manager, created = _manager()
+    config = _config()
+    lease = manager.acquire(config, use_batching=True)
+
+    lease.release(timeout=None, close_on_zero=False)
+
+    assert not created[0].closed
+    assert manager.reference_count(config, use_batching=True) == 0
+    assert manager.active_connection_count() == 1  # still cached, not evicted
+
+
+def test_acquire__after_gc_release__reuses_cached_bundle():
+    # A bundle left cached by a close_on_zero=False release is reused by the
+    # next same-identity acquire rather than rebuilt.
+    manager, created = _manager()
+    config = _config()
+    lease = manager.acquire(config, use_batching=True)
+    lease.release(timeout=None, close_on_zero=False)
+
+    lease_again = manager.acquire(config, use_batching=True)
+
+    assert len(created) == 1
+    assert lease_again.resources is created[0]
+    assert manager.reference_count(config, use_batching=True) == 1
+
+
+def test_close_all__after_gc_release__disposes_cached_bundle():
+    # Whatever a GC release leaves cached is still disposed at process exit.
+    manager, created = _manager()
+    config = _config()
+    lease = manager.acquire(config, use_batching=True)
+    lease.release(timeout=None, close_on_zero=False)
+
+    manager.close_all()
+
+    assert created[0].close_calls == [(None, True)]
+    assert manager.active_connection_count() == 0
+
+
+def test_release__called_twice__decrements_once():
+    manager, created = _manager()
+    config = _config()
+    lease_a = manager.acquire(config, use_batching=True)
+    lease_b = manager.acquire(config, use_batching=True)
+
+    lease_a.release(timeout=None, close_on_zero=True)
+    lease_a.release(
+        timeout=None, close_on_zero=True
+    )  # idempotent — must not decrement again
+
+    assert not created[0].closed
+    assert manager.reference_count(config, use_batching=True) == 1
+
+    lease_b.release(timeout=None, close_on_zero=True)
+    assert created[0].close_calls == [(None, True)]
+    assert manager.reference_count(config, use_batching=True) == 0
+
+
+def test_acquire__after_full_release__builds_fresh_bundle():
+    manager, created = _manager()
+    config = _config()
+    lease = manager.acquire(config, use_batching=True)
+    lease.release(timeout=None, close_on_zero=True)
+
+    lease_again = manager.acquire(config, use_batching=True)
+
+    assert len(created) == 2
+    assert lease_again.resources is created[1]
+
+
+def test_close_all__cached_bundles__closed_and_cleared():
+    manager, created = _manager()
+    manager.acquire(_config("ws-a"), use_batching=True)
+    manager.acquire(_config("ws-b"), use_batching=True)
+
+    manager.close_all()
+
+    assert manager.active_connection_count() == 0
+    # Each bundle is closed with its own configured flush timeout (None here).
+    assert all(bundle.close_calls == [(None, True)] for bundle in created)
+
+
+def test_acquire_release__concurrent_same_config__preserves_invariants():
+    manager, created = _manager()
+    config = _config()
+
+    thread_count = 16
+    iterations = 50
+    barrier = threading.Barrier(thread_count)
+    errors: List[str] = []
+
+    def worker() -> None:
+        barrier.wait()
+        for _ in range(iterations):
+            lease = manager.acquire(config, use_batching=True)
+            # While a reference is held, the bundle must never be torn down —
+            # the manager must not hand out a closing bundle.
+            if lease.resources.closed:
+                errors.append("acquired a bundle that was already closed")
+            lease.release(timeout=None, close_on_zero=True)
+
+    threads = [threading.Thread(target=worker) for _ in range(thread_count)]
+    for thread in threads:
+        thread.start()
+    for thread in threads:
+        thread.join()
+
+    assert errors == []
+    assert manager.reference_count(config, use_batching=True) == 0
+    assert manager.active_connection_count() == 0
+    # Every bundle that was ever built must have been closed exactly once.
+    for bundle in created:
+        assert len(bundle.close_calls) == 1
+
+
+def _bundle_with_mock_transport(flush_timeout=None):
+    streamer = mock.Mock()
+    file_upload_manager = mock.Mock()
+    httpx_client = mock.Mock()
+    bundle = connection_resources.SharedConnectionResourcesBundle(
+        httpx_client=httpx_client,
+        rest_client=mock.Mock(),
+        message_processor=mock.Mock(),
+        file_upload_manager=file_upload_manager,
+        replay_manager=mock.Mock(),
+        streamer=streamer,
+        flush_timeout=flush_timeout,
+    )
+    return bundle, streamer, file_upload_manager, httpx_client
+
+
+def test_bundle_close__flush_true__drains_streamer_and_upload_pool():
+    bundle, streamer, file_upload_manager, httpx_client = _bundle_with_mock_transport()
+
+    bundle.close(5, flush=True)
+
+    # flush=True finalizes pending data: the streamer drains its queue and
+    # flushes uploads, and the upload pool waits for in-flight uploads.
+    streamer.close.assert_called_once_with(5, flush=True)
+    file_upload_manager.close.assert_called_once_with(wait=True)
+    # On a durable close the replay thread is joined and uploads drained, so the
+    # httpx pool is released too — eviction leaks no sockets.
+    httpx_client.close.assert_called_once_with()
+
+
+def test_bundle_close__flush_false__stops_without_waiting():
+    bundle, streamer, file_upload_manager, httpx_client = _bundle_with_mock_transport()
+
+    bundle.close(None, flush=False)
+
+    streamer.close.assert_called_once_with(None, flush=False)
+    file_upload_manager.close.assert_called_once_with(wait=False)
+    # flush=False is fire-and-forget: the streamer leaves daemon threads to
+    # finish in-flight requests, so the shared httpx pool must NOT be closed here
+    # (closing it would race those requests). It's released at GC / process exit.
+    httpx_client.close.assert_not_called()
+
+
+def test_bundle_flush__drains_streamer_without_closing():
+    bundle, streamer, file_upload_manager, httpx_client = _bundle_with_mock_transport()
+
+    bundle.flush(4)
+
+    # flush() drains the queue but must not tear anything down — the bundle is
+    # still shared by other handles.
+    streamer.flush.assert_called_once_with(4)
+    streamer.close.assert_not_called()
+    file_upload_manager.close.assert_not_called()
+    httpx_client.close.assert_not_called()

sdks/python/tests/unit/api_objects/test_data_helpers.py

@@ -0,0 +1,75 @@
+from opik.api_objects import data_helpers
+
+
+def test_merge_tags_both_none():
+    """Test merge_tags with both inputs None."""
+    result = data_helpers.merge_tags(None, None)
+    assert result is None
+
+
+def test_merge_tags_existing_none():
+    """Test merge_tags with existing tags None."""
+    result = data_helpers.merge_tags(None, ["new_tag"])
+    assert result == ["new_tag"]
+
+
+def test_merge_tags_new_none():
+    """Test merge_tags with new tags None."""
+    result = data_helpers.merge_tags(["existing_tag"], None)
+    assert result == ["existing_tag"]
+
+
+def test_merge_tags_no_duplicates():
+    """Test merge_tags with no duplicates."""
+    result = data_helpers.merge_tags(["tag1"], ["tag2", "tag3"])
+    assert result == ["tag1", "tag2", "tag3"]
+
+
+def test_merge_tags_with_duplicates():
+    """Test merge_tags with duplicates."""
+    result = data_helpers.merge_tags(["tag1", "tag2"], ["tag2", "tag3"])
+    assert result == ["tag1", "tag2", "tag3"]
+
+
+def test_merge_tags_empty_lists():
+    """Test merge_tags with empty lists."""
+    result = data_helpers.merge_tags([], [])
+    assert result is None
+
+
+def test_merge_metadata_both_none():
+    """Test merge_metadata with both inputs None."""
+    result = data_helpers.merge_metadata(None, None)
+    assert result is None
+
+
+def test_merge_metadata_existing_none():
+    """Test merge_metadata with existing metadata None."""
+    result = data_helpers.merge_metadata(None, {"key": "value"})
+    assert result == {"key": "value"}
+
+
+def test_merge_metadata_new_none():
+    """Test merge_metadata with new metadata None."""
+    result = data_helpers.merge_metadata({"key": "value"}, None)
+    assert result == {"key": "value"}
+
+
+def test_merge_metadata_no_conflicts():
+    """Test merge_metadata with no key conflicts."""
+    result = data_helpers.merge_metadata({"key1": "value1"}, {"key2": "value2"})
+    assert result == {"key1": "value1", "key2": "value2"}
+
+
+def test_merge_metadata_with_conflicts():
+    """Test merge_metadata with key conflicts (new values win)."""
+    result = data_helpers.merge_metadata(
+        {"key": "old_value", "other": "kept"}, {"key": "new_value"}
+    )
+    assert result == {"key": "new_value", "other": "kept"}
+
+
+def test_merge_metadata_empty_dicts():
+    """Test merge_metadata with empty dictionaries."""
+    result = data_helpers.merge_metadata({}, {})
+    assert result is None

sdks/python/tests/unit/api_objects/test_environment.py

@@ -0,0 +1,335 @@
+"""Unit tests for the SDK ``environment`` plumbing.
+
+Covers:
+- ``Opik.trace(environment=...)`` is the only entry point that accepts an
+  explicit environment; the value flows into the emitted ``CreateTraceMessage``.
+- Spans (``Trace.span``, ``Span.span``, ``@track``-created spans) inherit the
+  parent trace's environment unconditionally.
+- A nested ``@track(environment=...)`` whose value differs from the enclosing
+  trace's environment logs a warning and is ignored — the parent trace's value
+  wins (mirrors how mismatched ``project_name`` is handled).
+- Backwards compat: existing call sites without ``environment`` still work.
+"""
+
+from unittest.mock import MagicMock
+
+import opik
+from opik.api_objects import opik_client
+from opik.message_processing import messages
+from opik import dict_utils
+
+
+def _capture_messages(client: opik_client.Opik) -> MagicMock:
+    mock_streamer = MagicMock()
+    client._streamer = mock_streamer
+    return mock_streamer
+
+
+def _create_trace_messages(streamer: MagicMock):
+    return [
+        c.args[0]
+        for c in streamer.put.call_args_list
+        if isinstance(c.args[0], messages.CreateTraceMessage)
+    ]
+
+
+def _create_span_messages(streamer: MagicMock):
+    return [
+        c.args[0]
+        for c in streamer.put.call_args_list
+        if isinstance(c.args[0], messages.CreateSpanMessage)
+    ]
+
+
+def test_opik_client__no_environment_set__messages_have_none_environment_and_payload_omits_it():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    client.trace(name="t")
+    client.span(name="s")
+
+    trace_msg = _create_trace_messages(streamer)[0]
+    span_msg = _create_span_messages(streamer)[0]
+
+    assert trace_msg.environment is None
+    assert span_msg.environment is None
+
+    cleaned = dict_utils.remove_none_from_dict(trace_msg.as_payload_dict())
+    assert "environment" not in cleaned
+
+
+def test_opik_client_trace__environment_kwarg_propagates_to_create_trace_message():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    client.trace(name="t", environment="staging")
+
+    assert _create_trace_messages(streamer)[0].environment == "staging"
+
+
+def test_opik_client_trace__span_inherits_trace_environment():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    trace = client.trace(name="t", environment="staging")
+    trace.span(name="s")
+
+    assert _create_span_messages(streamer)[0].environment == "staging"
+
+
+def test_opik_client_trace__nested_span_inherits_environment():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    trace = client.trace(name="t", environment="staging")
+    span = trace.span(name="parent-span")
+    span.span(name="child-span")
+
+    span_msgs = _create_span_messages(streamer)
+    assert len(span_msgs) == 2
+    assert all(m.environment == "staging" for m in span_msgs)
+
+
+def test_track_decorator__environment_propagates_to_root_trace_and_spans():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+    opik_client.set_global_client(client)
+
+    @opik.track(environment="staging")
+    def my_fn():
+        return 42
+
+    my_fn()
+
+    trace_msgs = _create_trace_messages(streamer)
+    span_msgs = _create_span_messages(streamer)
+    assert trace_msgs and span_msgs
+    assert all(m.environment == "staging" for m in trace_msgs)
+    assert all(m.environment == "staging" for m in span_msgs)
+
+
+def test_track_decorator__nested_environment_mismatch_warns_and_inherits_parent(
+    monkeypatch,
+):
+    from opik.api_objects import helpers as opik_helpers
+
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+    opik_client.set_global_client(client)
+
+    warnings = []
+    monkeypatch.setattr(
+        opik_helpers.LOGGER,
+        "warning",
+        lambda msg, *a, **kw: warnings.append(msg % a if a else msg),
+    )
+
+    @opik.track(environment="this-should-be-ignored")
+    def inner():
+        return 1
+
+    @opik.track(environment="staging")
+    def outer():
+        return inner()
+
+    outer()
+
+    span_msgs = _create_span_messages(streamer)
+    assert span_msgs
+    assert all(m.environment == "staging" for m in span_msgs)
+    assert any("this-should-be-ignored" in w and "staging" in w for w in warnings), (
+        warnings
+    )
+
+
+def test_create_trace_message__environment_field_default_is_none_for_backwards_compat():
+    msg = messages.CreateTraceMessage(
+        trace_id="t",
+        project_name="p",
+        name=None,
+        start_time=None,  # type: ignore[arg-type]
+        end_time=None,
+        input=None,
+        output=None,
+        metadata=None,
+        tags=None,
+        error_info=None,
+        thread_id=None,
+        last_updated_at=None,
+        source="sdk",
+    )
+    assert msg.environment is None
+
+
+def test_update_environment__colour_on_builtin_raises_error():
+    from opik.exceptions import EnvironmentConfigurationError
+
+    client = opik_client.Opik(project_name="test-project")
+
+    for env_name in ("production", "staging", "development"):
+        try:
+            client.update_environment(env_name, color="#ff0000")
+            assert False, f"expected EnvironmentConfigurationError for {env_name!r}"
+        except EnvironmentConfigurationError as e:
+            assert env_name in str(e)
+
+
+def test_update_environment__colour_on_builtin_not_called_when_no_colour():
+    from unittest.mock import patch
+    from opik.rest_api.types.environment_public import EnvironmentPublic
+
+    client = opik_client.Opik(project_name="test-project")
+
+    fake_env = EnvironmentPublic(id="abc", name="production")
+    with (
+        patch.object(client, "_find_environment_by_name", return_value=fake_env),
+        patch.object(
+            client._rest_client.environments,
+            "update_environment",
+            return_value=None,
+        ) as mock_update,
+        patch.object(
+            client._rest_client.environments,
+            "get_environment_by_id",
+            return_value=fake_env,
+        ),
+    ):
+        result = client.update_environment("production", description="new desc")
+        mock_update.assert_called_once()
+        assert result == fake_env
+
+
+def test_update_environment__colour_on_custom_env_is_allowed():
+    from unittest.mock import patch
+    from opik.rest_api.types.environment_public import EnvironmentPublic
+
+    client = opik_client.Opik(project_name="test-project")
+
+    fake_env = EnvironmentPublic(id="xyz", name="my-custom-env")
+    with (
+        patch.object(client, "_find_environment_by_name", return_value=fake_env),
+        patch.object(
+            client._rest_client.environments,
+            "update_environment",
+            return_value=None,
+        ) as mock_update,
+        patch.object(
+            client._rest_client.environments,
+            "get_environment_by_id",
+            return_value=fake_env,
+        ),
+    ):
+        result = client.update_environment("my-custom-env", color="#123456")
+        assert result == fake_env
+        mock_update.assert_called_once()
+        assert mock_update.call_args.kwargs.get("color") == "#123456"
+
+
+def test_create_environment__conflict_raises_environment_already_exists():
+    from unittest.mock import patch
+    from opik.exceptions import EnvironmentAlreadyExists
+    from opik.rest_api.errors import ConflictError
+
+    client = opik_client.Opik(project_name="test-project")
+
+    with patch.object(
+        client._rest_client.environments,
+        "create_environment",
+        side_effect=ConflictError(body={"message": "already exists"}),
+    ):
+        try:
+            client.create_environment("production")
+            assert False, "expected EnvironmentAlreadyExists"
+        except EnvironmentAlreadyExists as e:
+            assert "production" in str(e)
+
+
+def test_create_span_message__environment_field_default_is_none_for_backwards_compat():
+    msg = messages.CreateSpanMessage(
+        span_id="s",
+        trace_id="t",
+        project_name="p",
+        parent_span_id=None,
+        name=None,
+        start_time=None,  # type: ignore[arg-type]
+        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=None,
+        source="sdk",
+    )
+    assert msg.environment is None
+
+
+def _update_span_messages(streamer: MagicMock):
+    return [
+        c.args[0]
+        for c in streamer.put.call_args_list
+        if isinstance(c.args[0], messages.UpdateSpanMessage)
+    ]
+
+
+def _update_trace_messages(streamer: MagicMock):
+    return [
+        c.args[0]
+        for c in streamer.put.call_args_list
+        if isinstance(c.args[0], messages.UpdateTraceMessage)
+    ]
+
+
+def test_span_end__preserves_environment_in_update_message():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    trace = client.trace(name="t", environment="staging")
+    span = trace.span(name="s")
+    span.end()
+
+    update_msgs = _update_span_messages(streamer)
+    assert len(update_msgs) == 1
+    assert update_msgs[0].environment == "staging"
+
+
+def test_span_update__preserves_environment_in_update_message():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    trace = client.trace(name="t", environment="production")
+    span = trace.span(name="s")
+    span.update(output={"result": "ok"})
+
+    update_msgs = _update_span_messages(streamer)
+    assert len(update_msgs) == 1
+    assert update_msgs[0].environment == "production"
+
+
+def test_trace_end__preserves_environment_in_update_message():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    trace = client.trace(name="t", environment="staging")
+    trace.end()
+
+    update_msgs = _update_trace_messages(streamer)
+    assert len(update_msgs) == 1
+    assert update_msgs[0].environment == "staging"
+
+
+def test_trace_update__preserves_environment_in_update_message():
+    client = opik_client.Opik(project_name="test-project")
+    streamer = _capture_messages(client)
+
+    trace = client.trace(name="t", environment="production")
+    trace.update(output={"result": "ok"})
+
+    update_msgs = _update_trace_messages(streamer)
+    assert len(update_msgs) == 1
+    assert update_msgs[0].environment == "production"

sdks/python/tests/unit/api_objects/test_helpers.py

@@ -0,0 +1,19 @@
+import pytest
+
+from opik.api_objects import helpers
+
+
+@pytest.mark.parametrize(
+    "usage,metadata,create_metadata,expected",
+    [
+        (None, None, False, None),
+        (None, {"foo": "bar"}, False, {"foo": "bar"}),
+        ({}, None, False, None),
+        ({"foo": "bar"}, None, True, {"usage": {"foo": "bar"}}),
+    ],
+)
+def test_add_usage_to_metadata(usage, metadata, create_metadata, expected):
+    result = helpers.add_usage_to_metadata(
+        usage=usage, metadata=metadata, create_metadata=create_metadata
+    )
+    assert result == expected

sdks/python/tests/unit/api_objects/test_opik_client_sharing.py

@@ -0,0 +1,119 @@
+"""Sharing behaviour of Opik clients over ref-counted connection resources.
+
+These assert only through the public client surface (`rest_client`, `trace()`,
+`flush()`, `end()`): clients with a matching connection config reuse one
+transport, distinct configs get their own, and ending one client does not
+disable another that shares the connection. The ref-counting/eviction mechanics
+themselves are covered at the manager level in
+``test_connection_resource_manager.py``.
+"""
+
+import gc
+import threading
+import time
+import weakref
+from unittest import mock
+
+from opik.api_objects import opik_client
+
+
+def _make_client(**kwargs) -> opik_client.Opik:
+    return opik_client.Opik(_show_misconfiguration_message=False, **kwargs)
+
+
+def test_get_global_client__concurrent_cold_start__creates_single_client():
+    # Regression: get_global_client() must create the singleton once under
+    # concurrency. When several threads hit the cold-start path together (e.g. a
+    # tracer's _opik_client property accessed from parallel pipelines), each
+    # building its own client would race the shared connection-resource manager
+    # and, under a streamer-sharing test backend, close a streamer still in use —
+    # hanging a later flush().
+    opik_client.reset_global_client(end_client=False)
+
+    thread_count = 8
+    barrier = threading.Barrier(thread_count)
+    constructed = []
+    results = []
+    results_lock = threading.Lock()
+
+    def slow_construct(*args, **kwargs):
+        time.sleep(0.02)  # widen the window a racy implementation would lose in
+        client = object()
+        constructed.append(client)
+        return client
+
+    def worker() -> None:
+        barrier.wait()  # release all threads into the cold-start path together
+        client = opik_client.get_global_client()
+        with results_lock:
+            results.append(client)
+
+    try:
+        with mock.patch.object(opik_client, "Opik", side_effect=slow_construct):
+            threads = [threading.Thread(target=worker) for _ in range(thread_count)]
+            for thread in threads:
+                thread.start()
+            for thread in threads:
+                thread.join()
+
+        assert len(constructed) == 1  # built exactly once despite the race
+        assert {id(client) for client in results} == {id(constructed[0])}
+    finally:
+        opik_client.reset_global_client(end_client=False)
+
+
+def test_opik_clients__matching_connection_config__share_one_rest_client():
+    client_a = _make_client()
+    client_b = _make_client()
+    try:
+        # A shared connection is observable through the public REST client:
+        # both handles expose the very same underlying client object.
+        assert client_a.rest_client is client_b.rest_client
+    finally:
+        client_a.end(flush=False)
+        client_b.end(flush=False)
+
+
+def test_opik_clients__distinct_connection_config__use_separate_rest_clients():
+    client_a = _make_client()
+    client_b = _make_client(host="http://localhost:39999/api")
+    try:
+        assert client_a.rest_client is not client_b.rest_client
+    finally:
+        client_a.end(flush=False)
+        client_b.end(flush=False)
+
+
+def test_opik_client__logs_after_co_located_client_ended__data_still_delivered(
+    fake_backend,
+):
+    keeper = _make_client()
+    transient = _make_client()  # shares keeper's connection
+
+    # Ending one client releases only its reference; the shared transport must
+    # stay alive for the other handle.
+    transient.end(flush=False)
+
+    keeper.trace(name="after-sibling-end")
+    keeper.flush()
+
+    assert [trace.name for trace in fake_backend.trace_trees] == ["after-sibling-end"]
+
+    keeper.end(flush=False)
+
+
+def test_opik_client__dropped_without_end__is_garbage_collected(fake_backend):
+    # No lingering strong reference (a cached manager entry, the GC finalizer, a
+    # background thread, or the global singleton) should keep a dropped client
+    # alive. After `del` + `gc.collect()`, its weakref must no longer resolve.
+    # A unique host gives this handle its own isolated bundle so the assertion is
+    # about this client alone.
+    opik_client.reset_global_client(end_client=False)
+
+    client = _make_client(host="http://localhost:39998/api")
+    client_ref = weakref.ref(client)
+
+    del client
+    gc.collect()
+
+    assert client_ref() is None

sdks/python/tests/unit/api_objects/test_resolve_project_name.py

@@ -0,0 +1,86 @@
+import logging
+from unittest.mock import patch
+
+import pytest
+
+from opik.api_objects import helpers
+from opik.config import OPIK_PROJECT_DEFAULT_NAME
+from opik._logging import LOG_ONCE_CACHE
+
+WARNING_MESSAGE = (
+    'No project name configured. Traces are being logged to "Default Project".\n'
+    "Set OPIK_PROJECT_NAME environment variable or pass project_name to the Opik client\n"
+    "to log to a specific project.\n"
+    "See https://www.comet.com/docs/opik/tracing/advanced/sdk_configuration"
+)
+
+
+@pytest.fixture(autouse=True)
+def _clear_log_once_cache():
+    LOG_ONCE_CACHE.discard(WARNING_MESSAGE)
+    yield
+
+
+class TestResolveProjectName:
+    def test_explicit_project_name_returned(self):
+        result = helpers.resolve_project_name(
+            explicitly_passed_value="My Project",
+            value_from_config=OPIK_PROJECT_DEFAULT_NAME,
+        )
+        assert result == "My Project"
+
+    def test_explicit_default_project_name_no_warning(self):
+        with patch(
+            "opik.api_objects.helpers.opik_logging.log_once_at_level"
+        ) as mock_log:
+            result = helpers.resolve_project_name(
+                explicitly_passed_value=OPIK_PROJECT_DEFAULT_NAME,
+                value_from_config=OPIK_PROJECT_DEFAULT_NAME,
+            )
+        assert result == OPIK_PROJECT_DEFAULT_NAME
+        mock_log.assert_not_called()
+
+    def test_context_project_returned(self):
+        result = helpers.resolve_project_name(
+            explicitly_passed_value=None,
+            value_from_config=OPIK_PROJECT_DEFAULT_NAME,
+            value_from_context="Context Project",
+        )
+        assert result == "Context Project"
+
+    def test_context_project_no_warning(self):
+        with patch(
+            "opik.api_objects.helpers.opik_logging.log_once_at_level"
+        ) as mock_log:
+            helpers.resolve_project_name(
+                explicitly_passed_value=None,
+                value_from_config=OPIK_PROJECT_DEFAULT_NAME,
+                value_from_context="Context Project",
+            )
+        mock_log.assert_not_called()
+
+    def test_default_fallback_warns(self):
+        with patch(
+            "opik.api_objects.helpers.opik_logging.log_once_at_level"
+        ) as mock_log:
+            result = helpers.resolve_project_name(
+                explicitly_passed_value=None,
+                value_from_config=OPIK_PROJECT_DEFAULT_NAME,
+            )
+        assert result == OPIK_PROJECT_DEFAULT_NAME
+        mock_log.assert_called_once_with(
+            logging_level=logging.WARNING,
+            message=WARNING_MESSAGE,
+            logger=helpers.LOGGER,
+        )
+
+    def test_non_default_config_project_no_warning(self):
+        with patch(
+            "opik.api_objects.helpers.opik_logging.log_once_at_level"
+        ) as mock_log:
+            result = helpers.resolve_project_name(
+                explicitly_passed_value=None,
+                value_from_config="Custom Project",
+            )
+        assert result == "Custom Project"
+        mock_log.assert_not_called()

sdks/python/tests/unit/api_objects/test_rest_stream_parser.py

@@ -0,0 +1,163 @@
+import json
+
+import httpx
+import pytest
+
+from opik.rest_api.types import span_public as rest_api_types
+from opik.api_objects import rest_stream_parser
+
+SPANS_STREAM_JSON = [
+    {
+        "id": "0195f6f1-9da2-7630-b285-60cf5580372f",
+        "project_id": "0195f6f1-9c82-751f-a1ad-54fec4b5c7d8",
+        "trace_id": "0195f6f1-9c3f-7b65-a334-56173d19bc00",
+        "parent_span_id": "0195f6f1-9d9f-77c0-8063-17c59f6d9875",
+        "name": "synthesize",
+        "type": "general",
+        "start_time": "2025-01-03T11:03:17.875608Z",
+        "end_time": "2025-01-03T11:03:18.591814Z",
+        "input": {"query_str": "If Opik had a motto, what would it be?"},
+        "output": {
+            "output": '"Empowering continuous improvement through community-driven innovation."'
+        },
+        "created_at": "2025-04-02T14:39:44.412550Z",
+        "last_updated_at": "2025-04-02T14:39:44.412550Z",
+        "created_by": "admin",
+        "last_updated_by": "admin",
+        "duration": 716.206,
+    },
+    {
+        "id": "0195f6f1-9d9f-77c0-8063-17c59f6d9875",
+        "project_id": "0195f6f1-9c82-751f-a1ad-54fec4b5c7d8",
+        "trace_id": "0195f6f1-9c3f-7b65-a334-56173d19bc00",
+        "name": "query",
+        "type": "general",
+        "start_time": "2025-01-03T11:03:17.505783Z",
+        "end_time": "2025-01-03T11:03:18.591897Z",
+        "input": {"query_str": "If Opik had a motto, what would it be?"},
+        "output": {
+            "output": '"Empowering continuous improvement through community-driven innovation."'
+        },
+        "created_at": "2025-04-02T14:39:44.412550Z",
+        "last_updated_at": "2025-04-02T14:39:44.412550Z",
+        "created_by": "admin",
+        "last_updated_by": "admin",
+        "duration": 1086.114,
+    },
+]
+
+
+@pytest.fixture
+def spans_stream_source():
+    spans_stream = [
+        f"{json.dumps(span)}\r\n".encode("utf-8") for span in SPANS_STREAM_JSON
+    ]
+    yield spans_stream
+
+
+def test_read_and_parse_stream__span(spans_stream_source):
+    spans = rest_stream_parser.read_and_parse_stream(
+        spans_stream_source, item_class=rest_api_types.SpanPublic
+    )
+    assert len(spans) == 2
+    for i, span in enumerate(spans):
+        expected = rest_api_types.SpanPublic.model_validate(SPANS_STREAM_JSON[i])
+        assert span == expected
+
+
+def test_read_and_parse_stream__limit_samples(spans_stream_source):
+    spans = rest_stream_parser.read_and_parse_stream(
+        spans_stream_source, item_class=rest_api_types.SpanPublic, nb_samples=1
+    )
+    assert len(spans) == 1
+    expected = rest_api_types.SpanPublic.model_validate(SPANS_STREAM_JSON[0])
+    assert spans[0] == expected
+
+
+def test_read_and_parse_full_stream__happy_flow(spans_stream_source):
+    spans = rest_stream_parser.read_and_parse_full_stream(
+        read_source=lambda current_batch_size, last_retrieved_id: spans_stream_source,
+        parsed_item_class=rest_api_types.SpanPublic,
+        max_results=10,
+    )
+    assert len(spans) == 2
+    for i, span in enumerate(spans):
+        expected = rest_api_types.SpanPublic.model_validate(SPANS_STREAM_JSON[i])
+        assert span == expected
+
+
+def test_read_and_parse_full_stream__no_error__requested_batch_sizes_not_halved(
+    spans_stream_source,
+):
+    requested_batch_sizes = []
+
+    def read_source(current_batch_size, last_retrieved_id):
+        requested_batch_sizes.append(current_batch_size)
+        return spans_stream_source
+
+    rest_stream_parser.read_and_parse_full_stream(
+        read_source=read_source,
+        parsed_item_class=rest_api_types.SpanPublic,
+        max_results=None,
+        max_endpoint_batch_size=400,
+    )
+
+    # Two spans returned (< page size) ends the loop after one request; the
+    # requested page size is the configured one, never shrunk.
+    assert requested_batch_sizes == [400]
+
+
+def test_read_and_parse_full_stream__size_correlated_error__halves_page_and_retries_same_cursor(
+    spans_stream_source,
+):
+    requested = []
+    calls = {"n": 0}
+
+    def read_source(current_batch_size, last_retrieved_id):
+        requested.append((current_batch_size, last_retrieved_id))
+        calls["n"] += 1
+        if calls["n"] == 1:
+            raise httpx.RemoteProtocolError("incomplete chunked read")
+        return spans_stream_source
+
+    spans = rest_stream_parser.read_and_parse_full_stream(
+        read_source=read_source,
+        parsed_item_class=rest_api_types.SpanPublic,
+        max_results=None,
+        max_endpoint_batch_size=400,
+    )
+
+    # First request at 400 fails; retried at 200 from the SAME (None) cursor.
+    assert requested == [(400, None), (200, None)]
+    assert len(spans) == 2
+
+
+def test_read_and_parse_full_stream__shrink_floor_reached__reraises(
+    spans_stream_source,
+):
+    def read_source(current_batch_size, last_retrieved_id):
+        raise httpx.ReadTimeout("timed out")
+
+    # Start at the floor so the first failure can't shrink further.
+    with pytest.raises(httpx.ReadTimeout):
+        rest_stream_parser.read_and_parse_full_stream(
+            read_source=read_source,
+            parsed_item_class=rest_api_types.SpanPublic,
+            max_results=None,
+            max_endpoint_batch_size=rest_stream_parser.MIN_ENDPOINT_BATCH_SIZE,
+        )
+
+
+def test_read_and_parse_full_stream__non_size_correlated_error__propagates(
+    spans_stream_source,
+):
+    def read_source(current_batch_size, last_retrieved_id):
+        raise ValueError("not a connection error")
+
+    with pytest.raises(ValueError):
+        rest_stream_parser.read_and_parse_full_stream(
+            read_source=read_source,
+            parsed_item_class=rest_api_types.SpanPublic,
+            max_results=None,
+            max_endpoint_batch_size=400,
+        )

sdks/python/tests/unit/api_objects/test_search_helpers.py

@@ -0,0 +1,216 @@
+import json
+import logging
+from unittest import mock
+
+import pytest
+
+from opik.api_objects import search_helpers
+from opik.rest_api.core.api_error import ApiError
+from opik.rest_api.types import span_public, trace_public
+
+
+SPAN_RECORD = {
+    "id": "0195f6f1-9da2-7630-b285-60cf5580372f",
+    "project_id": "0195f6f1-9c82-751f-a1ad-54fec4b5c7d8",
+    "trace_id": "0195f6f1-9c3f-7b65-a334-56173d19bc00",
+    "name": "synthesize",
+    "type": "general",
+    "start_time": "2025-01-03T11:03:17.875608Z",
+    "end_time": "2025-01-03T11:03:18.591814Z",
+    "input": {"q": "x"},
+    "output": {"a": "y"},
+    "created_at": "2025-04-02T14:39:44.412550Z",
+    "last_updated_at": "2025-04-02T14:39:44.412550Z",
+    "created_by": "admin",
+    "last_updated_by": "admin",
+    "duration": 1.0,
+}
+
+TRACE_RECORD = {
+    "id": "0195f6f1-9c3f-7b65-a334-56173d19bc00",
+    "project_id": "0195f6f1-9c82-751f-a1ad-54fec4b5c7d8",
+    "name": "root",
+    "start_time": "2025-01-03T11:03:17.505783Z",
+    "end_time": "2025-01-03T11:03:18.591897Z",
+    "input": {"q": "x"},
+    "output": {"a": "y"},
+    "created_at": "2025-04-02T14:39:44.412550Z",
+    "last_updated_at": "2025-04-02T14:39:44.412550Z",
+    "created_by": "admin",
+    "last_updated_by": "admin",
+    "duration": 1.0,
+}
+
+
+def _stream(record: dict):
+    """Mimic the bytes-iterator returned by the Fern-generated search clients."""
+    yield f"{json.dumps(record)}\r\n".encode("utf-8")
+
+
+@pytest.fixture
+def patched_sleep():
+    # Patch the rest_helpers retry-delay indirection, NOT the global time.sleep:
+    # a global patch turns background daemon threads' pacing sleep into a busy
+    # loop that can starve the interpreter and hang the suite.
+    with mock.patch("opik.api_objects.rest_helpers._sleep") as sleep_mock:
+        yield sleep_mock
+
+
+def _sleep_called_with(mock_sleep: mock.Mock, seconds) -> bool:
+    return any(call.args == (seconds,) for call in mock_sleep.call_args_list)
+
+
+def test_search_spans__429_response__retries_with_reset_header(
+    patched_sleep, capture_log
+):
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.side_effect = [
+        ApiError(status_code=429, headers={"RateLimit-Reset": "2"}),
+        _stream(SPAN_RECORD),
+    ]
+
+    result = search_helpers.search_spans_with_filters(
+        rest_client=rest_client,
+        trace_id=None,
+        project_name="proj",
+        filters=None,
+        max_results=10,
+        truncate=False,
+    )
+
+    assert len(result) == 1
+    assert result[0] == span_public.SpanPublic.model_validate(SPAN_RECORD)
+    assert rest_client.spans.search_spans.call_count == 2
+    assert _sleep_called_with(patched_sleep, 2)
+    assert any(
+        "search_spans" in record.message and record.levelno == logging.WARNING
+        for record in capture_log.records
+    )
+
+
+def test_search_traces__429_response__retries_with_reset_header(
+    patched_sleep, capture_log
+):
+    rest_client = mock.Mock()
+    rest_client.traces.search_traces.side_effect = [
+        ApiError(status_code=429, headers={"RateLimit-Reset": "3"}),
+        _stream(TRACE_RECORD),
+    ]
+
+    result = search_helpers.search_traces_with_filters(
+        rest_client=rest_client,
+        project_name="proj",
+        filters=None,
+        max_results=10,
+        truncate=False,
+    )
+
+    assert len(result) == 1
+    assert result[0] == trace_public.TracePublic.model_validate(TRACE_RECORD)
+    assert rest_client.traces.search_traces.call_count == 2
+    assert _sleep_called_with(patched_sleep, 3)
+    assert any(
+        "search_traces" in record.message and record.levelno == logging.WARNING
+        for record in capture_log.records
+    )
+
+
+def test_search_spans__non_429_error__propagates_without_retry():
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.side_effect = ApiError(
+        status_code=400, body="bad request"
+    )
+
+    with pytest.raises(ApiError) as exc_info:
+        search_helpers.search_spans_with_filters(
+            rest_client=rest_client,
+            trace_id=None,
+            project_name="proj",
+            filters=None,
+            max_results=10,
+            truncate=False,
+        )
+
+    assert exc_info.value.status_code == 400
+    # call_count == 1 already proves no retry happened — no need to assert on
+    # patched_sleep, which is polluted by background-thread sleeps.
+    assert rest_client.spans.search_spans.call_count == 1
+
+
+def test_search_traces__non_429_error__propagates_without_retry():
+    rest_client = mock.Mock()
+    rest_client.traces.search_traces.side_effect = ApiError(
+        status_code=500, body="boom"
+    )
+
+    with pytest.raises(ApiError) as exc_info:
+        search_helpers.search_traces_with_filters(
+            rest_client=rest_client,
+            project_name="proj",
+            filters=None,
+            max_results=10,
+            truncate=False,
+        )
+
+    assert exc_info.value.status_code == 500
+    assert rest_client.traces.search_traces.call_count == 1
+
+
+def test_search_spans__max_batch_size__passed_as_limit_to_rest_client():
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.return_value = _stream(SPAN_RECORD)
+
+    search_helpers.search_spans_with_filters(
+        rest_client=rest_client,
+        trace_id=None,
+        project_name="proj",
+        filters=None,
+        max_results=10_000,
+        truncate=False,
+        max_batch_size=400,
+    )
+
+    # max_results (10_000) is the total cap; max_batch_size (400) is the
+    # per-request page size that reaches the Fern client as ``limit``.
+    _, kwargs = rest_client.spans.search_spans.call_args
+    assert kwargs["limit"] == 400
+
+
+def test_search_traces__max_batch_size__passed_as_limit_to_rest_client():
+    rest_client = mock.Mock()
+    rest_client.traces.search_traces.return_value = _stream(TRACE_RECORD)
+
+    search_helpers.search_traces_with_filters(
+        rest_client=rest_client,
+        project_name="proj",
+        filters=None,
+        max_results=10_000,
+        truncate=False,
+        max_batch_size=250,
+    )
+
+    _, kwargs = rest_client.traces.search_traces.call_args
+    assert kwargs["limit"] == 250
+
+
+def test_search_spans__429_without_reset_header__falls_back_to_one_second_sleep(
+    patched_sleep,
+):
+    rest_client = mock.Mock()
+    rest_client.spans.search_spans.side_effect = [
+        ApiError(status_code=429, headers={}),
+        _stream(SPAN_RECORD),
+    ]
+
+    result = search_helpers.search_spans_with_filters(
+        rest_client=rest_client,
+        trace_id=None,
+        project_name="proj",
+        filters=None,
+        max_results=10,
+        truncate=False,
+    )
+
+    assert len(result) == 1
+    assert rest_client.spans.search_spans.call_count == 2
+    assert _sleep_called_with(patched_sleep, 1)

sdks/python/tests/unit/api_objects/trace/test_trace_data.py

@@ -0,0 +1,67 @@
+import datetime
+
+from opik.api_objects import trace
+from ....testlib import assert_equal, ANY_BUT_NONE
+
+
+def test_trace_data__as_start_parameters__expected_parameters_are_set():
+    trace_data = trace.TraceData(
+        name="name",
+        project_name="project_name",
+        metadata={"foo": "bar"},
+        input={"input": "input"},
+        tags=["one", "two"],
+        created_by="evaluation",
+        source="sdk",
+    )
+
+    expected_parameters = {
+        "id": ANY_BUT_NONE,
+        "start_time": ANY_BUT_NONE,
+        "project_name": "project_name",
+        "name": "name",
+        "metadata": {"foo": "bar"},
+        "input": {"input": "input"},
+        "tags": ["one", "two"],
+        "source": "sdk",
+    }
+
+    assert_equal(expected_parameters, trace_data.as_start_parameters)
+
+
+def test_trace_data__as_parameters__expected_parameters_are_set():
+    trace_data = trace.TraceData(
+        name="name",
+        end_time=datetime.datetime.now(),
+        metadata={"foo": "bar"},
+        input={"input": "input"},
+        output={"output": "output"},
+        tags=["one", "two"],
+        feedback_scores=[{"name": "score_name", "value": 0.5}],
+        project_name="project_name",
+        error_info={},
+        thread_id="thread_id",
+        attachments=[],
+        created_by="evaluation",
+        source="sdk",
+    )
+
+    expected_parameters = {
+        "id": ANY_BUT_NONE,
+        "name": "name",
+        "start_time": ANY_BUT_NONE,
+        "end_time": ANY_BUT_NONE,
+        "metadata": ANY_BUT_NONE,
+        "input": ANY_BUT_NONE,
+        "output": ANY_BUT_NONE,
+        "tags": ANY_BUT_NONE,
+        "feedback_scores": ANY_BUT_NONE,
+        "project_name": "project_name",
+        "error_info": ANY_BUT_NONE,
+        "thread_id": ANY_BUT_NONE,
+        "attachments": ANY_BUT_NONE,
+        "source": "sdk",
+        "environment": None,
+    }
+
+    assert_equal(expected_parameters, trace_data.as_parameters)

sdks/python/tests/unit/cli/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for CLI functionality."""

sdks/python/tests/unit/cli/_migrate_helpers.py

@@ -0,0 +1,196 @@
+"""Shared mock helpers for ``opik migrate`` test modules.
+
+Used by both ``test_migrate_dataset_exclude_versions.py`` (Slice 1 paths) and
+``test_migrate_dataset_version_replay.py`` (Slice 2 paths). Lives as a plain
+module (not conftest.py) because these are helper classes, not pytest
+fixtures.
+"""
+
+from __future__ import annotations
+
+from typing import Dict, List, Optional
+from unittest.mock import MagicMock
+
+
+class _DatasetRow:
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        description: Optional[str] = None,
+        items: int = 0,
+        type: Optional[str] = "dataset",
+        visibility: Optional[str] = "private",
+        tags: Optional[List[str]] = None,
+        # ``project_id=None`` represents a workspace-scoped dataset (V1
+        # entity, or anything left at workspace scope after auto-migration).
+        # Tests that want a project-scoped source pass an explicit id.
+        project_id: Optional[str] = None,
+    ) -> None:
+        self.id = id
+        self.name = name
+        self.description = description
+        self.dataset_items_count = items
+        self.type = type
+        self.visibility = visibility
+        self.tags = tags
+        self.project_id = project_id
+
+
+class _Page:
+    def __init__(self, content: List[_DatasetRow]) -> None:
+        self.content = content
+
+
+def _named(name: str) -> MagicMock:
+    obj = MagicMock()
+    obj.name = name
+    return obj
+
+
+def _planner_rest_client(
+    find_side_effects: List[_Page],
+    *,
+    target_project_exists: bool = True,
+    workspace_project_names: Optional[List[str]] = None,
+) -> MagicMock:
+    """Build a rest_client mock for direct planner unit tests."""
+    rest_client = MagicMock()
+    if target_project_exists:
+        target_project = MagicMock()
+        target_project.id = "target-project-id"
+        rest_client.projects.retrieve_project.return_value = target_project
+    else:
+        from opik.rest_api.core.api_error import ApiError
+
+        rest_client.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={}
+        )
+        # Suggestion lookup queries find_projects on the 404 path.
+        candidates = [_named(name) for name in (workspace_project_names or [])]
+        rest_client.projects.find_projects.return_value = _Page(candidates)
+    rest_client.datasets.find_datasets.side_effect = find_side_effects
+    return rest_client
+
+
+def _planner_client(rest_client: MagicMock) -> MagicMock:
+    """Wrap a planner rest_client mock as an ``opik.Opik``-shaped client.
+
+    The planner now takes the high-level client (so the resolver can route
+    ``get_project_by_id`` through ``client.get_project`` instead of the
+    Fern surface). Tests still build the low-level rest_client mock and
+    drive its side_effects, then wrap with this helper to satisfy the new
+    planner signature.
+
+    ``client.get_project(id=...)`` delegates to
+    ``rest_client.projects.get_project_by_id(id=...)`` so any per-test
+    stub on the rest_client side still flows through unchanged.
+    """
+    client = MagicMock()
+    client.rest_client = rest_client
+    client.get_project = MagicMock(
+        side_effect=lambda id: rest_client.projects.get_project_by_id(id=id)
+    )
+    return client
+
+
+def _build_fake_client(
+    *,
+    source_rows: List[_DatasetRow],
+    destination_rows: List[_DatasetRow],
+    items: List[Dict[str, object]],
+    target_project_exists: bool = True,
+    stale_temp_rows: Optional[List[_DatasetRow]] = None,
+) -> MagicMock:
+    """Construct an opik.Opik mock matching the executor's call surface.
+
+    The planner makes three workspace ``find_datasets`` lookups in a fixed
+    order: (1) source resolution, (2) the ``<name>_v1`` rename-target
+    collision pre-flight, and (3) the ``<name>__migrating`` stale-temp
+    lookup (OPIK-7162). A side-effect list drives all three deterministically.
+    ``destination_rows`` feeds the ``_v1`` check; the stale-temp lookup
+    returns empty by default (no leftover temp) unless ``stale_temp_rows``
+    is supplied.
+
+    ``items`` is a list of dicts; we materialize them as DatasetItem
+    dataclasses for the streaming mock (matches the real `__internal_api__
+    stream_items_as_dataclasses__` shape) so per-item fidelity assertions
+    have somewhere to land.
+    """
+    from opik.api_objects.dataset import dataset_item
+
+    rest_client = MagicMock()
+    if target_project_exists:
+        target_project = MagicMock()
+        target_project.id = "target-project-id"
+        rest_client.projects.retrieve_project.return_value = target_project
+    else:
+        from opik.rest_api.core.api_error import ApiError
+
+        rest_client.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={}
+        )
+    rest_client.datasets.find_datasets.side_effect = [
+        _Page(source_rows),
+        _Page(destination_rows),
+        _Page(stale_temp_rows or []),
+    ]
+    rest_client.datasets.update_dataset = MagicMock()
+    rest_client.datasets.create_dataset = MagicMock()
+    rest_client.datasets.delete_dataset = MagicMock()
+
+    client = MagicMock()
+    client.rest_client = rest_client
+    client._workspace = "default"
+
+    # Build DatasetItem dataclasses from the provided dicts so the executor's
+    # dataclass-form stream returns a realistic shape. Top-level fields like
+    # `description` / `source` / `trace_id` / `span_id` can be passed via the
+    # dict (other keys get stuffed into `data`/extra).
+    top_level = {
+        "id",
+        "trace_id",
+        "span_id",
+        "source",
+        "description",
+        "evaluators",
+        "execution_policy",
+    }
+    source_items: List[dataset_item.DatasetItem] = []
+    for raw in items:
+        kwargs = {k: v for k, v in raw.items() if k in top_level and k != "id"}
+        data = {k: v for k, v in raw.items() if k not in top_level}
+        ds_item = dataset_item.DatasetItem(**kwargs, **data)
+        if "id" in raw:
+            ds_item.id = raw["id"]  # type: ignore[assignment]
+        source_items.append(ds_item)
+
+    # MagicMock treats dunder-prefixed names as magic and blocks them by
+    # default; pre-attach plain MagicMocks so attribute access works.
+    source_dataset = MagicMock()
+    stream_mock = MagicMock(return_value=iter(source_items))
+    source_dataset.__internal_api__stream_items_as_dataclasses__ = stream_mock
+
+    dest_dataset = MagicMock()
+    insert_mock = MagicMock()
+    dest_dataset.__internal_api__insert_items_as_dataclasses__ = insert_mock
+
+    # Under the OPIK-7162 ordering the destination is written under the temp
+    # name ``<orig>__migrating`` for the whole copy, then promoted to the
+    # original name at the end. The executor resolves the destination by name
+    # in ``_replay_versions`` (temp name) and again in ``PromoteDestination``
+    # (temp name -> id). The source is never resolved via ``get_dataset`` any
+    # more (item reads stream by name on the rest_client), so route the temp
+    # name (and anything else) to the destination dataset.
+    source_orig_name = source_rows[0].name if source_rows else ""
+    temp_name = f"{source_orig_name}__migrating"
+
+    def _get_dataset(name: str, project_name: Optional[str] = None) -> MagicMock:
+        if name == temp_name:
+            return dest_dataset
+        return source_dataset
+
+    client.get_dataset.side_effect = _get_dataset
+    client.create_dataset = MagicMock()
+    client.delete_dataset = MagicMock()
+    return client, source_dataset, dest_dataset

sdks/python/tests/unit/cli/_migrate_prompt_helpers.py

@@ -0,0 +1,125 @@
+"""Shared mock helpers for ``opik migrate prompt`` test modules.
+
+Used by ``test_migrate_prompt_planner.py``,
+``test_migrate_prompt_executor.py``, and
+``test_migrate_prompt_version_replay.py``. Lives as a plain module (not
+a conftest.py) because these are helper classes, not pytest fixtures.
+"""
+
+from __future__ import annotations
+
+from typing import Any, List, Optional
+from unittest.mock import MagicMock
+
+
+class _PromptRow:
+    """Minimal stand-in for ``PromptPublic`` rows returned by ``get_prompts``."""
+
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        *,
+        description: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        # None = workspace-scoped (legacy v1 prompts that never had a
+        # project). Tests that want a project-scoped source pass an id.
+        project_id: Optional[str] = None,
+        template_structure: Optional[str] = "text",
+    ) -> None:
+        self.id = id
+        self.name = name
+        self.description = description
+        self.tags = tags
+        self.project_id = project_id
+        self.template_structure = template_structure
+
+
+class _PromptVersionRow:
+    """Stand-in for ``PromptVersionPublic`` rows from ``get_prompt_versions``.
+
+    Field set mirrors the Fern model: ``id``, ``prompt_id``, ``commit``,
+    ``template`` (required), plus optional metadata / type / change_description
+    / tags / environments / template_structure.
+    """
+
+    def __init__(
+        self,
+        id: str,
+        prompt_id: str,
+        commit: str,
+        template: str,
+        *,
+        metadata: Optional[Any] = None,
+        type: Optional[str] = None,
+        change_description: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        environments: Optional[List[str]] = None,
+        template_structure: Optional[str] = "text",
+    ) -> None:
+        self.id = id
+        self.prompt_id = prompt_id
+        self.commit = commit
+        self.template = template
+        self.metadata = metadata
+        self.type = type
+        self.change_description = change_description
+        self.tags = tags
+        self.environments = environments
+        self.template_structure = template_structure
+
+
+class _Page:
+    def __init__(self, content: List[Any]) -> None:
+        self.content = content
+
+
+def _named(name: str) -> MagicMock:
+    obj = MagicMock()
+    obj.name = name
+    return obj
+
+
+def _planner_rest_client(
+    find_prompts_side_effects: List[_Page],
+    *,
+    target_project_exists: bool = True,
+    workspace_project_names: Optional[List[str]] = None,
+) -> MagicMock:
+    """Build a rest_client mock for direct prompt-planner unit tests.
+
+    ``find_prompts_side_effects`` is the side-effect list for
+    ``rest_client.prompts.get_prompts``; the planner calls it twice (once
+    for source resolution, once for the rename-collision preflight) so
+    callers pass two pages.
+    """
+    rest_client = MagicMock()
+    if target_project_exists:
+        target_project = MagicMock()
+        target_project.id = "target-project-id"
+        rest_client.projects.retrieve_project.return_value = target_project
+    else:
+        from opik.rest_api.core.api_error import ApiError
+
+        rest_client.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={}
+        )
+        candidates = [_named(n) for n in (workspace_project_names or [])]
+        rest_client.projects.find_projects.return_value = _Page(candidates)
+    rest_client.prompts.get_prompts.side_effect = find_prompts_side_effects
+    return rest_client
+
+
+def _planner_client(rest_client: MagicMock) -> MagicMock:
+    """Wrap a rest_client mock as an ``opik.Opik``-shaped client.
+
+    The resolver routes project-id lookups through ``client.get_project``
+    (high-level surface) so the project-name resolution still benefits
+    from any test-side stubs on ``rest_client.projects.get_project_by_id``.
+    """
+    client = MagicMock()
+    client.rest_client = rest_client
+    client.get_project = MagicMock(
+        side_effect=lambda id: rest_client.projects.get_project_by_id(id=id)
+    )
+    return client

sdks/python/tests/unit/cli/local_runner/test_stop.py

@@ -0,0 +1,313 @@
+"""Tests for opik.cli.local_runner.stop and the `opik <type> stop` CLI surface."""
+
+import json
+import os
+import signal
+from pathlib import Path
+from typing import Optional
+from unittest.mock import patch
+
+import pytest
+from click.testing import CliRunner
+
+from opik.cli.local_runner import stop as stop_module
+from opik.cli.local_runner.pairing import RunnerType
+from opik.cli.local_runner.stop import do_stop
+from opik.cli.main import cli
+from opik.runner import pid_file
+from opik.runner.pid_file import RunnerInfo
+
+
+@pytest.fixture
+def tmp_runners_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
+    runners = tmp_path / "runners"
+    monkeypatch.setattr(pid_file, "_RUNNERS_DIR", runners)
+    return runners
+
+
+def _write_runner_file(
+    dir: Path,
+    runner_id: str,
+    runner_type: str,
+    project_name: str,
+    pid: int,
+    workspace: Optional[str] = None,
+) -> Path:
+    dir.mkdir(parents=True, exist_ok=True)
+    path = dir / f"{runner_type}-{runner_id}.json"
+    path.write_text(
+        json.dumps(
+            {
+                "pid": pid,
+                "runner_id": runner_id,
+                "runner_type": runner_type,
+                "project_name": project_name,
+                "workspace": workspace,
+                "started_at": 1.0,
+            }
+        )
+    )
+    return path
+
+
+class TestDoStopFilters:
+    def test_do_stop__no_filter__usage_error(self, tmp_runners_dir: Path) -> None:
+        # No --project / --runner / --all → usage error before any signaling.
+        with pytest.raises(Exception) as exc_info:
+            do_stop(
+                runner_type=RunnerType.CONNECT,
+                project_name=None,
+                all_flag=False,
+                runner_id_filter=None,
+            )
+        # click.UsageError prints "Error: ..." and is a Click exception.
+        assert "specify" in str(exc_info.value).lower()
+
+    def test_do_stop__no_match__prints_no_match_message(
+        self, tmp_runners_dir: Path, capsys: pytest.CaptureFixture
+    ) -> None:
+        do_stop(
+            runner_type=RunnerType.CONNECT,
+            project_name="missing",
+            all_flag=False,
+            runner_id_filter=None,
+        )
+        out = capsys.readouterr().out
+        assert "No local 'connect' runners found" in out
+
+    def test_do_stop__project_filter__signals_matching_only(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        _write_runner_file(tmp_runners_dir, "r-1", "connect", "alpha", os.getpid())
+        _write_runner_file(tmp_runners_dir, "r-2", "connect", "beta", os.getpid())
+
+        seen = []
+
+        def fake_signal(info: RunnerInfo):
+            seen.append(info.runner_id)
+            return True, ""
+
+        with patch.object(stop_module, "_signal_until_gone", side_effect=fake_signal):
+            do_stop(
+                runner_type=RunnerType.CONNECT,
+                project_name="alpha",
+                all_flag=False,
+                runner_id_filter=None,
+            )
+
+        assert seen == ["r-1"]
+        # The matched runner's pid file is cleaned up.
+        assert not (tmp_runners_dir / "connect-r-1.json").exists()
+        # Non-matched is left alone.
+        assert (tmp_runners_dir / "connect-r-2.json").exists()
+
+    def test_do_stop__all_flag__signals_only_target_runner_type(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        _write_runner_file(tmp_runners_dir, "r-1", "connect", "alpha", os.getpid())
+        _write_runner_file(tmp_runners_dir, "r-2", "connect", "beta", os.getpid())
+        # An endpoint runner must NOT be touched by `connect stop --all`.
+        _write_runner_file(tmp_runners_dir, "r-3", "endpoint", "alpha", os.getpid())
+
+        seen = []
+        with patch.object(
+            stop_module,
+            "_signal_until_gone",
+            lambda info: (seen.append(info.runner_id) or (True, "")),
+        ):
+            do_stop(
+                runner_type=RunnerType.CONNECT,
+                project_name=None,
+                all_flag=True,
+                runner_id_filter=None,
+            )
+        assert sorted(seen) == ["r-1", "r-2"]
+        assert (tmp_runners_dir / "endpoint-r-3.json").exists()
+
+    def test_do_stop__runner_id_filter__signals_matching_only(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        _write_runner_file(tmp_runners_dir, "r-1", "connect", "alpha", os.getpid())
+        _write_runner_file(tmp_runners_dir, "r-2", "connect", "alpha", os.getpid())
+
+        seen = []
+        with patch.object(
+            stop_module,
+            "_signal_until_gone",
+            lambda info: (seen.append(info.runner_id) or (True, "")),
+        ):
+            do_stop(
+                runner_type=RunnerType.CONNECT,
+                project_name=None,
+                all_flag=False,
+                runner_id_filter="r-2",
+            )
+        assert seen == ["r-2"]
+
+
+class TestSignalUntilGone:
+    """Logic tests for the SIGTERM → wait → SIGKILL flow.
+
+    Real-subprocess kill is exercised by `tests/unit/runner/test_pid_file.py::TestSignalRoundtrip`;
+    those cases call `proc.wait()` to reap zombies, which pytest's parent-child
+    relationship requires. Here we mock the alive-check + signal so we can assert
+    the escalation order without fighting zombie state.
+    """
+
+    def _make_info(self, tmp_runners_dir: Path, pid: int = 12345) -> RunnerInfo:
+        return RunnerInfo(
+            pid=pid,
+            runner_id="r-1",
+            runner_type="connect",
+            project_name="p",
+            workspace=None,
+            started_at=0.0,
+            path=tmp_runners_dir / "connect-r-1.json",
+        )
+
+    def test_signal_until_gone__target_exits_after_sigterm__returns_ok(
+        self, tmp_runners_dir: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setattr(stop_module, "_POLL_INTERVAL_SECONDS", 0.01)
+        sent_signals = []
+        monkeypatch.setattr(
+            stop_module.os, "kill", lambda pid, sig: sent_signals.append((pid, sig))
+        )
+        # First poll says alive, next says gone.
+        alive_seq = iter([True, False])
+        monkeypatch.setattr(
+            stop_module.pid_file, "is_pid_alive", lambda _pid: next(alive_seq)
+        )
+
+        ok, _ = stop_module._signal_until_gone(self._make_info(tmp_runners_dir))
+        assert ok
+        assert sent_signals == [(12345, signal.SIGTERM)]
+
+    def test_signal_until_gone__sigterm_ignored__escalates_to_sigkill(
+        self, tmp_runners_dir: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        monkeypatch.setattr(stop_module, "_SIGTERM_GRACE_SECONDS", 0.05)
+        monkeypatch.setattr(stop_module, "_POLL_INTERVAL_SECONDS", 0.01)
+        sent_signals = []
+        state = {"alive": True}
+
+        def fake_kill(pid: int, sig: int) -> None:
+            sent_signals.append((pid, sig))
+            if sig == signal.SIGKILL:
+                state["alive"] = False
+
+        monkeypatch.setattr(stop_module.os, "kill", fake_kill)
+        monkeypatch.setattr(
+            stop_module.pid_file, "is_pid_alive", lambda _pid: state["alive"]
+        )
+
+        ok, _ = stop_module._signal_until_gone(self._make_info(tmp_runners_dir))
+        assert ok
+        assert [sig for _, sig in sent_signals] == [signal.SIGTERM, signal.SIGKILL]
+
+    def test_signal_until_gone__already_dead_pid__returns_ok(
+        self, tmp_runners_dir: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        def fake_kill(_pid: int, _sig: int) -> None:
+            raise ProcessLookupError()
+
+        monkeypatch.setattr(stop_module.os, "kill", fake_kill)
+        ok, reason = stop_module._signal_until_gone(self._make_info(tmp_runners_dir))
+        assert ok
+        assert "exited" in reason
+
+    def test_signal_until_gone__permission_error__returns_failure(
+        self, tmp_runners_dir: Path, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        def fake_kill(_pid: int, _sig: int) -> None:
+            raise PermissionError("not yours")
+
+        monkeypatch.setattr(stop_module.os, "kill", fake_kill)
+        ok, reason = stop_module._signal_until_gone(self._make_info(tmp_runners_dir))
+        assert not ok
+        assert "permission" in reason.lower()
+
+
+class TestStopExitCode:
+    def test_do_stop__signal_failed__raises_systemexit_1(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        _write_runner_file(tmp_runners_dir, "r-1", "connect", "alpha", os.getpid())
+        with patch.object(
+            stop_module, "_signal_until_gone", return_value=(False, "denied")
+        ):
+            with pytest.raises(SystemExit) as exc_info:
+                do_stop(
+                    runner_type=RunnerType.CONNECT,
+                    project_name="alpha",
+                    all_flag=False,
+                    runner_id_filter=None,
+                )
+            assert exc_info.value.code == 1
+
+
+class TestConnectStopCli:
+    def test_cli_connect_stop__no_filter__exits_with_usage_error(self) -> None:
+        runner = CliRunner()
+        result = runner.invoke(cli, ["connect", "stop"])
+        assert result.exit_code == 2
+        assert "Specify" in result.output or "specify" in result.output
+
+    def test_cli_connect_stop__with_project__routes_to_do_stop(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        runner = CliRunner()
+        with patch("opik.cli.local_runner.stop.do_stop") as mock_stop:
+            result = runner.invoke(cli, ["connect", "stop", "--project", "p"])
+            assert result.exit_code == 0, result.output
+            mock_stop.assert_called_once()
+            assert mock_stop.call_args.kwargs == dict(
+                runner_type=RunnerType.CONNECT,
+                project_name="p",
+                all_flag=False,
+                runner_id_filter=None,
+            )
+
+    def test_cli_connect_stop__all_flag__sets_all_flag_kwarg(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        runner = CliRunner()
+        with patch("opik.cli.local_runner.stop.do_stop") as mock_stop:
+            result = runner.invoke(cli, ["connect", "stop", "--all"])
+            assert result.exit_code == 0, result.output
+            assert mock_stop.call_args.kwargs["all_flag"] is True
+
+    def test_cli_connect_stop__runner_flag__sets_runner_id_filter(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        runner = CliRunner()
+        with patch("opik.cli.local_runner.stop.do_stop") as mock_stop:
+            result = runner.invoke(cli, ["connect", "stop", "--runner", "r-1"])
+            assert result.exit_code == 0, result.output
+            assert mock_stop.call_args.kwargs["runner_id_filter"] == "r-1"
+
+
+class TestEndpointStopCli:
+    def test_cli_endpoint_stop__with_project__routes_to_do_stop(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        runner = CliRunner()
+        with patch("opik.cli.local_runner.stop.do_stop") as mock_stop:
+            result = runner.invoke(cli, ["endpoint", "stop", "--project", "p"])
+            assert result.exit_code == 0, result.output
+            mock_stop.assert_called_once()
+            assert mock_stop.call_args.kwargs == dict(
+                runner_type=RunnerType.ENDPOINT,
+                project_name="p",
+                all_flag=False,
+                runner_id_filter=None,
+            )
+
+    def test_cli_endpoint_stop__all_flag__sets_all_flag_kwarg(
+        self, tmp_runners_dir: Path
+    ) -> None:
+        runner = CliRunner()
+        with patch("opik.cli.local_runner.stop.do_stop") as mock_stop:
+            result = runner.invoke(cli, ["endpoint", "stop", "--all"])
+            assert result.exit_code == 0, result.output
+            assert mock_stop.call_args.kwargs["all_flag"] is True

sdks/python/tests/unit/cli/test_cli.py

@@ -0,0 +1,445 @@
+"""Tests for CLI commands."""
+
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+from contextlib import contextmanager
+
+from click.testing import CliRunner
+
+from opik.cli import cli
+
+
+class TestDownloadCommand:
+    """Test the download CLI command."""
+
+    def test_export_group_help(self):
+        """Test that the export group shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["export", "--help"])
+        assert result.exit_code == 0
+        assert "Export data from an Opik project" in result.output
+        assert "dataset" in result.output
+        assert "traces" in result.output
+        assert "experiment" in result.output
+
+    def test_export_dataset_help(self):
+        """Test that the export dataset command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["export", "default", "proj", "dataset", "--help"])
+        assert result.exit_code == 0
+        assert "Export a dataset by exact name" in result.output
+        assert "--force" in result.output
+
+    def test_export_traces_help(self):
+        """Test that the export traces command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["export", "default", "proj", "traces", "--help"])
+        assert result.exit_code == 0
+        assert "Export the project's traces" in result.output
+
+    def test_export_experiment_help(self):
+        """Test that the export experiment command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(
+            cli, ["export", "default", "proj", "experiment", "--help"]
+        )
+        assert result.exit_code == 0
+        assert "Export an experiment by exact name" in result.output
+
+
+class TestUploadCommand:
+    """Test the upload CLI command."""
+
+    def test_import_group_help(self):
+        """Test that the import group shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["import", "--help"])
+        assert result.exit_code == 0
+        assert "Import data into an Opik project" in result.output
+        assert "dataset" in result.output
+        assert "traces" in result.output
+        assert "experiment" in result.output
+
+    def test_import_dataset_help(self):
+        """Test that the import dataset command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["import", "default", "proj", "dataset", "--help"])
+        assert result.exit_code == 0
+        assert "Import datasets from" in result.output
+        assert "--dry-run" in result.output
+
+    def test_import_traces_help(self):
+        """Test that the import traces command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["import", "default", "proj", "traces", "--help"])
+        assert result.exit_code == 0
+        assert "Import the project's traces" in result.output
+        assert "--dry-run" in result.output
+
+    def test_import_experiment_help(self):
+        """Test that the import experiment command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(
+            cli, ["import", "default", "proj", "experiment", "--help"]
+        )
+        assert result.exit_code == 0
+        assert "Import experiments from" in result.output
+        assert "--dry-run" in result.output
+
+
+class TestSmokeTestCommand:
+    """Test the smoke-test functionality via healthcheck command."""
+
+    def test_smoke_test_help(self):
+        """Test that the healthcheck --smoke-test command shows help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["healthcheck", "--help"])
+        assert result.exit_code == 0
+        assert "--smoke-test" in result.output
+        assert "--project-name" in result.output
+        assert "Project name for the smoke test" in result.output
+        assert "WORKSPACE" in result.output
+        assert "Run a smoke test to verify Opik integration" in result.output
+
+    def test_smoke_test_minimal_args_parsing(self):
+        """Test that healthcheck --smoke-test command requires workspace value."""
+        runner = CliRunner()
+        # Test that help shows workspace is required
+        result = runner.invoke(cli, ["healthcheck", "--help"])
+        assert result.exit_code == 0
+        assert "WORKSPACE" in result.output
+        # Test that missing workspace value causes error
+        result = runner.invoke(cli, ["healthcheck", "--smoke-test"])
+        assert result.exit_code != 0
+        assert "Error" in result.output or "Missing" in result.output
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.smoke_test.opik.Opik")
+    @patch("opik.cli.healthcheck.smoke_test.opik.start_as_current_trace")
+    @patch("opik.cli.healthcheck.smoke_test.opik_context.update_current_trace")
+    @patch("opik.cli.healthcheck.smoke_test.opik_context.update_current_span")
+    @patch("opik.cli.healthcheck.smoke_test.create_opik_logo_image")
+    @patch("opik.cli.healthcheck.smoke_test.track")
+    def test_smoke_test_with_workspace_and_project_name(
+        self,
+        mock_track,
+        mock_create_logo,
+        mock_update_span,
+        mock_update_trace,
+        mock_start_trace,
+        mock_opik_class,
+        mock_healthcheck_run,
+    ):
+        """Test healthcheck --smoke-test command with workspace and --project-name arguments."""
+        # Setup mocks
+        mock_client = MagicMock()
+        # Mock search_traces to raise an exception immediately to skip verification polling
+        # This prevents the verification from running and triggering real client creation
+        mock_client.search_traces = MagicMock(
+            side_effect=AttributeError("Mock client - search_traces not available")
+        )
+        mock_opik_class.return_value = mock_client
+
+        # Mock the context manager for start_as_current_trace
+        @contextmanager
+        def mock_trace_context(*args, **kwargs):
+            yield MagicMock()
+
+        mock_start_trace.return_value = mock_trace_context()
+
+        # Mock create_opik_logo_image to return a fake path
+        mock_logo_path = Path("/tmp/fake_logo.png")
+        mock_create_logo.return_value = mock_logo_path
+
+        # Mock the track decorator to return the function unchanged
+        def track_decorator(*args, **kwargs):
+            def decorator(func):
+                return func
+
+            return decorator
+
+        mock_track.side_effect = track_decorator
+
+        # Run the command
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            [
+                "healthcheck",
+                "--smoke-test",
+                "test-workspace",
+                "--project-name",
+                "test-project",
+            ],
+            catch_exceptions=False,
+        )
+
+        # Assertions
+        assert result.exit_code == 0
+        # Verify client was created with correct arguments
+        mock_opik_class.assert_called_once()
+        call_kwargs = mock_opik_class.call_args[1]
+        assert call_kwargs["workspace"] == "test-workspace"
+        assert call_kwargs["project_name"] == "test-project"
+        # Verify trace was started
+        mock_start_trace.assert_called_once()
+        # Verify explicit client was flushed and ended
+        mock_client.flush.assert_called_once()
+        mock_client.end.assert_called_once()
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.smoke_test.opik.Opik")
+    @patch("opik.cli.healthcheck.smoke_test.opik.start_as_current_trace")
+    @patch("opik.cli.healthcheck.smoke_test.opik_context.update_current_trace")
+    @patch("opik.cli.healthcheck.smoke_test.opik_context.update_current_span")
+    @patch("opik.cli.healthcheck.smoke_test.create_opik_logo_image")
+    @patch("opik.cli.healthcheck.smoke_test.track")
+    def test_smoke_test_with_default_project_name(
+        self,
+        mock_track,
+        mock_create_logo,
+        mock_update_span,
+        mock_update_trace,
+        mock_start_trace,
+        mock_opik_class,
+        mock_healthcheck_run,
+    ):
+        """Test healthcheck --smoke-test command with workspace only (uses default project name)."""
+        # Setup mocks
+        mock_client = MagicMock()
+        # Mock search_traces to raise an exception immediately to skip verification polling
+        # This prevents the verification from running and triggering real client creation
+        mock_client.search_traces = MagicMock(
+            side_effect=AttributeError("Mock client - search_traces not available")
+        )
+        mock_opik_class.return_value = mock_client
+
+        # Mock the context manager for start_as_current_trace
+        @contextmanager
+        def mock_trace_context(*args, **kwargs):
+            yield MagicMock()
+
+        mock_start_trace.return_value = mock_trace_context()
+
+        # Mock create_opik_logo_image to return a fake path
+        mock_logo_path = Path("/tmp/fake_logo.png")
+        mock_create_logo.return_value = mock_logo_path
+
+        # Mock the track decorator to return the function unchanged
+        def track_decorator(*args, **kwargs):
+            def decorator(func):
+                return func
+
+            return decorator
+
+        mock_track.side_effect = track_decorator
+
+        # Run the command without --project-name (should use default)
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["healthcheck", "--smoke-test", "test-workspace"],
+            catch_exceptions=False,
+        )
+
+        # Assertions
+        assert result.exit_code == 0
+        # Verify client was created with default project name
+        mock_opik_class.assert_called_once()
+        call_kwargs = mock_opik_class.call_args[1]
+        assert call_kwargs["workspace"] == "test-workspace"
+        assert call_kwargs["project_name"] == "smoke-test-project"  # Default value
+        # Verify trace was started with default project name
+        mock_start_trace.assert_called_once()
+        trace_call_kwargs = mock_start_trace.call_args[1]
+        assert trace_call_kwargs["project_name"] == "smoke-test-project"
+
+
+class TestCheckPermissionsCommand:
+    """Test the --check-permissions functionality via healthcheck command."""
+
+    def test_check_permissions_option__in_healthcheck_help__shown(self):
+        """Test that --check-permissions option is shown in healthcheck help."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["healthcheck", "--help"])
+        assert result.exit_code == 0
+        assert "--check-permissions" in result.output
+        assert "WORKSPACE" in result.output
+
+    def test_check_permissions__missing_workspace__raises_error(self):
+        """Test that --check-permissions without a workspace value raises an error."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["healthcheck", "--check-permissions"])
+        assert result.exit_code != 0
+        assert "Error" in result.output or "Missing" in result.output
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.cli.check_user_permissions.run")
+    def test_check_permissions__workspace_provided__calls_run_with_workspace(
+        self,
+        mock_permissions_run,
+        mock_standard_run,
+    ):
+        """Test that --check-permissions pass the workspace to check_user_permissions.run."""
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["healthcheck", "--check-permissions", "my-workspace"],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_permissions_run.assert_called_once_with(
+            api_key=None, workspace="my-workspace"
+        )
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.cli.check_user_permissions.run")
+    def test_check_permissions__api_key_in_cli_context__passes_api_key_to_run(
+        self,
+        mock_permissions_run,
+        mock_standard_run,
+    ):
+        """Test that --check-permissions forward the api_key from CLI context."""
+        from opik.cli import cli as opik_cli
+
+        runner = CliRunner()
+        result = runner.invoke(
+            opik_cli,
+            [
+                "--api-key",
+                "test-api-key",
+                "healthcheck",
+                "--check-permissions",
+                "my-workspace",
+            ],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        mock_permissions_run.assert_called_once_with(
+            api_key="test-api-key", workspace="my-workspace"
+        )
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.cli.check_user_permissions.run")
+    def test_check_permissions__empty_workspace__raises_bad_parameter(
+        self,
+        mock_permissions_run,
+        mock_standard_run,
+    ):
+        """Test that --check-permissions with an empty workspace string raises BadParameter."""
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["healthcheck", "--check-permissions", ""],
+        )
+
+        assert result.exit_code != 0
+        assert (
+            "Error" in result.output
+            or "requires a non-empty workspace name" in result.output
+        )
+        mock_permissions_run.assert_not_called()
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.cli.check_user_permissions.run")
+    def test_check_permissions__no_flag__not_called(
+        self,
+        mock_permissions_run,
+        mock_standard_run,
+    ):
+        """Test that check_user_permissions.run is not called when --check-permissions is absent."""
+        runner = CliRunner()
+        result = runner.invoke(cli, ["healthcheck"], catch_exceptions=False)
+
+        assert result.exit_code == 0
+        mock_permissions_run.assert_not_called()
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.check_user_permissions.get_user_permissions")
+    @patch("opik.cli.healthcheck.check_user_permissions.config.OpikConfig")
+    def test_check_permissions__api_returns_user_and_workspace__displays_info(
+        self,
+        mock_opik_config,
+        mock_get_permissions,
+        mock_standard_run,
+    ):
+        """Test that user and workspace info from API response are printed."""
+        mock_config = MagicMock()
+        mock_config.api_key = "test-api-key"
+        mock_config.workspace = "my-workspace"
+        mock_config.url_override = "https://opik.example.com"
+        mock_opik_config.return_value = mock_config
+
+        mock_get_permissions.return_value = {
+            "user_name": "alice",
+            "workspace_name": "my-workspace",
+            "permissions": [
+                {"permission_name": "read", "permission_value": True},
+                {"permission_name": "write", "permission_value": False},
+            ],
+        }
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["healthcheck", "--check-permissions", "my-workspace"],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        assert "alice" in result.output
+        assert "my-workspace" in result.output
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.check_user_permissions.get_user_permissions")
+    @patch("opik.cli.healthcheck.check_user_permissions.config.OpikConfig")
+    def test_check_permissions__connection_error__prints_error(
+        self,
+        mock_opik_config,
+        mock_get_permissions,
+        mock_standard_run,
+    ):
+        """Test that a ConnectionError from the API prints an error message."""
+        mock_config = MagicMock()
+        mock_config.api_key = "test-api-key"
+        mock_config.workspace = "my-workspace"
+        mock_config.url_override = "https://opik.example.com"
+        mock_opik_config.return_value = mock_config
+
+        mock_get_permissions.side_effect = ConnectionError("Network error: timeout")
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["healthcheck", "--check-permissions", "my-workspace"],
+            catch_exceptions=False,
+        )
+
+        assert result.exit_code == 0
+        assert "Failed to fetch user permissions" in result.output
+
+    @patch("opik.cli.healthcheck.cli.standard_check.run")
+    @patch("opik.cli.healthcheck.check_user_permissions.config.OpikConfig")
+    def test_check_permissions__missing_api_key__raises_value_error(
+        self,
+        mock_opik_config,
+        mock_standard_run,
+    ):
+        """Test that missing api_key (not in CLI context and not in config) raises ValueError."""
+        mock_config = MagicMock()
+        mock_config.api_key = None
+        mock_config.workspace = None
+        mock_config.url_override = "https://opik.example.com"
+        mock_opik_config.return_value = mock_config
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["healthcheck", "--check-permissions", "my-workspace"],
+        )
+
+        assert result.exit_code != 0
+        assert isinstance(result.exception, ValueError)
+        assert "API key is required" in str(result.exception)

sdks/python/tests/unit/cli/test_configure_cli.py

@@ -0,0 +1,54 @@
+"""Tests for the ``opik configure`` command group."""
+
+import pathlib
+from unittest import mock
+
+from click.testing import CliRunner
+
+from opik.cli import cli
+from opik.cli import configure as configure_cli
+
+
+def test_configure_status__prints_config_summary():
+    runner = CliRunner()
+    config = mock.Mock(
+        config_file_exists=True,
+        config_file_fullpath=pathlib.Path("/home/u/.opik.config"),
+        url_override="https://dev.comet.com/opik/api/",
+        workspace="my-ws",
+    )
+    with mock.patch.object(
+        configure_cli.opik_config, "OpikConfig", return_value=config
+    ):
+        result = runner.invoke(cli, ["configure", "status"])
+
+    assert result.exit_code == 0
+    assert "Your Opik configuration" in result.output
+    assert "https://dev.comet.com/opik/api/" in result.output
+    assert "my-ws" in result.output
+
+
+def test_configure_status__not_configured__points_to_configure():
+    runner = CliRunner()
+    config = mock.Mock(
+        config_file_exists=False,
+        config_file_fullpath=pathlib.Path("/home/u/.opik.config"),
+    )
+    with mock.patch.object(
+        configure_cli.opik_config, "OpikConfig", return_value=config
+    ):
+        result = runner.invoke(cli, ["configure", "status"])
+
+    assert result.exit_code == 0
+    assert "not found" in result.output
+    assert "opik configure" in result.output
+
+
+def test_configure_no_subcommand__runs_configurator():
+    runner = CliRunner()
+    with mock.patch.object(configure_cli, "run_interactive_configure") as spy:
+        result = runner.invoke(cli, ["configure", "--use-local"])
+
+    assert result.exit_code == 0
+    spy.assert_called_once()
+    assert spy.call_args.kwargs["use_local"] is True

sdks/python/tests/unit/cli/test_connect.py

@@ -0,0 +1,230 @@
+from unittest.mock import MagicMock, patch
+
+import httpx
+from click.testing import CliRunner
+
+from opik.cli.main import cli
+from opik.cli.local_runner.pairing import PairingResult, RunnerType
+from opik.rest_api.core.api_error import ApiError
+
+
+class TestConnect:
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__with_project__calls_pairing_and_supervisor(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-abc",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(cli, ["connect", "--project", "my-proj"])
+        assert result.exit_code == 0
+
+        mock_run_pairing.assert_called_once()
+        call_kwargs = mock_run_pairing.call_args[1]
+        assert call_kwargs["project_name"] == "my-proj"
+        assert call_kwargs["runner_type"] == RunnerType.CONNECT
+
+        mock_launch.assert_called_once()
+        launch_kwargs = mock_launch.call_args[1]
+        assert launch_kwargs["command"] is None
+
+    def test_connect__no_project__shows_error(self):
+        runner = CliRunner()
+        # Any arg flips the `_EndpointGroup`-style fallback into the hidden _run
+        # subcommand, where `--project` is enforced as required by Click. Bare
+        # `opik connect` prints group help instead; covered by the next test.
+        result = runner.invoke(cli, ["connect", "--name", "test-runner"])
+        assert result.exit_code == 2
+        assert "--project" in result.output
+
+    def test_connect__no_args__shows_help(self):
+        runner = CliRunner()
+        result = runner.invoke(cli, ["connect"])
+        assert result.exit_code == 0
+        # Help body lists `stop` and the docstring mentions `--project`.
+        assert "stop" in result.output
+        assert "--project" in result.output
+
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__network_failure__shows_clean_error(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.side_effect = httpx.ConnectError("Connection refused")
+
+        runner = CliRunner()
+        result = runner.invoke(cli, ["connect", "--project", "my-proj"])
+        assert result.exit_code != 0
+        assert "Could not connect to Opik backend" in result.output
+        assert "https://api.test/" in result.output
+        assert "Run: opik configure" in result.output
+
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__api_error__shows_error_body(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.side_effect = ApiError(status_code=409, body="conflict")
+
+        runner = CliRunner()
+        result = runner.invoke(cli, ["connect", "--project", "my-proj"])
+        assert result.exit_code != 0
+        assert "conflict" in result.output
+
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__tui_stopped_on_pairing_failure(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.side_effect = ApiError(status_code=500, body="boom")
+
+        runner = CliRunner()
+        runner.invoke(cli, ["connect", "--project", "my-proj"])
+
+        tui_instance = mock_tui_cls.return_value
+        tui_instance.stop.assert_called_once()
+
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__workspace_and_api_key_passed__forwarded_to_opik_constructor(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        client.config.config_file_exists = True
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-abc",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            [
+                "connect",
+                "--project",
+                "my-proj",
+                "--workspace",
+                "my-ws",
+                "--api-key",
+                "my-key",
+            ],
+        )
+        assert result.exit_code == 0
+
+        mock_opik_cls.assert_called_once_with(
+            project_name="my-proj",
+            api_key="my-key",
+            workspace="my-ws",
+            _show_misconfiguration_message=False,
+        )
+
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__local_api_key_overrides_global(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        client.config.config_file_exists = True
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-abc",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            [
+                "--api-key",
+                "global-key",
+                "connect",
+                "--project",
+                "my-proj",
+                "--api-key",
+                "local-key",
+            ],
+        )
+        assert result.exit_code == 0
+
+        mock_opik_cls.assert_called_once_with(
+            project_name="my-proj",
+            api_key="local-key",
+            workspace=None,
+            _show_misconfiguration_message=False,
+        )
+
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_connect__no_local_api_key__falls_back_to_global(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        client.config.config_file_exists = True
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-abc",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["--api-key", "global-key", "connect", "--project", "my-proj"],
+        )
+        assert result.exit_code == 0
+
+        mock_opik_cls.assert_called_once_with(
+            project_name="my-proj",
+            api_key="global-key",
+            workspace=None,
+            _show_misconfiguration_message=False,
+        )

sdks/python/tests/unit/cli/test_endpoint.py

@@ -0,0 +1,240 @@
+from unittest.mock import MagicMock, patch
+
+import httpx
+from click.testing import CliRunner
+
+from opik.cli.main import cli
+from opik.cli.local_runner.pairing import PairingResult, RunnerType
+from opik.rest_api.core.api_error import ApiError
+
+
+class TestEndpoint:
+    @patch("opik.runner.snapshot.has_entrypoint", return_value=True)
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_endpoint__with_command__calls_pairing_and_supervisor(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls, _mock_ep
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-xyz",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli, ["endpoint", "--project", "my-proj", "--", "echo", "hello"]
+        )
+        assert result.exit_code == 0
+
+        mock_run_pairing.assert_called_once()
+        call_kwargs = mock_run_pairing.call_args[1]
+        assert call_kwargs["project_name"] == "my-proj"
+        assert call_kwargs["runner_type"] == RunnerType.ENDPOINT
+
+        mock_launch.assert_called_once()
+        launch_kwargs = mock_launch.call_args[1]
+        assert launch_kwargs["command"] == ["echo", "hello"]
+
+    def test_endpoint__no_project__shows_error(self):
+        runner = CliRunner()
+        result = runner.invoke(cli, ["endpoint", "--", "echo", "hello"])
+        assert result.exit_code == 2
+        assert "--project" in result.output
+
+    def test_endpoint__no_command__shows_error(self):
+        runner = CliRunner()
+        result = runner.invoke(cli, ["endpoint", "--project", "my-proj"])
+        assert result.exit_code == 2
+
+    def test_endpoint__nonexistent_binary__shows_error(self):
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            ["endpoint", "--project", "my-proj", "--", "nonexistent-binary-xyz-12345"],
+        )
+        assert result.exit_code == 2
+        assert "not found" in result.output.lower()
+
+    @patch("opik.runner.snapshot.has_entrypoint", return_value=True)
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_endpoint__network_failure__shows_clean_error(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls, _mock_ep
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.side_effect = httpx.ConnectError("Connection refused")
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli, ["endpoint", "--project", "my-proj", "--", "echo", "hello"]
+        )
+        assert result.exit_code != 0
+        assert "Could not connect to Opik backend" in result.output
+        assert "https://api.test/" in result.output
+        assert "Run: opik configure" in result.output
+
+    @patch("opik.runner.snapshot.has_entrypoint", return_value=True)
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_endpoint__tui_stopped_on_pairing_failure(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls, _mock_ep
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.side_effect = ApiError(status_code=500, body="boom")
+
+        runner = CliRunner()
+        runner.invoke(cli, ["endpoint", "--project", "my-proj", "--", "echo", "hello"])
+
+        tui_instance = mock_tui_cls.return_value
+        tui_instance.stop.assert_called_once()
+
+    @patch("opik.runner.snapshot.has_entrypoint", return_value=True)
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_endpoint__workspace_and_api_key_passed__forwarded_to_opik_constructor(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls, _mock_ep
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        client.config.config_file_exists = True
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-xyz",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            [
+                "endpoint",
+                "--project",
+                "my-proj",
+                "--workspace",
+                "my-ws",
+                "--api-key",
+                "my-key",
+                "--",
+                "echo",
+                "hello",
+            ],
+        )
+        assert result.exit_code == 0
+
+        mock_opik_cls.assert_called_once_with(
+            project_name="my-proj",
+            api_key="my-key",
+            workspace="my-ws",
+            _show_misconfiguration_message=False,
+        )
+
+    @patch("opik.runner.snapshot.has_entrypoint", return_value=True)
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_endpoint__local_api_key_overrides_global(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls, _mock_ep
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        client.config.config_file_exists = True
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-xyz",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            [
+                "--api-key",
+                "global-key",
+                "endpoint",
+                "--project",
+                "my-proj",
+                "--api-key",
+                "local-key",
+                "--",
+                "echo",
+                "hello",
+            ],
+        )
+        assert result.exit_code == 0
+
+        mock_opik_cls.assert_called_once_with(
+            project_name="my-proj",
+            api_key="local-key",
+            workspace=None,
+            _show_misconfiguration_message=False,
+        )
+
+    @patch("opik.runner.snapshot.has_entrypoint", return_value=True)
+    @patch("opik.cli.local_runner._run.RunnerTUI")
+    @patch("opik.cli.local_runner.pairing.launch_supervisor")
+    @patch("opik.cli.local_runner.pairing.run_pairing")
+    @patch("opik.cli.local_runner._run.Opik")
+    def test_endpoint__no_local_api_key__falls_back_to_global(
+        self, mock_opik_cls, mock_run_pairing, mock_launch, mock_tui_cls, _mock_ep
+    ):
+        client = MagicMock()
+        client.config.url_override = "https://api.test/"
+        client.config.config_file_exists = True
+        mock_opik_cls.return_value = client
+
+        mock_run_pairing.return_value = PairingResult(
+            runner_id="r-xyz",
+            project_name="my-proj",
+            project_id="p-123",
+            bridge_key=b"\x00" * 32,
+        )
+
+        runner = CliRunner()
+        result = runner.invoke(
+            cli,
+            [
+                "--api-key",
+                "global-key",
+                "endpoint",
+                "--project",
+                "my-proj",
+                "--",
+                "echo",
+                "hello",
+            ],
+        )
+        assert result.exit_code == 0
+
+        mock_opik_cls.assert_called_once_with(
+            project_name="my-proj",
+            api_key="global-key",
+            workspace=None,
+            _show_misconfiguration_message=False,
+        )

sdks/python/tests/unit/cli/test_export_error_handling.py

@@ -0,0 +1,282 @@
+"""Unit tests for export error-handling fixes.
+
+Covers three HIGH-priority fixes to prevent silent data loss during export:
+
+  HIGH-1  experiment.export_traces_by_ids: when get_trace_content raises 429,
+          the internal _fetch_trace_data retry decorator must kick in so the
+          trace is eventually exported, not silently dropped. Only genuine 404s
+          should be skipped without retrying.
+
+  HIGH-2  experiment.export_experiment_by_id: must re-raise on error so that
+          the all.py caller can set had_errors=True (previously it swallowed
+          all exceptions and returned a zero-stats tuple).
+
+  HIGH-3  all._fetch_experiments_page_raw: 429 must trigger a retry via the
+          _export_rest_retry decorator instead of silently returning ([], 0)
+          and losing an entire page of experiments.
+"""
+
+from unittest.mock import MagicMock, patch
+
+import pytest
+import httpx
+
+from opik.rest_api.core.api_error import ApiError
+from opik import exceptions as opik_exceptions
+
+_EXPERIMENT_MODULE = "opik.cli.exports.experiment"
+_ALL_MODULE = "opik.cli.exports.all"
+
+
+# ---------------------------------------------------------------------------
+# HIGH-1: export_traces_by_ids retries 429 rather than dropping the trace
+# ---------------------------------------------------------------------------
+
+
+class TestExportTracesByIdsRetry:
+    def _make_mock_trace(self, trace_id: str) -> MagicMock:
+        t = MagicMock()
+        t.model_dump.return_value = {"id": trace_id}
+        return t
+
+    def test_export_traces_by_ids__429_is_retried_not_dropped(self, tmp_path):
+        """A 429 ApiError must trigger the retry decorator so the trace is
+        exported, not silently dropped."""
+        from opik.cli.exports.experiment import export_traces_by_ids
+
+        call_count = 0
+
+        def get_trace_side_effect(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise ApiError(status_code=429, headers={})
+            return self._make_mock_trace("t1")
+
+        client = MagicMock()
+        client.get_trace_content.side_effect = get_trace_side_effect
+        client.search_spans.return_value = []
+
+        with patch(f"{_EXPERIMENT_MODULE}._fetch_trace_data.retry.sleep"):
+            exported, skipped = export_traces_by_ids(
+                client, ["t1"], tmp_path, "proj", None, "json", False, False
+            )
+
+        assert exported == 1, "429 must not silently drop the trace"
+        assert call_count == 2, "429 must trigger a retry"
+
+    def test_export_traces_by_ids__rate_limited_exception_is_retried(self, tmp_path):
+        """OpikCloudRequestsRateLimited must also be retried."""
+        from opik.cli.exports.experiment import export_traces_by_ids
+
+        call_count = 0
+
+        def get_trace_side_effect(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise opik_exceptions.OpikCloudRequestsRateLimited(
+                    headers={}, retry_after=1.0
+                )
+            return self._make_mock_trace("t2")
+
+        client = MagicMock()
+        client.get_trace_content.side_effect = get_trace_side_effect
+        client.search_spans.return_value = []
+
+        with patch(f"{_EXPERIMENT_MODULE}._fetch_trace_data.retry.sleep"):
+            exported, skipped = export_traces_by_ids(
+                client, ["t2"], tmp_path, "proj", None, "json", False, False
+            )
+
+        assert exported == 1, "rate-limit exception must not silently drop the trace"
+        assert call_count == 2
+
+    def test_export_traces_by_ids__404_is_skipped_not_retried(self, tmp_path):
+        """A genuine 404 (trace not found) should be skipped cleanly without retrying."""
+        from opik.cli.exports.experiment import export_traces_by_ids
+
+        call_count = 0
+
+        def get_trace_side_effect(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            raise ApiError(status_code=404, headers={})
+
+        client = MagicMock()
+        client.get_trace_content.side_effect = get_trace_side_effect
+
+        exported, skipped = export_traces_by_ids(
+            client, ["missing-id"], tmp_path, "proj", None, "json", False, False
+        )
+
+        assert exported == 0, "404 trace must not be exported"
+        assert call_count == 1, "404 must not trigger a retry"
+
+    def test_export_traces_by_ids__success__exports_trace(self, tmp_path):
+        """Happy path: trace is written to disk and counted as exported."""
+        from opik.cli.exports.experiment import export_traces_by_ids
+
+        client = MagicMock()
+        client.get_trace_content.return_value = self._make_mock_trace("t3")
+        client.search_spans.return_value = []
+
+        exported, skipped = export_traces_by_ids(
+            client, ["t3"], tmp_path, "proj", None, "json", False, False
+        )
+
+        assert exported == 1
+
+
+# ---------------------------------------------------------------------------
+# HIGH-2: export_experiment_by_id propagates errors instead of swallowing them
+# ---------------------------------------------------------------------------
+
+
+class TestExportExperimentByIdPropagatesErrors:
+    def test_export_experiment_by_id__api_error_propagates(self, tmp_path):
+        """An unexpected error must propagate so the caller can set had_errors=True."""
+        from opik.cli.exports.experiment import export_experiment_by_id
+
+        client = MagicMock()
+        # Simulate a non-retriable API error (e.g. 403 Forbidden)
+        client.get_experiment_by_id.side_effect = ApiError(status_code=403, headers={})
+
+        with pytest.raises(ApiError):
+            export_experiment_by_id(
+                client=client,
+                output_dir=tmp_path,
+                project_name="proj",
+                experiment_id="exp-1",
+                max_traces=None,
+                force=True,
+                debug=False,
+                format="json",
+            )
+
+    def test_export_all_experiments__experiment_failure_sets_had_errors(self, tmp_path):
+        """When export_experiment_by_id raises, _export_all_experiments must set
+        had_errors=True rather than silently counting it as skipped."""
+        from opik.cli.exports.all import _export_all_experiments
+
+        client = MagicMock()
+        experiments_dir = tmp_path / "experiments"
+        experiments_dir.mkdir()
+
+        # Two experiments: first succeeds, second raises.
+        exp1 = MagicMock()
+        exp1.id = "e1"
+        exp1.name = "exp-one"
+        exp2 = MagicMock()
+        exp2.id = "e2"
+        exp2.name = "exp-two"
+
+        def export_side_effect(
+            client, output_dir, project_name, exp_id, *args, **kwargs
+        ):
+            if exp_id == "e2":
+                raise RuntimeError("transient failure")
+            m = MagicMock()
+            m.get_all_trace_ids.return_value = []
+            return {}, 1, m
+
+        with (
+            patch(
+                f"{_ALL_MODULE}._paginate_experiments",
+                return_value=[exp1, exp2],
+            ),
+            patch(
+                f"{_ALL_MODULE}.export_experiment_by_id",
+                side_effect=export_side_effect,
+            ),
+            patch(f"{_ALL_MODULE}.export_collected_trace_ids", return_value=(0, 0)),
+        ):
+            _, _, _, _, had_errors = _export_all_experiments(
+                client=client,
+                project_dir=tmp_path,
+                project_name="proj",
+                project_id=None,
+                experiments_dir=experiments_dir,
+                max_results=None,
+                force=False,
+                debug=False,
+                format="json",
+            )
+
+        assert had_errors is True, "a failed experiment must set had_errors=True"
+
+
+# ---------------------------------------------------------------------------
+# HIGH-3: _fetch_experiments_page_raw retries 429 instead of returning [], 0
+# ---------------------------------------------------------------------------
+
+
+class TestFetchExperimentsPageRawRetry:
+    def _make_httpx_response(self, status_code: int, headers=None, json_body=None):
+        """Build a minimal httpx.Response for raise_for_status() tests."""
+        headers = headers or {}
+        body = json_body or {"content": [], "total": 0}
+        import json as _json
+
+        # httpx.Response.raise_for_status() requires a request to be attached.
+        request = httpx.Request("GET", "http://test/v1/private/experiments")
+        response = httpx.Response(
+            status_code=status_code,
+            headers=headers,
+            content=_json.dumps(body).encode(),
+            request=request,
+        )
+        return response
+
+    def test_fetch_experiments_page_raw__429_raises_rate_limited_not_empty_list(
+        self,
+    ):
+        """A 429 response must raise OpikCloudRequestsRateLimited (for the retry
+        decorator), NOT silently return ([], 0)."""
+        from opik.cli.exports.all import _fetch_experiments_page_raw
+
+        client = MagicMock()
+        httpx_mock = MagicMock()
+        client.rest_client.experiments._raw_client._client_wrapper.httpx_client = (
+            httpx_mock
+        )
+        httpx_mock.request.return_value = self._make_httpx_response(
+            429, headers={"Retry-After": "5"}
+        )
+
+        with pytest.raises(opik_exceptions.OpikCloudRequestsRateLimited):
+            _fetch_experiments_page_raw.__wrapped__(client, 1, None)
+
+    def test_fetch_experiments_page_raw__429_is_retried_by_decorator(self):
+        """When the underlying call raises 429, the decorator retries rather
+        than letting a silent ([], 0) escape to the caller."""
+        from opik.cli.exports.all import _fetch_experiments_page_raw
+
+        call_count = 0
+
+        def request_side_effect(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                return self._make_httpx_response(429, headers={"Retry-After": "1"})
+            return self._make_httpx_response(
+                200,
+                json_body={
+                    "content": [{"id": "exp-1", "name": "my-exp"}],
+                    "total": 1,
+                },
+            )
+
+        client = MagicMock()
+        httpx_mock = MagicMock()
+        client.rest_client.experiments._raw_client._client_wrapper.httpx_client = (
+            httpx_mock
+        )
+        httpx_mock.request.side_effect = request_side_effect
+
+        with patch(f"{_ALL_MODULE}._fetch_experiments_page_raw.retry.sleep"):
+            items, total = _fetch_experiments_page_raw(client, 1, None)
+
+        assert call_count == 2, "429 must trigger a retry"
+        assert len(items) == 1, "successful retry must return the page contents"
+        assert total == 1

sdks/python/tests/unit/cli/test_export_project_rate_limiting.py

@@ -0,0 +1,788 @@
+"""Tests for rate-limiting resilience in project export.
+
+Covers the changes introduced to handle server-side 429 throttling:
+
+  Retry wait behaviour (via export_traces public API)
+    - 429 with Retry-After header → sleep honours header value
+    - 429 with Retry-After exceeding cap → sleep falls back to 30 s
+    - 429 with no Retry-After → sleep is 30 s
+    - non-429 transient error → sleep is in exponential backoff range [2, 60] s
+
+  export_traces (page-fetch failures)
+    - 429 on page N → page is skipped, remaining pages still fetched, had_errors=True
+    - all pages succeed → had_errors=False, all traces exported
+    - middle page failure → subsequent pages still exported
+
+  export_traces (span-fetch failures)
+    - bulk span fetch raises repeatedly → traces exported with empty spans, had_errors=True
+
+  export_traces (inter-page delay)
+    - time.sleep NOT called between successful full pages (inter-page sleep removed)
+    - time.sleep not called after partial (terminal) page
+
+  export_single_project (had_errors propagation)
+    - returns 4-tuple; had_errors=False on clean run, True when errors occurred
+    - first element is 0 when no traces exported or skipped, 1 otherwise
+"""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from opik.rest_api.core.api_error import ApiError
+
+_MODULE = "opik.cli.exports.project"
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_mock_trace(trace_id: str) -> MagicMock:
+    t = MagicMock()
+    t.id = trace_id
+    t.name = f"trace-{trace_id}"
+    t.model_dump.return_value = {
+        "id": trace_id,
+        "name": f"trace-{trace_id}",
+        "start_time": None,
+        "end_time": None,
+        "input": {},
+        "output": {},
+        "metadata": {},
+        "tags": [],
+        "feedback_scores": [],
+        "error_info": None,
+        "thread_id": None,
+        "created_at": None,
+        "created_by": None,
+        "last_updated_at": None,
+        "last_updated_by": None,
+        "visibility_mode": None,
+        "ttft": None,
+        "project_name": "test-project",
+    }
+    return t
+
+
+def _make_page(traces, total=None):
+    page = MagicMock()
+    page.content = traces
+    page.total = total if total is not None else len(traces)
+    return page
+
+
+def _make_full_page(n=500, offset=0):
+    """Return a page with exactly *n* mock traces (matches the internal page_size=500).
+
+    Use *offset* to generate distinct trace IDs across pages.
+    """
+    return _make_page([_make_mock_trace(f"bulk-{offset + i}") for i in range(n)])
+
+
+def _uuid7(n: int) -> str:
+    """Deterministic, valid UUIDv7 string for trace index *n*.
+
+    Version nibble is 7 and the variant nibble is RFC-4122 (8), so
+    ``uuid.UUID(...).version == 7`` — this is what makes the exporter build a
+    bounded ``trace_id`` range filter instead of falling back to a full scan.
+    """
+    high = f"{0x0190A1B2C3D4 + n:012x}"
+    return f"{high[:8]}-{high[8:12]}-7000-8000-{n:012x}"
+
+
+def _make_full_page_uuid7(n=500, offset=0):
+    """Like ``_make_full_page`` but with UUIDv7 trace ids (bounded-scan path)."""
+    return _make_page([_make_mock_trace(_uuid7(offset + i)) for i in range(n)])
+
+
+# ---------------------------------------------------------------------------
+# Retry wait behaviour — tested via export_traces (public API)
+# ---------------------------------------------------------------------------
+
+
+class TestExportTracesRetryWaitBehaviour:
+    """Verify that the export-specific retry decorator applies the correct wait
+    durations when the underlying page-fetch raises errors.  Tests go through
+    the public export_traces entrypoint so the actual retry decorator fires;
+    time.sleep is patched to capture what tenacity passes to it.
+    """
+
+    def _run_with_first_call_raising(self, exc, tmp_path):
+        """Run export_traces where the first page fetch raises *exc* then succeeds."""
+        from opik.cli.exports.project import export_traces
+
+        trace = _make_mock_trace("t1")
+        call_count = 0
+
+        def get_traces_side_effect(*args, **kwargs):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                raise exc
+            return _make_page([trace])
+
+        mock_client = MagicMock()
+        mock_client.rest_client.traces.get_traces_by_project.side_effect = (
+            get_traces_side_effect
+        )
+
+        with patch(f"{_MODULE}._fetch_traces_page.retry.sleep") as mock_sleep:
+            export_traces(
+                client=mock_client,
+                project_name="proj",
+                project_dir=tmp_path,
+                max_results=None,
+                filter_string=None,
+                show_progress=False,
+            )
+
+        return [c.args[0] for c in mock_sleep.call_args_list]
+
+    def test_export_traces__page_fetch_429_with_retry_after_header__sleep_honours_header(
+        self, tmp_path
+    ):
+        exc = ApiError(status_code=429, headers={"retry-after": "45"})
+        sleep_calls = self._run_with_first_call_raising(exc, tmp_path)
+        # Jitter adds 0–5 s on top of the header value.
+        assert any(45.0 <= s <= 50.0 for s in sleep_calls)
+
+    def test_export_traces__page_fetch_429_retry_after_exceeds_cap__sleep_clamped_to_max(
+        self, tmp_path
+    ):
+        # 200 s > _EXPORT_MAX_RETRY_AFTER_SECONDS (120 s) → clamped to 120 s
+        exc = ApiError(status_code=429, headers={"retry-after": "200"})
+        sleep_calls = self._run_with_first_call_raising(exc, tmp_path)
+        # Jitter adds 0–5 s on top of the 120 s cap.
+        assert any(120.0 <= s <= 125.0 for s in sleep_calls)
+
+    def test_export_traces__page_fetch_429_retry_after_http_date__sleep_honours_header(
+        self, tmp_path
+    ):
+        import email.utils
+        import time
+
+        # Build an HTTP-date Retry-After 60 seconds in the future.
+        future = email.utils.formatdate(timeval=time.time() + 60, usegmt=True)
+        exc = ApiError(status_code=429, headers={"retry-after": future})
+        sleep_calls = self._run_with_first_call_raising(exc, tmp_path)
+        # Should sleep approximately 60 s (allow ±2 s for execution) plus 0–5 s jitter.
+        assert any(58.0 <= s <= 67.0 for s in sleep_calls)
+
+    def test_export_traces__page_fetch_429_no_retry_after_header__sleep_is_30s(
+        self, tmp_path
+    ):
+        exc = ApiError(status_code=429, headers={})
+        sleep_calls = self._run_with_first_call_raising(exc, tmp_path)
+        # Jitter adds 0–5 s on top of the 30 s fallback.
+        assert any(30.0 <= s <= 35.0 for s in sleep_calls)
+
+    def test_export_traces__page_fetch_non_429_transient_error__sleep_is_exponential_backoff(
+        self, tmp_path
+    ):
+        exc = ApiError(status_code=503)
+        sleep_calls = self._run_with_first_call_raising(exc, tmp_path)
+        # Backoff is exponential starting at 2 s, capped at 60 s.
+        assert any(2.0 <= s <= 60.0 for s in sleep_calls)
+
+
+# ---------------------------------------------------------------------------
+# export_traces — page-fetch failures
+# ---------------------------------------------------------------------------
+
+
+class TestExportTracesPageFetchFailures:
+    def test_export_traces__first_page_429__skips_page_and_continues(self):
+        from opik.cli.exports.project import export_traces
+
+        trace_p2 = _make_mock_trace("t-p2")
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                mock_fetch.side_effect = [
+                    ApiError(status_code=429, body="rate limited"),
+                    _make_page([trace_p2]),
+                    _make_page([]),
+                ]
+                with patch(f"{_MODULE}.time.sleep"):
+                    exported, skipped, had_errors = export_traces(
+                        client=mock_client,
+                        project_name="proj",
+                        project_dir=Path(tmp),
+                        max_results=None,
+                        filter_string=None,
+                        show_progress=False,
+                    )
+
+        assert had_errors is True
+        assert exported == 1
+        assert skipped == 0
+
+    def test_export_traces__all_pages_succeed__happyflow(self):
+        from opik.cli.exports.project import export_traces
+
+        traces = [_make_mock_trace("t1"), _make_mock_trace("t2")]
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                mock_fetch.side_effect = [
+                    _make_page(traces),
+                    _make_page([]),
+                ]
+                with patch(f"{_MODULE}.time.sleep"):
+                    exported, skipped, had_errors = export_traces(
+                        client=mock_client,
+                        project_name="proj",
+                        project_dir=Path(tmp),
+                        max_results=None,
+                        filter_string=None,
+                        show_progress=False,
+                    )
+
+        assert had_errors is False
+        assert exported == 2
+
+    def test_export_traces__middle_page_fails__subsequent_pages_still_exported(self):
+        """A failed page in the middle must not prevent later pages from being fetched."""
+        from opik.cli.exports.project import export_traces
+
+        # Use page_size=2 so pages with 2 traces are "full" (2 == page_size).
+        # Keeps file-write count tiny (4 files) while still exercising the
+        # pagination loop — the loop only breaks early on a *short* page.
+        page1 = _make_full_page(2, offset=0)
+        page3 = _make_full_page(2, offset=2)
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                mock_fetch.side_effect = [
+                    page1,  # page 1: OK, 2 traces
+                    ApiError(status_code=429),  # page 2: rate limited
+                    page3,  # page 3: OK, 2 traces
+                    _make_page([]),  # page 4: end
+                ]
+                with patch(f"{_MODULE}._fetch_spans_page", return_value=_make_page([])):
+                    with patch(f"{_MODULE}.time.sleep"):
+                        exported, skipped, had_errors = export_traces(
+                            client=mock_client,
+                            project_name="proj",
+                            project_dir=Path(tmp),
+                            max_results=None,
+                            filter_string=None,
+                            show_progress=False,
+                            page_size=2,
+                        )
+
+        assert had_errors is True
+        assert exported == 4  # page 1 and page 3 traces both exported
+
+    def test_export_traces__consecutive_failures_reach_cap__export_aborted(self):
+        """After MAX_CONSECUTIVE_PAGE_FAILURES pages in a row fail, the loop aborts."""
+        from opik.cli.exports.project import (
+            export_traces,
+            MAX_CONSECUTIVE_PAGE_FAILURES,
+        )
+
+        mock_client = MagicMock()
+
+        # Every page raises a transient error
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                mock_fetch.side_effect = ApiError(status_code=503)
+                with patch(f"{_MODULE}.time.sleep"):
+                    exported, skipped, had_errors = export_traces(
+                        client=mock_client,
+                        project_name="proj",
+                        project_dir=Path(tmp),
+                        max_results=None,
+                        filter_string=None,
+                        show_progress=False,
+                    )
+
+        assert had_errors is True
+        assert exported == 0
+        # Should have stopped after MAX_CONSECUTIVE_PAGE_FAILURES attempts
+        assert mock_fetch.call_count == MAX_CONSECUTIVE_PAGE_FAILURES
+
+    def test_export_traces__permanent_api_error__raises_instead_of_skipping(self):
+        """A 400 Bad Request must not be silently skipped — it should propagate."""
+        from opik.cli.exports.project import export_traces
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                mock_fetch.side_effect = ApiError(status_code=400, body="bad request")
+                with patch(f"{_MODULE}.time.sleep"):
+                    with pytest.raises(ApiError) as exc_info:
+                        export_traces(
+                            client=mock_client,
+                            project_name="proj",
+                            project_dir=Path(tmp),
+                            max_results=None,
+                            filter_string=None,
+                            show_progress=False,
+                        )
+
+        assert exc_info.value.status_code == 400
+        # Should have raised immediately, not retried
+        assert mock_fetch.call_count == 1
+
+
+# ---------------------------------------------------------------------------
+# export_traces — span-fetch failures
+# ---------------------------------------------------------------------------
+
+
+class TestExportTracesSpanFetchFailures:
+    def test_export_traces__bulk_span_fetch_fails__traces_exported_with_empty_spans(
+        self,
+    ):
+        """When the bulk span page fetch keeps failing, traces are still written
+        (with empty span lists) and had_errors is set to True."""
+        from opik.cli.exports.project import export_traces
+
+        t1 = _make_mock_trace("t1")
+        t2 = _make_mock_trace("t2")
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_trace_fetch:
+                mock_trace_fetch.side_effect = [
+                    _make_page([t1, t2]),
+                    _make_page([]),
+                ]
+                with patch(f"{_MODULE}._fetch_spans_page") as mock_span_fetch:
+                    mock_span_fetch.side_effect = ApiError(status_code=503)
+                    with patch(f"{_MODULE}.time.sleep"):
+                        exported, skipped, had_errors = export_traces(
+                            client=mock_client,
+                            project_name="proj",
+                            project_dir=Path(tmp),
+                            max_results=None,
+                            filter_string=None,
+                            show_progress=False,
+                        )
+
+        assert had_errors is True
+        assert exported == 2  # both trace files written, just without spans
+        assert skipped == 0
+
+
+# ---------------------------------------------------------------------------
+# export_traces — bounded-chunk (memory) behaviour
+# ---------------------------------------------------------------------------
+
+
+class TestExportTracesChunking:
+    def test_export_traces__large_project_flushes_in_bounded_chunks(self):
+        """A multi-page export writes traces incrementally, one bounded chunk at
+        a time, instead of buffering every trace and span before the first write.
+        Incremental flushing is what keeps peak memory from scaling with the
+        whole project (the previous all-at-once design OOM'd on large projects).
+
+        The export is driven purely through the public ``export_traces`` API,
+        mocking only the public REST client boundary
+        (``traces.get_traces_by_project`` / ``spans.get_spans_by_project``) — no
+        private helpers are patched.  Correctness is asserted on the returned
+        tuple and the final ``trace_*.json`` artifacts.  The bounded-chunk
+        property is asserted through robust invariants observed at the public
+        span endpoint (files already on disk when each span fetch fires), not a
+        brittle exact-call-order snapshot: the assertions hold regardless of the
+        exact chunk boundary the implementation picks.
+        """
+        from opik.cli.exports.project import export_traces
+
+        # UUIDv7 ids so each chunk builds a bounded trace_id range filter and
+        # scans spans per chunk (the common, memory-bounded path).
+        page1 = _make_full_page_uuid7(2, offset=0)
+        page2 = _make_full_page_uuid7(2, offset=2)
+        page3 = _make_full_page_uuid7(2, offset=4)
+        mock_client = MagicMock()
+        expected_ids = [_uuid7(i) for i in range(6)]
+
+        with tempfile.TemporaryDirectory() as tmp:
+            project_dir = Path(tmp)
+            files_on_disk_at_each_span_fetch = []
+
+            def span_fetch_side_effect(*args, **kwargs):
+                # How many trace files already exist when a span fetch begins.
+                # A growing, sub-total count proves earlier chunks were flushed
+                # to disk before later chunks are processed (incremental), rather
+                # than one all-at-once sweep after buffering everything.
+                files_on_disk_at_each_span_fetch.append(
+                    len(list(project_dir.glob("trace_*.json")))
+                )
+                return _make_page([])
+
+            mock_client.rest_client.traces.get_traces_by_project.side_effect = [
+                page1,
+                page2,
+                page3,
+                _make_page([]),
+            ]
+            mock_client.rest_client.spans.get_spans_by_project.side_effect = (
+                span_fetch_side_effect
+            )
+
+            with patch(f"{_MODULE}.time.sleep"):
+                exported, skipped, had_errors = export_traces(
+                    client=mock_client,
+                    project_name="proj",
+                    project_dir=project_dir,
+                    max_results=None,
+                    filter_string=None,
+                    show_progress=False,
+                    page_size=2,
+                )
+
+            # Public result contract.
+            assert (exported, skipped, had_errors) == (6, 0, False)
+
+            # Final artifacts: exactly the six expected trace files were written.
+            written_ids = sorted(
+                p.name[len("trace_") : -len(".json")]
+                for p in project_dir.glob("trace_*.json")
+            )
+            assert written_ids == sorted(expected_ids)
+
+        # Bounded-chunk invariants (robust to the exact chunk boundary):
+        #  - more than one span-fetch session → work is chunked, not one sweep;
+        #  - counts never decrease → files accumulate as chunks flush;
+        #  - a later fetch sees files already written → earlier chunks were
+        #    flushed to disk before the last chunk was processed, so only ~one
+        #    chunk's worth of traces is held in memory at a time.
+        assert len(files_on_disk_at_each_span_fetch) > 1
+        assert files_on_disk_at_each_span_fetch == sorted(
+            files_on_disk_at_each_span_fetch
+        )
+        assert files_on_disk_at_each_span_fetch[0] == 0
+        assert files_on_disk_at_each_span_fetch[-1] > 0
+
+    def test_export_traces__non_uuid7_ids__span_scan_bounded_per_chunk(self):
+        """When trace ids aren't UUIDv7 the span scan can't be bounded by a
+        trace_id range, so the backend streams back every span in the project.
+
+        Rather than cache that whole-project span set — which would reintroduce
+        the unbounded memory peak this chunked export exists to prevent — each
+        chunk runs its own scan and keeps only its own traces' spans.  Peak
+        memory therefore stays bounded to one chunk regardless of project size;
+        the trade-off is one unfiltered scan per chunk on this rare fallback.
+        """
+        from opik.cli.exports.project import export_traces
+
+        # Non-UUIDv7 ids ("bulk-N") force the unfiltered-scan fallback.
+        page1 = _make_full_page(2, offset=0)
+        page2 = _make_full_page(2, offset=2)
+        page3 = _make_full_page(2, offset=4)
+        mock_client = MagicMock()
+        expected_ids = [f"bulk-{i}" for i in range(6)]
+
+        # Each per-chunk scan pages GET /spans starting at page 1; a fresh page-1
+        # request marks a new scan session.
+        scan_sessions = []
+
+        def span_fetch_side_effect(*args, **kwargs):
+            page = kwargs.get("page")
+            if page == 1:
+                scan_sessions.append([])
+            scan_sessions[-1].append(page)
+            return _make_page([])
+
+        mock_client.rest_client.traces.get_traces_by_project.side_effect = [
+            page1,
+            page2,
+            page3,
+            _make_page([]),
+        ]
+        mock_client.rest_client.spans.get_spans_by_project.side_effect = (
+            span_fetch_side_effect
+        )
+
+        with tempfile.TemporaryDirectory() as tmp:
+            project_dir = Path(tmp)
+            with patch(f"{_MODULE}.time.sleep"):
+                exported, skipped, had_errors = export_traces(
+                    client=mock_client,
+                    project_name="proj",
+                    project_dir=project_dir,
+                    max_results=None,
+                    filter_string=None,
+                    show_progress=False,
+                    page_size=2,
+                )
+
+            assert (exported, skipped, had_errors) == (6, 0, False)
+            written_ids = sorted(
+                p.name[len("trace_") : -len(".json")]
+                for p in project_dir.glob("trace_*.json")
+            )
+            assert written_ids == sorted(expected_ids)
+
+        # One unfiltered scan per flushed chunk (3) — not a single cached
+        # whole-project scan — which is what keeps peak memory bounded to a chunk.
+        assert len(scan_sessions) == 3
+
+    def test_export_traces__span_scan_error_isolated_to_its_chunk(self):
+        """A transient span-scan failure is confined to the chunk that hit it.
+
+        With per-chunk scanning there is no shared span cache to poison: the
+        failed chunk is exported with partial spans and had_errors is set, while
+        later chunks scan independently and export normally.  The failure is
+        injected at the public REST boundary (``spans.get_spans_by_project``),
+        per the SDK testing guidelines, not at a private transport helper.
+        """
+        from opik.cli.exports.project import export_traces
+
+        # Non-UUIDv7 ids force the unfiltered per-chunk scan fallback.
+        page1 = _make_full_page(2, offset=0)
+        page2 = _make_full_page(2, offset=2)
+        page3 = _make_full_page(2, offset=4)
+        mock_client = MagicMock()
+        expected_ids = [f"bulk-{i}" for i in range(6)]
+
+        with tempfile.TemporaryDirectory() as tmp:
+            project_dir = Path(tmp)
+
+            def get_spans_side_effect(*args, **kwargs):
+                # Fail every span page while the first chunk is being scanned
+                # (nothing flushed yet); once earlier chunks have written their
+                # files, succeed — proving the error didn't leak across chunks.
+                if not any(project_dir.glob("trace_*.json")):
+                    raise ApiError(status_code=503)
+                return _make_page([])
+
+            mock_client.rest_client.traces.get_traces_by_project.side_effect = [
+                page1,
+                page2,
+                page3,
+                _make_page([]),
+            ]
+            mock_client.rest_client.spans.get_spans_by_project.side_effect = (
+                get_spans_side_effect
+            )
+
+            # Neutralise the retry decorator's sleep (public boundary triggers
+            # the real 8-attempt retry) and the inter-page delay.
+            with patch(f"{_MODULE}._fetch_spans_page.retry.sleep"):
+                with patch(f"{_MODULE}.time.sleep"):
+                    exported, skipped, had_errors = export_traces(
+                        client=mock_client,
+                        project_name="proj",
+                        project_dir=project_dir,
+                        max_results=None,
+                        filter_string=None,
+                        show_progress=False,
+                        page_size=2,
+                    )
+
+            # The failed chunk still wrote its traces (with partial spans) and the
+            # failure is flagged; later chunks were unaffected, so all 6 landed.
+            assert (exported, skipped) == (6, 0)
+            assert had_errors is True
+            written_ids = sorted(
+                p.name[len("trace_") : -len(".json")]
+                for p in project_dir.glob("trace_*.json")
+            )
+            assert written_ids == sorted(expected_ids)
+
+
+# ---------------------------------------------------------------------------
+# export_traces — inter-page delay
+# ---------------------------------------------------------------------------
+
+
+class TestExportTracesInterPageDelay:
+    def test_export_traces__multiple_full_pages__sleep_not_called_on_success(self):
+        """sleep is NOT called between successful pages — the inter-page delay was
+        removed because server-side 429s are already handled by the retry decorator."""
+        from opik.cli.exports.project import export_traces, _PAGE_FETCH_DELAY_SECONDS
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                mock_fetch.side_effect = [
+                    _make_full_page(2, offset=0),  # page 1: full, IDs 0–1
+                    _make_full_page(2, offset=2),  # page 2: full, IDs 2–3
+                    _make_page([]),  # page 3: empty — end
+                ]
+                with patch(f"{_MODULE}._fetch_spans_page", return_value=_make_page([])):
+                    with patch(f"{_MODULE}.time.sleep") as mock_sleep:
+                        export_traces(
+                            client=mock_client,
+                            project_name="proj",
+                            project_dir=Path(tmp),
+                            max_results=None,
+                            filter_string=None,
+                            show_progress=False,
+                            page_size=2,
+                        )
+
+        from unittest.mock import call as mock_call
+
+        delay_calls = mock_sleep.call_args_list.count(
+            mock_call(_PAGE_FETCH_DELAY_SECONDS)
+        )
+        assert delay_calls == 0
+
+    def test_export_traces__single_partial_page__sleep_not_called(self):
+        """A partial (< 100 trace) page means we're at the end — no sleep needed."""
+        from opik.cli.exports.project import export_traces, _PAGE_FETCH_DELAY_SECONDS
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            with patch(f"{_MODULE}._fetch_traces_page") as mock_fetch:
+                # Single partial page — loop breaks immediately on short-page check
+                mock_fetch.side_effect = [
+                    _make_page([_make_mock_trace("t1")]),
+                ]
+                with patch(f"{_MODULE}.time.sleep") as mock_sleep:
+                    export_traces(
+                        client=mock_client,
+                        project_name="proj",
+                        project_dir=Path(tmp),
+                        max_results=None,
+                        filter_string=None,
+                        show_progress=False,
+                    )
+
+        from unittest.mock import call as mock_call
+
+        delay_calls = mock_sleep.call_args_list.count(
+            mock_call(_PAGE_FETCH_DELAY_SECONDS)
+        )
+        assert delay_calls == 0
+
+
+# ---------------------------------------------------------------------------
+# export_single_project — had_errors and project_exported propagation
+# ---------------------------------------------------------------------------
+
+
+class TestExportSingleProjectReturnValues:
+    def _make_project(self, name="test-proj"):
+        p = MagicMock()
+        p.name = name
+        p.id = "proj-id-1"
+        return p
+
+    def test_export_single_project__clean_run__had_errors_false(self):
+        from opik.cli.exports.project import export_single_project
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            projects_dir = Path(tmp) / "projects"
+            projects_dir.mkdir()
+
+            with (
+                patch(f"{_MODULE}.export_traces") as mock_export,
+                patch(f"{_MODULE}.time.sleep"),
+            ):
+                mock_export.return_value = (3, 0, False)
+                result = export_single_project(
+                    client=mock_client,
+                    project_name="test-proj",
+                    project_dir=projects_dir,
+                    filter_string=None,
+                    max_results=None,
+                    force=False,
+                    debug=False,
+                    format="json",
+                )
+
+        assert len(result) == 4
+        proj_count, _exported, _skipped, had_errors = result
+        assert had_errors is False
+        assert proj_count == 1  # traces were exported
+
+    def test_export_single_project__run_with_errors__had_errors_true(self):
+        from opik.cli.exports.project import export_single_project
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            projects_dir = Path(tmp) / "projects"
+            projects_dir.mkdir()
+
+            with (
+                patch(f"{_MODULE}.export_traces") as mock_export,
+                patch(f"{_MODULE}.time.sleep"),
+            ):
+                mock_export.return_value = (2, 1, True)
+                result = export_single_project(
+                    client=mock_client,
+                    project_name="test-proj",
+                    project_dir=projects_dir,
+                    filter_string=None,
+                    max_results=None,
+                    force=False,
+                    debug=False,
+                    format="json",
+                )
+
+        assert len(result) == 4
+        _proj_count, _exported, _skipped, had_errors = result
+        assert had_errors is True
+
+    def test_export_single_project__no_traces_found__project_count_is_zero(self):
+        """When there are no traces at all, the project-exported flag must be 0."""
+        from opik.cli.exports.project import export_single_project
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            projects_dir = Path(tmp) / "projects"
+            projects_dir.mkdir()
+
+            with (
+                patch(f"{_MODULE}.export_traces") as mock_export,
+                patch(f"{_MODULE}.time.sleep"),
+            ):
+                mock_export.return_value = (0, 0, False)  # nothing exported or skipped
+                result = export_single_project(
+                    client=mock_client,
+                    project_name="test-proj",
+                    project_dir=projects_dir,
+                    filter_string=None,
+                    max_results=None,
+                    force=False,
+                    debug=False,
+                    format="json",
+                )
+
+        proj_count, _exported, _skipped, _had_errors = result
+        assert proj_count == 0
+
+    def test_export_single_project__exception_during_export__returns_had_errors_true(
+        self,
+    ):
+        from opik.cli.exports.project import export_single_project
+
+        mock_client = MagicMock()
+
+        with tempfile.TemporaryDirectory() as tmp:
+            projects_dir = Path(tmp) / "projects"
+            projects_dir.mkdir()
+
+            with patch(f"{_MODULE}.export_traces") as mock_export:
+                mock_export.side_effect = RuntimeError("unexpected failure")
+                result = export_single_project(
+                    client=mock_client,
+                    project_name="test-proj",
+                    project_dir=projects_dir,
+                    filter_string=None,
+                    max_results=None,
+                    force=False,
+                    debug=False,
+                    format="json",
+                )
+
+        assert result == (0, 0, 0, True)

sdks/python/tests/unit/cli/test_import_experiment.py

@@ -0,0 +1,900 @@
+"""Unit tests for experiment import functionality."""
+
+import json
+import sys
+import types
+from pathlib import Path
+from typing import Dict, Any
+from unittest.mock import Mock, MagicMock, patch
+import pytest
+
+# Mock the problematic imports before importing
+# Mock the prompt import that's causing issues
+sys.modules["opik.api_objects.prompt.prompt"] = MagicMock()
+
+# Now we can import normally
+from opik.cli.imports.experiment import (  # noqa: E402
+    ExperimentData,
+    load_experiment_data,
+    recreate_experiment,
+    _import_traces_for_project,
+)
+from opik.cli.imports.utils import (  # noqa: E402
+    translate_trace_id as utils_translate_trace_id,
+    sort_spans_topologically,
+)
+
+
+class TestExperimentData:
+    """Test ExperimentData dataclass."""
+
+    def test_experiment_data_from_dict(self) -> None:
+        """Test creating ExperimentData from dictionary."""
+        data = {
+            "experiment": {
+                "id": "exp-123",
+                "name": "test-experiment",
+                "dataset_name": "test-dataset",
+            },
+            "items": [
+                {"id": "item-1", "trace_id": "trace-1"},
+                {"id": "item-2", "trace_id": "trace-2"},
+            ],
+            "downloaded_at": "2024-01-01T00:00:00",
+        }
+
+        exp_data = ExperimentData.from_dict(data)
+
+        assert exp_data.experiment["id"] == "exp-123"
+        assert exp_data.experiment["name"] == "test-experiment"
+        assert len(exp_data.items) == 2
+        assert exp_data.downloaded_at == "2024-01-01T00:00:00"
+
+    def test_experiment_data_from_dict_minimal(self) -> None:
+        """Test creating ExperimentData with minimal data."""
+        data = {
+            "experiment": {"id": "exp-123", "dataset_name": "test-dataset"},
+            "items": [],
+        }
+
+        exp_data = ExperimentData.from_dict(data)
+
+        assert exp_data.experiment["id"] == "exp-123"
+        assert exp_data.items == []
+        assert exp_data.downloaded_at is None
+
+    def test_load_experiment_data_from_file(self, tmp_path: Path) -> None:
+        """Test loading experiment data from JSON file."""
+        experiment_file = tmp_path / "experiment_test.json"
+        data = {
+            "experiment": {
+                "id": "exp-123",
+                "name": "test-experiment",
+                "dataset_name": "test-dataset",
+            },
+            "items": [{"id": "item-1"}],
+        }
+
+        with open(experiment_file, "w") as f:
+            json.dump(data, f)
+
+        exp_data = load_experiment_data(experiment_file)
+
+        assert isinstance(exp_data, ExperimentData)
+        assert exp_data.experiment["id"] == "exp-123"
+        assert len(exp_data.items) == 1
+
+
+class TestTranslateTraceId:
+    """Test translate_trace_id function."""
+
+    def test_translate_trace_id_found(self) -> None:
+        """Test translating trace ID when mapping exists."""
+        trace_id_map = {"old-trace-1": "new-trace-1", "old-trace-2": "new-trace-2"}
+
+        result = utils_translate_trace_id("old-trace-1", trace_id_map)
+
+        assert result == "new-trace-1"
+
+    def test_translate_trace_id_not_found(self) -> None:
+        """Test translating trace ID when mapping doesn't exist."""
+        trace_id_map = {"old-trace-1": "new-trace-1"}
+
+        result = utils_translate_trace_id("old-trace-2", trace_id_map)
+
+        assert result is None
+
+    def test_translate_trace_id_empty_map(self) -> None:
+        """Test translating trace ID with empty map."""
+        trace_id_map: Dict[str, str] = {}
+
+        result = utils_translate_trace_id("old-trace-1", trace_id_map)
+
+        assert result is None
+
+    def test_translate_trace_id_requires_dict(self) -> None:
+        """Test that translate_trace_id requires Dict, not Optional."""
+        # This test verifies the type signature is correct
+        # If someone tries to pass None, type checker should catch it
+        trace_id_map: Dict[str, str] = {}  # Required, not Optional
+
+        result = utils_translate_trace_id("trace-1", trace_id_map)
+
+        assert result is None
+
+
+class TestRecreateExperiment:
+    """Test recreate_experiment function."""
+
+    @staticmethod
+    def _extract_items_arg_from_call_args(call_args: Any) -> Any:
+        """Helper to extract the items argument from call_args.
+
+        Handles both positional and keyword arguments.
+        """
+        if hasattr(call_args, "args") and call_args.args:
+            return call_args.args[0]
+        if hasattr(call_args, "kwargs"):
+            if "items" in call_args.kwargs:
+                return call_args.kwargs["items"]
+            for value in call_args.kwargs.values():
+                if isinstance(value, list) and len(value) > 0:
+                    return value
+        return None
+
+    @pytest.fixture
+    def mock_client(self) -> Mock:
+        """Create a mock Opik client."""
+        client = Mock()
+        # Ensure flush returns True to indicate success
+        client.flush = Mock(return_value=True)
+
+        # Mock dataset
+        mock_dataset = Mock()
+        mock_dataset.name = "test-dataset"
+        mock_dataset.__internal_api__insert_items_as_dataclasses__ = Mock()
+
+        # Mock experiment
+        mock_experiment = Mock()
+        mock_experiment.insert = Mock()
+        mock_experiment.id = "exp-123"
+
+        client.get_or_create_dataset = Mock(return_value=mock_dataset)
+        client.create_experiment = Mock(return_value=mock_experiment)
+        # Mock REST client for experiment items creation
+        client._rest_client = Mock()
+        client._rest_client.experiments = Mock()
+        client._rest_client.experiments.create_experiment_items = Mock()
+
+        return client
+
+    @pytest.fixture
+    def experiment_data(self) -> ExperimentData:
+        """Create sample experiment data."""
+        return ExperimentData(
+            experiment={
+                "id": "exp-123",
+                "name": "test-experiment",
+                "dataset_name": "test-dataset",
+                "type": "regular",
+            },
+            items=[
+                {
+                    "trace_id": "trace-1",
+                    "dataset_item_id": "ds-item-1",
+                    "dataset_item_data": {
+                        "input": "test input",
+                        "expected_output": "test output",
+                    },
+                },
+                {
+                    "trace_id": "trace-2",
+                    "dataset_item_id": "ds-item-2",
+                    "dataset_item_data": {"input": "test input 2"},
+                },
+            ],
+        )
+
+    def test_recreate_experiment_requires_trace_id_map(
+        self, mock_client: Mock, experiment_data: ExperimentData
+    ) -> None:
+        """Test that recreate_experiment requires trace_id_map (not Optional)."""
+        # This test verifies the type signature
+        trace_id_map: Dict[str, str] = {
+            "trace-1": "new-trace-1",
+            "trace-2": "new-trace-2",
+        }
+        dataset_item_id_map: Dict[str, str] = {
+            "ds-item-1": "new-ds-item-1",
+            "ds-item-2": "new-ds-item-2",
+        }
+
+        # Should not accept None - type checker would catch this
+        # We test that it works with a dict
+        with patch("opik.cli.imports.experiment.id_helpers_module") as mock_id_helpers:
+            mock_id_helpers.generate_id = Mock(return_value="generated-id")
+
+            recreate_experiment(
+                mock_client,
+                experiment_data,
+                "test-project",
+                trace_id_map,  # Required, not Optional
+                dataset_item_id_map,  # Required for mapping dataset items
+                dry_run=False,
+                debug=False,
+            )
+
+            # Verify experiment items were created via REST API
+            assert mock_client._rest_client.experiments.create_experiment_items.called
+
+    def test_recreate_experiment_batches_dataset_items(
+        self, mock_client: Mock, experiment_data: ExperimentData
+    ) -> None:
+        """Test that experiment items are created in batch, not one at a time."""
+        with patch("opik.cli.imports.experiment.id_helpers_module") as mock_id_helpers:
+            mock_id_helpers.generate_id = Mock(side_effect=["exp-item-1", "exp-item-2"])
+
+            trace_id_map = {"trace-1": "new-trace-1", "trace-2": "new-trace-2"}
+            dataset_item_id_map = {
+                "ds-item-1": "new-ds-item-1",
+                "ds-item-2": "new-ds-item-2",
+            }
+
+            recreate_experiment(
+                mock_client,
+                experiment_data,
+                "test-project",
+                trace_id_map,
+                dataset_item_id_map,
+                dry_run=False,
+                debug=False,
+            )
+
+            # Verify batch insert was called ONCE with all items
+            assert (
+                mock_client._rest_client.experiments.create_experiment_items.call_count
+                == 1
+            )
+
+            # Verify it was called with a list of items (batch)
+            call_args = (
+                mock_client._rest_client.experiments.create_experiment_items.call_args
+            )
+            assert call_args is not None
+
+            # Extract experiment_items argument from call_args
+            experiment_items_arg = None
+            if hasattr(call_args, "kwargs") and "experiment_items" in call_args.kwargs:
+                experiment_items_arg = call_args.kwargs["experiment_items"]
+            elif hasattr(call_args, "args") and call_args.args:
+                experiment_items_arg = call_args.args[0]
+
+            assert experiment_items_arg is not None, (
+                f"Could not find experiment_items in call_args: {call_args}"
+            )
+            assert len(experiment_items_arg) == 2, (
+                f"Expected 2 items in batch, got {len(experiment_items_arg)}"
+            )
+
+    def test_recreate_experiment_uses_module_names_correctly(
+        self, mock_client: Mock, experiment_data: ExperimentData
+    ) -> None:
+        """Test that module names (id_helpers_module) are used correctly."""
+        with patch("opik.cli.imports.experiment.id_helpers_module") as mock_id_helpers:
+            mock_id_helpers.generate_id = Mock(return_value="generated-id")
+
+            trace_id_map = {"trace-1": "new-trace-1"}
+            dataset_item_id_map = {"ds-item-1": "new-ds-item-1"}
+
+            recreate_experiment(
+                mock_client,
+                experiment_data,
+                "test-project",
+                trace_id_map,
+                dataset_item_id_map,
+                dry_run=False,
+                debug=False,
+            )
+
+            # Verify id_helpers module is used (not checked for None)
+            assert mock_id_helpers.generate_id.called
+
+    def test_recreate_experiment_handles_empty_trace_id_map(
+        self, mock_client: Mock, experiment_data: ExperimentData
+    ) -> None:
+        """Test that empty trace_id_map is handled correctly."""
+        trace_id_map: Dict[str, str] = {}  # Empty but valid
+        dataset_item_id_map: Dict[str, str] = {}  # Empty but valid
+
+        recreate_experiment(
+            mock_client,
+            experiment_data,
+            "test-project",
+            trace_id_map,
+            dataset_item_id_map,
+            dry_run=False,
+            debug=False,
+        )
+
+        # Should still create experiment and dataset, but skip items
+        assert mock_client.get_or_create_dataset.called
+        assert mock_client.create_experiment.called
+
+    def test_recreate_experiment_dry_run(
+        self, mock_client: Mock, experiment_data: ExperimentData
+    ) -> None:
+        """Test dry run mode."""
+        trace_id_map = {"trace-1": "new-trace-1"}
+        dataset_item_id_map = {"ds-item-1": "new-ds-item-1"}
+
+        result = recreate_experiment(
+            mock_client,
+            experiment_data,
+            "test-project",
+            trace_id_map,
+            dataset_item_id_map,
+            dry_run=True,
+            debug=False,
+        )
+
+        assert result is True
+        # Should not create anything in dry run
+        assert not mock_client.get_or_create_dataset.called
+        assert not mock_client.create_experiment.called
+
+    def test_recreate_experiment_chunks_items_within_be_cap(
+        self, mock_client: Mock
+    ) -> None:
+        # BE rejects a single ``create_experiment_items`` POST whose item
+        # count exceeds ``ExperimentItemsBatch``'s ``@Size(max=…)``. The
+        # actual cap value is BE-configured and may change over time, so
+        # we read it from the module-level constant the chunker uses and
+        # assert behavior against that, not against a hardcoded literal.
+        # Contract: every batch ``<= cap``, total across batches ==
+        # full item count, multiple batches when items > cap.
+        from opik.cli.imports.experiment import _EXPERIMENT_ITEMS_INSERT_BATCH_SIZE
+
+        cap = _EXPERIMENT_ITEMS_INSERT_BATCH_SIZE
+        # 2.5 × cap exercises the multi-batch path with a partial-last-
+        # batch remainder; any value above ``cap`` works for the contract.
+        n_items = cap * 5 // 2
+        experiment_data = ExperimentData(
+            experiment={
+                "id": "exp-big",
+                "name": "big-experiment",
+                "dataset_name": "test-dataset",
+                "type": "regular",
+            },
+            items=[
+                {
+                    "trace_id": f"trace-{i}",
+                    "dataset_item_id": f"ds-item-{i}",
+                    "dataset_item_data": {"input": f"test input {i}"},
+                }
+                for i in range(n_items)
+            ],
+        )
+        trace_id_map = {f"trace-{i}": f"new-trace-{i}" for i in range(n_items)}
+        dataset_item_id_map = {
+            f"ds-item-{i}": f"new-ds-item-{i}" for i in range(n_items)
+        }
+
+        with patch("opik.cli.imports.experiment.id_helpers_module") as mock_id_helpers:
+            mock_id_helpers.generate_id = Mock(
+                side_effect=[f"exp-item-{i}" for i in range(n_items)]
+            )
+
+            recreate_experiment(
+                mock_client,
+                experiment_data,
+                "test-project",
+                trace_id_map,
+                dataset_item_id_map,
+                dry_run=False,
+                debug=False,
+            )
+
+            create_calls = mock_client._rest_client.experiments.create_experiment_items.call_args_list
+            # ceil(n_items / cap) batched POSTs.
+            expected_batches = (n_items + cap - 1) // cap
+            assert len(create_calls) == expected_batches, (
+                f"expected {expected_batches} chunked create_experiment_items "
+                f"calls for {n_items} items at cap={cap}, got {len(create_calls)}"
+            )
+            # Every batch must respect the BE cap; total must equal the
+            # full item count (nothing dropped, nothing duplicated).
+            batch_sizes = [
+                len(call.kwargs.get("experiment_items", []))
+                if call.kwargs.get("experiment_items") is not None
+                else len(call.args[0])
+                for call in create_calls
+            ]
+            assert all(size <= cap for size in batch_sizes), (
+                f"every batch must respect the BE cap (={cap}), got {batch_sizes}"
+            )
+            assert sum(batch_sizes) == n_items, (
+                f"all {n_items} items must end up across the chunked batches, "
+                f"got sum={sum(batch_sizes)}"
+            )
+
+
+class TestImportTracesWithSpans:
+    """Test trace import with span parent_span_id preservation."""
+
+    @pytest.fixture
+    def mock_client(self) -> Mock:
+        """Create a mock Opik client."""
+        client = Mock()
+        client.flush = Mock()
+
+        # Mock trace creation
+        mock_trace = Mock()
+        mock_trace.id = "new-trace-1"
+        client.trace = Mock(return_value=mock_trace)
+
+        # Mock span creation
+        mock_spans = []
+        for i in range(3):
+            mock_span = Mock()
+            mock_span.id = f"new-span-{i + 1}"
+            mock_spans.append(mock_span)
+
+        client.span = Mock(side_effect=mock_spans)
+
+        return client
+
+    def test_import_traces_preserves_span_hierarchy(
+        self, mock_client: Mock, tmp_path: Path
+    ) -> None:
+        """Test that span parent_span_id relationships are preserved."""
+        # Create test trace file with spans
+        projects_dir = tmp_path / "projects" / "test-project"
+        projects_dir.mkdir(parents=True)
+
+        trace_data = {
+            "trace": {
+                "id": "original-trace-1",
+                "name": "test-trace",
+                "input": {},
+                "output": {},
+            },
+            "spans": [
+                {
+                    "id": "span-1",
+                    "name": "root-span",
+                    "parent_span_id": None,  # Root span
+                    "input": {},
+                    "output": {},
+                },
+                {
+                    "id": "span-2",
+                    "name": "child-span",
+                    "parent_span_id": "span-1",  # Child of span-1
+                    "input": {},
+                    "output": {},
+                },
+                {
+                    "id": "span-3",
+                    "name": "grandchild-span",
+                    "parent_span_id": "span-2",  # Child of span-2
+                    "input": {},
+                    "output": {},
+                },
+            ],
+        }
+
+        trace_file = projects_dir / "trace_original-trace-1.json"
+        with open(trace_file, "w") as f:
+            json.dump(trace_data, f)
+
+        # Import traces
+        trace_id_map, _ = _import_traces_for_project(
+            mock_client, projects_dir, "test-project", dry_run=False, debug=False
+        )
+
+        # Verify spans were created
+        assert mock_client.span.call_count == 3
+
+        # Verify spans were created in correct order (root first, then children)
+        span_calls = mock_client.span.call_args_list
+
+        # First span should be root (no parent_span_id)
+        first_call = span_calls[0]
+        assert first_call.kwargs.get("parent_span_id") is None
+
+        # Second span should have parent_span_id set to first span's new ID
+        # Note: We can't easily verify the exact ID mapping without more complex mocking,
+        # but we can verify that parent_span_id is being passed
+        second_call = span_calls[1]
+        # The parent_span_id should be set (not None) since span-1 was created first
+        # and its new ID should be in span_id_map
+        assert "parent_span_id" in second_call.kwargs
+
+        # Third span (grandchild) should have parent_span_id set to second span's new ID
+        third_call = span_calls[2]
+        assert "parent_span_id" in third_call.kwargs
+        assert third_call.kwargs.get("parent_span_id") is not None
+
+        # Verify trace was created
+        assert mock_client.trace.called
+        assert "original-trace-1" in trace_id_map
+
+    def test_import_traces_preserves_deep_hierarchy(
+        self, mock_client: Mock, tmp_path: Path
+    ) -> None:
+        """Test that deep hierarchies (4+ levels) are preserved correctly."""
+        projects_dir = tmp_path / "projects" / "test-project"
+        projects_dir.mkdir(parents=True)
+
+        # Create a 4-level hierarchy: root -> child -> grandchild -> great-grandchild
+        # Spans are intentionally in wrong order to test sorting
+        trace_data = {
+            "trace": {
+                "id": "original-trace-1",
+                "name": "test-trace",
+                "input": {},
+                "output": {},
+            },
+            "spans": [
+                {
+                    "id": "span-4",
+                    "name": "great-grandchild",
+                    "parent_span_id": "span-3",
+                    "input": {},
+                    "output": {},
+                },
+                {
+                    "id": "span-2",
+                    "name": "child",
+                    "parent_span_id": "span-1",
+                    "input": {},
+                    "output": {},
+                },
+                {
+                    "id": "span-1",
+                    "name": "root",
+                    "parent_span_id": None,
+                    "input": {},
+                    "output": {},
+                },
+                {
+                    "id": "span-3",
+                    "name": "grandchild",
+                    "parent_span_id": "span-2",
+                    "input": {},
+                    "output": {},
+                },
+            ],
+        }
+
+        trace_file = projects_dir / "trace_original-trace-1.json"
+        with open(trace_file, "w") as f:
+            json.dump(trace_data, f)
+
+        # Import traces
+        trace_id_map, _ = _import_traces_for_project(
+            mock_client, projects_dir, "test-project", dry_run=False, debug=False
+        )
+
+        # Verify all spans were created
+        assert mock_client.span.call_count == 4
+
+        span_calls = mock_client.span.call_args_list
+
+        # Verify order: root -> child -> grandchild -> great-grandchild
+        # First span should be root (no parent)
+        assert span_calls[0].kwargs.get("parent_span_id") is None
+
+        # Second span should be child (has parent)
+        assert span_calls[1].kwargs.get("parent_span_id") is not None
+
+        # Third span should be grandchild (has parent)
+        assert span_calls[2].kwargs.get("parent_span_id") is not None
+
+        # Fourth span should be great-grandchild (has parent)
+        assert span_calls[3].kwargs.get("parent_span_id") is not None
+
+        # Verify trace was created
+        assert mock_client.trace.called
+        assert "original-trace-1" in trace_id_map
+
+    def test_import_traces_sorts_spans_correctly(
+        self, mock_client: Mock, tmp_path: Path
+    ) -> None:
+        """Test that spans are sorted (root spans first, then children)."""
+        projects_dir = tmp_path / "projects" / "test-project"
+        projects_dir.mkdir(parents=True)
+
+        trace_data = {
+            "trace": {
+                "id": "original-trace-1",
+                "name": "test-trace",
+                "input": {},
+                "output": {},
+            },
+            "spans": [
+                {
+                    "id": "span-2",
+                    "name": "child-span",
+                    "parent_span_id": "span-1",  # Child - should come after root
+                    "input": {},
+                    "output": {},
+                },
+                {
+                    "id": "span-1",
+                    "name": "root-span",
+                    "parent_span_id": None,  # Root - should come first
+                    "input": {},
+                    "output": {},
+                },
+            ],
+        }
+
+        trace_file = projects_dir / "trace_original-trace-1.json"
+        with open(trace_file, "w") as f:
+            json.dump(trace_data, f)
+
+        # Import traces
+        _, _ = _import_traces_for_project(
+            mock_client, projects_dir, "test-project", dry_run=False, debug=False
+        )  # Returns (trace_id_map, stats), but we don't need them for this test
+
+        # Verify spans were created in correct order
+        span_calls = mock_client.span.call_args_list
+
+        # First span should be root (no parent)
+        assert span_calls[0].kwargs.get("parent_span_id") is None
+
+        # Second span should have parent_span_id
+        assert span_calls[1].kwargs.get("parent_span_id") is not None
+
+
+class TestTopologicalSort:
+    """Test the topological sort function for spans."""
+
+    def test_sort_spans_simple_hierarchy(self) -> None:
+        """Test sorting with a simple 2-level hierarchy."""
+        spans = [
+            {"id": "span-2", "name": "child", "parent_span_id": "span-1"},
+            {"id": "span-1", "name": "root", "parent_span_id": None},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Root should come first
+        assert sorted_spans[0]["id"] == "span-1"
+        assert sorted_spans[0]["parent_span_id"] is None
+        # Child should come second
+        assert sorted_spans[1]["id"] == "span-2"
+        assert sorted_spans[1]["parent_span_id"] == "span-1"
+
+    def test_sort_spans_multi_level_hierarchy(self) -> None:
+        """Test sorting with a 4-level hierarchy (root -> child -> grandchild -> great-grandchild)."""
+        spans = [
+            {"id": "span-4", "name": "great-grandchild", "parent_span_id": "span-3"},
+            {"id": "span-2", "name": "child", "parent_span_id": "span-1"},
+            {"id": "span-1", "name": "root", "parent_span_id": None},
+            {"id": "span-3", "name": "grandchild", "parent_span_id": "span-2"},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Verify order: root -> child -> grandchild -> great-grandchild
+        assert sorted_spans[0]["id"] == "span-1"
+        assert sorted_spans[0]["parent_span_id"] is None
+
+        assert sorted_spans[1]["id"] == "span-2"
+        assert sorted_spans[1]["parent_span_id"] == "span-1"
+
+        assert sorted_spans[2]["id"] == "span-3"
+        assert sorted_spans[2]["parent_span_id"] == "span-2"
+
+        assert sorted_spans[3]["id"] == "span-4"
+        assert sorted_spans[3]["parent_span_id"] == "span-3"
+
+    def test_sort_spans_multiple_roots(self) -> None:
+        """Test sorting with multiple root spans."""
+        spans = [
+            {"id": "span-3", "name": "child-of-2", "parent_span_id": "span-2"},
+            {"id": "span-1", "name": "root-1", "parent_span_id": None},
+            {"id": "span-2", "name": "root-2", "parent_span_id": None},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Both roots should come before the child
+        root_ids = {sorted_spans[0]["id"], sorted_spans[1]["id"]}
+        assert root_ids == {"span-1", "span-2"}
+        assert sorted_spans[0]["parent_span_id"] is None
+        assert sorted_spans[1]["parent_span_id"] is None
+
+        # Child should come last
+        assert sorted_spans[2]["id"] == "span-3"
+        assert sorted_spans[2]["parent_span_id"] == "span-2"
+
+    def test_sort_spans_multiple_children(self) -> None:
+        """Test sorting with a root that has multiple children."""
+        spans = [
+            {"id": "span-3", "name": "child-2", "parent_span_id": "span-1"},
+            {"id": "span-1", "name": "root", "parent_span_id": None},
+            {"id": "span-2", "name": "child-1", "parent_span_id": "span-1"},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Root should come first
+        assert sorted_spans[0]["id"] == "span-1"
+        assert sorted_spans[0]["parent_span_id"] is None
+
+        # Both children should come after root (order doesn't matter for siblings)
+        child_ids = {sorted_spans[1]["id"], sorted_spans[2]["id"]}
+        assert child_ids == {"span-2", "span-3"}
+        assert sorted_spans[1]["parent_span_id"] == "span-1"
+        assert sorted_spans[2]["parent_span_id"] == "span-1"
+
+    def test_sort_spans_missing_parent(self) -> None:
+        """Test sorting when a span references a non-existent parent."""
+        spans = [
+            {"id": "span-1", "name": "root", "parent_span_id": None},
+            {"id": "span-2", "name": "orphan", "parent_span_id": "nonexistent"},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Both should be treated as roots (orphan becomes root)
+        assert len(sorted_spans) == 2
+        # Both should have no parent or invalid parent
+        for span in sorted_spans:
+            assert (
+                span["parent_span_id"] is None
+                or span["parent_span_id"] == "nonexistent"
+            )
+
+    def test_sort_spans_empty_list(self) -> None:
+        """Test sorting with empty list."""
+        spans: list = []
+        sorted_spans = sort_spans_topologically(spans)
+        assert sorted_spans == []
+
+    def test_sort_spans_single_root(self) -> None:
+        """Test sorting with single root span."""
+        spans = [{"id": "span-1", "name": "root", "parent_span_id": None}]
+        sorted_spans = sort_spans_topologically(spans)
+        assert len(sorted_spans) == 1
+        assert sorted_spans[0]["id"] == "span-1"
+        assert sorted_spans[0]["parent_span_id"] is None
+
+    def test_sort_spans_all_have_parents(self) -> None:
+        """Test sorting when all spans have parents (no explicit root).
+
+        This tests the fix for the bug where empty root_spans would cause
+        the function to return an empty list, silently dropping all spans.
+        """
+        spans = [
+            {"id": "span-1", "name": "child-1", "parent_span_id": "span-2"},
+            {"id": "span-2", "name": "child-2", "parent_span_id": "span-1"},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Should not return empty list - all spans should be included
+        assert len(sorted_spans) == 2
+        # Verify all spans are present
+        span_ids = {span["id"] for span in sorted_spans}
+        assert span_ids == {"span-1", "span-2"}
+
+    def test_sort_spans_cycle(self) -> None:
+        """Test sorting with a cycle in the span graph.
+
+        This tests that cycles don't cause infinite loops and all spans
+        are still included in the result.
+        """
+        spans = [
+            {"id": "span-1", "name": "span-1", "parent_span_id": "span-2"},
+            {"id": "span-2", "name": "span-2", "parent_span_id": "span-3"},
+            {"id": "span-3", "name": "span-3", "parent_span_id": "span-1"},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # Should not return empty list - all spans should be included
+        assert len(sorted_spans) == 3
+        # Verify all spans are present
+        span_ids = {span["id"] for span in sorted_spans}
+        assert span_ids == {"span-1", "span-2", "span-3"}
+
+    def test_sort_spans_disconnected_components(self) -> None:
+        """Test sorting with disconnected components (multiple separate graphs).
+
+        This tests that spans not reachable from root spans are still included.
+        """
+        spans = [
+            {"id": "span-1", "name": "root-1", "parent_span_id": None},
+            {"id": "span-2", "name": "child-1", "parent_span_id": "span-1"},
+            {"id": "span-3", "name": "disconnected-1", "parent_span_id": "span-4"},
+            {"id": "span-4", "name": "disconnected-2", "parent_span_id": "span-3"},
+        ]
+
+        sorted_spans = sort_spans_topologically(spans)
+
+        # All spans should be included
+        assert len(sorted_spans) == 4
+        # Verify all spans are present
+        span_ids = {span["id"] for span in sorted_spans}
+        assert span_ids == {"span-1", "span-2", "span-3", "span-4"}
+        # Root span should come first
+        assert sorted_spans[0]["id"] == "span-1"
+        assert sorted_spans[0]["parent_span_id"] is None
+
+
+class TestModuleNameUsage:
+    """Test that module names are used correctly (not checked for None)."""
+
+    def test_module_names_are_modules_not_variables(self) -> None:
+        """Test that dataset_item_module and id_helpers_module are modules."""
+        from opik.cli.imports.experiment import dataset_item_module, id_helpers_module
+
+        # Modules should exist and be importable
+        assert dataset_item_module is not None
+        assert id_helpers_module is not None
+
+        # They should be modules, not None
+        assert isinstance(dataset_item_module, types.ModuleType)
+        assert isinstance(id_helpers_module, types.ModuleType)
+
+
+class TestProjectTraceImport:
+    """Test importing a project's traces from its project directory."""
+
+    @pytest.fixture
+    def mock_client(self) -> Mock:
+        """Create a minimal mock Opik client."""
+        client = Mock()
+        client.flush = Mock()
+        mock_trace = Mock()
+        mock_trace.id = "new-trace-id"
+        client.trace = Mock(return_value=mock_trace)
+        client.span = Mock()
+        return client
+
+    def test_trace_files_in_project_dir_are_imported(
+        self, mock_client: Mock, tmp_path: Path
+    ) -> None:
+        """trace_{id}.json files directly under the project dir are imported and
+        appear in the returned trace_id_map, and the trace is created in the
+        named project."""
+        project_dir = tmp_path / "projects" / "my-project"
+        project_dir.mkdir(parents=True)
+
+        trace_id = "abc123"
+        trace_data: Dict[str, Any] = {
+            "trace": {"id": trace_id, "name": "t", "input": {}, "output": {}},
+            "spans": [],
+        }
+        trace_file = project_dir / f"trace_{trace_id}.json"
+        with open(trace_file, "w") as f:
+            json.dump(trace_data, f)
+
+        trace_id_map, _ = _import_traces_for_project(
+            mock_client, project_dir, "my-project", dry_run=False, debug=False
+        )
+
+        # The original trace ID must appear in the returned map
+        assert trace_id in trace_id_map
+        # The trace must be created in the named project (no "default" fallback)
+        assert mock_client.trace.call_args.kwargs.get("project_name") == "my-project"
+
+    def test_empty_project_dir_returns_empty_map(
+        self, mock_client: Mock, tmp_path: Path
+    ) -> None:
+        """A project dir with no trace files yields an empty map without error."""
+        project_dir = tmp_path / "projects" / "empty-project"
+        project_dir.mkdir(parents=True)
+
+        trace_id_map, stats = _import_traces_for_project(
+            mock_client, project_dir, "empty-project", dry_run=False, debug=False
+        )
+
+        assert trace_id_map == {}
+        assert stats["traces"] == 0
+        assert not mock_client.trace.called

sdks/python/tests/unit/cli/test_mcp.py

@@ -0,0 +1,208 @@
+"""Tests for the ``opik mcp configure`` command."""
+
+import pathlib
+from unittest.mock import patch
+
+from click.testing import CliRunner
+
+from opik.cli import cli
+from opik.cli import mcp as mcp_cli
+from opik.config import OpikConfig
+
+
+def _config(**overrides) -> OpikConfig:
+    values = dict(url_override="https://www.comet.com/opik/api/", workspace="acme-ai")
+    values.update(overrides)
+    return OpikConfig(**values)
+
+
+class TestResolveSetupParams:
+    def test_cloud__no_url_flags(self):
+        params = mcp_cli._resolve_setup_params(
+            _config(api_key="key", url_override="https://www.comet.com/opik/api/")
+        )
+        assert params["use_local"] is False
+        assert params["self_hosted_comet"] is False
+        assert params["api_url"] == "https://www.comet.com/opik/api/"
+
+    def test_self_hosted_comet__detected_from_opik_api_path(self):
+        params = mcp_cli._resolve_setup_params(
+            _config(api_key="key", url_override="https://opik.acme.com/opik/api/")
+        )
+        assert params["self_hosted_comet"] is True
+        assert params["use_local"] is False
+        assert params["base_url"] == "https://opik.acme.com/"
+
+    def test_localhost__is_use_local(self):
+        params = mcp_cli._resolve_setup_params(
+            _config(api_key=None, url_override="http://localhost:5173/api/")
+        )
+        assert params["use_local"] is True
+        assert params["self_hosted_comet"] is False
+
+    def test_self_hosted_oss__non_opik_path_is_use_local(self):
+        params = mcp_cli._resolve_setup_params(
+            _config(api_key=None, url_override="https://opik.acme.com/api/")
+        )
+        assert params["use_local"] is True
+        assert params["self_hosted_comet"] is False
+
+
+class TestInstallCommand:
+    def test_install__reads_config_and_calls_setup(self):
+        runner = CliRunner()
+        with (
+            patch.object(
+                mcp_cli.opik_config, "OpikConfig", return_value=_config(api_key="key")
+            ),
+            patch.object(
+                mcp_cli.interactive_helpers, "is_interactive", return_value=True
+            ),
+            patch.object(mcp_cli.mcp_installer, "setup_mcp_server") as setup_spy,
+        ):
+            result = runner.invoke(cli, ["mcp", "configure"])
+
+        assert result.exit_code == 0
+        setup_spy.assert_called_once()
+        assert setup_spy.call_args.kwargs["api_key"] == "key"
+        assert setup_spy.call_args.kwargs["workspace"] == "acme-ai"
+        assert setup_spy.call_args.kwargs["force_local_server"] is False
+
+    def test_install__local_server_flag__forces_local(self):
+        runner = CliRunner()
+        with (
+            patch.object(
+                mcp_cli.opik_config, "OpikConfig", return_value=_config(api_key="key")
+            ),
+            patch.object(
+                mcp_cli.interactive_helpers, "is_interactive", return_value=True
+            ),
+            patch.object(mcp_cli.mcp_installer, "setup_mcp_server") as setup_spy,
+        ):
+            result = runner.invoke(cli, ["mcp", "configure", "--local-server"])
+
+        assert result.exit_code == 0
+        setup_spy.assert_called_once()
+        assert setup_spy.call_args.kwargs["force_local_server"] is True
+
+    def test_install__non_interactive__errors(self):
+        runner = CliRunner()
+        with (
+            patch.object(
+                mcp_cli.interactive_helpers, "is_interactive", return_value=False
+            ),
+            patch.object(mcp_cli.mcp_installer, "setup_mcp_server") as setup_spy,
+        ):
+            result = runner.invoke(cli, ["mcp", "configure"])
+
+        assert result.exit_code != 0
+        assert "interactive terminal" in result.output
+        setup_spy.assert_not_called()
+
+    def test_install__no_config_user_declines__errors(self):
+        runner = CliRunner()
+        with (
+            patch.object(
+                mcp_cli.opik_config, "OpikConfig", return_value=_config(api_key=None)
+            ),
+            patch.object(
+                mcp_cli.interactive_helpers, "is_interactive", return_value=True
+            ),
+            patch.object(
+                mcp_cli.configure_cli, "run_interactive_configure"
+            ) as configure_spy,
+            patch.object(mcp_cli.mcp_installer, "setup_mcp_server") as setup_spy,
+        ):
+            result = runner.invoke(cli, ["mcp", "configure"], input="n\n")
+
+        assert result.exit_code != 0
+        assert "opik configure" in result.output
+        configure_spy.assert_not_called()
+        setup_spy.assert_not_called()
+
+    def test_install__no_config_user_accepts__runs_configure_then_installs(self):
+        runner = CliRunner()
+        configs = iter([_config(api_key=None), _config(api_key="new-key")])
+        with (
+            patch.object(
+                mcp_cli.opik_config, "OpikConfig", side_effect=lambda: next(configs)
+            ),
+            patch.object(
+                mcp_cli.interactive_helpers, "is_interactive", return_value=True
+            ),
+            patch.object(
+                mcp_cli.configure_cli, "run_interactive_configure"
+            ) as configure_spy,
+            patch.object(mcp_cli.mcp_installer, "setup_mcp_server") as setup_spy,
+        ):
+            result = runner.invoke(cli, ["mcp", "configure"], input="y\n")
+
+        assert result.exit_code == 0
+        configure_spy.assert_called_once_with(install_mcp=False)
+        setup_spy.assert_called_once()
+        assert setup_spy.call_args.kwargs["api_key"] == "new-key"
+
+    def test_status__lists_sdk_env_and_host_drift(self):
+        runner = CliRunner()
+        host = mcp_cli.mcp_status.HostStatus(
+            display_name="Claude Code",
+            config_path=pathlib.Path("/home/u/.claude.json"),
+            detected=True,
+            registered=True,
+            transport=mcp_cli.mcp_status.TRANSPORT_LOCAL,
+            points_to="http://localhost:5173/api/",
+            workspace="default",
+            in_sync=False,
+        )
+        with (
+            patch.object(
+                mcp_cli.opik_config, "OpikConfig", return_value=_config(api_key="key")
+            ),
+            patch.object(
+                mcp_cli.mcp_status, "collect_host_statuses", return_value=[host]
+            ),
+        ):
+            result = runner.invoke(cli, ["mcp", "status"])
+
+        assert result.exit_code == 0
+        assert "Your Opik configuration" in result.output
+        assert "configured for 1 AI assistant" in result.output
+        assert "Claude Code" in result.output
+        assert "OUT OF SYNC with your Opik configuration" in result.output
+        assert "http://localhost:5173/api/" in result.output
+        assert "default" in result.output
+
+    def test_status__none_configured__suggests_configure(self):
+        runner = CliRunner()
+        with (
+            patch.object(
+                mcp_cli.opik_config, "OpikConfig", return_value=_config(api_key="key")
+            ),
+            patch.object(mcp_cli.mcp_status, "collect_host_statuses", return_value=[]),
+        ):
+            result = runner.invoke(cli, ["mcp", "status"])
+
+        assert result.exit_code == 0
+        assert "not configured for any AI assistant" in result.output
+        assert "opik mcp configure" in result.output
+
+    def test_install__local_without_api_key__proceeds(self):
+        runner = CliRunner()
+        with (
+            patch.object(
+                mcp_cli.opik_config,
+                "OpikConfig",
+                return_value=_config(
+                    api_key=None, url_override="http://localhost:5173/api/"
+                ),
+            ),
+            patch.object(
+                mcp_cli.interactive_helpers, "is_interactive", return_value=True
+            ),
+            patch.object(mcp_cli.mcp_installer, "setup_mcp_server") as setup_spy,
+        ):
+            result = runner.invoke(cli, ["mcp", "configure"])
+
+        assert result.exit_code == 0
+        setup_spy.assert_called_once()
+        assert setup_spy.call_args.kwargs["use_local"] is True

sdks/python/tests/unit/cli/test_migrate_dataset_checkpoint.py

@@ -0,0 +1,553 @@
+"""Tests for ``opik migrate dataset`` checkpoint/resume (OPIK-7168).
+
+Two layers:
+
+* ``TestMigrationCheckpoint`` -- the checkpoint store in isolation
+  (``opik.cli.migrate.checkpoint``): key derivation, atomic round-trip,
+  in-flight tracking, corrupt/foreign-schema tolerance, delete lifecycle.
+
+* ``TestCascadeResume`` -- the cascade's resume behaviour
+  (``cascade_experiments`` with a ``checkpoint``): skip already-completed
+  experiments, re-migrate an interrupted experiment after deleting its partial
+  destination data, and seed the progress callback so a resumed run reports the
+  right completed count instead of starting at 0.
+
+The cascade tests reuse the elaborate REST/client fakes from
+``test_migrate_dataset_experiments_cascade`` so the resume behaviour is
+exercised against the same call surface the real cascade drives.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, List, Tuple
+from unittest.mock import MagicMock
+
+import pytest
+
+from opik.cli.migrate import checkpoint as checkpoint_module
+from opik.cli.migrate.checkpoint import (
+    SCHEMA_VERSION,
+    MigrationCheckpoint,
+    checkpoint_key,
+    checkpoint_path,
+    load_or_create,
+)
+from opik.cli.migrate.datasets.experiments import cascade_experiments
+
+from .test_migrate_dataset_experiments_cascade import (
+    _Experiment,
+    _ExperimentItem,
+    _Trace,
+    _audit,
+    _cascade_rest_client,
+    _client_with_recreate_capture,
+)
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint store (unit)
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(autouse=True)
+def _isolate_checkpoint_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+    """Point the checkpoint store at a tmp dir instead of the real ~/.opik.
+
+    The checkpoint now lives at a fixed per-user path (``checkpoint_dir``),
+    independent of cwd/--audit-log. Redirecting it here keeps the unit tests
+    hermetic and off the developer's home directory.
+    """
+    cp_dir = tmp_path / "migrate-checkpoints"
+    cp_dir.mkdir(parents=True, exist_ok=True)
+    monkeypatch.setattr(checkpoint_module, "checkpoint_dir", lambda: cp_dir)
+
+
+class TestMigrationCheckpoint:
+    def test_checkpoint_key__stable_and_distinct_per_tuple(self) -> None:
+        key = checkpoint_key("ws", "proj", "ds")
+        # Deterministic: same tuple -> same key across calls.
+        assert key == checkpoint_key("ws", "proj", "ds")
+        # Distinct tuples -> distinct keys, including the ambiguous shift of a
+        # separator across the three components (the \x00 join prevents
+        # "a"+"bc" colliding with "ab"+"c").
+        assert checkpoint_key("a", "bc", "d") != checkpoint_key("ab", "c", "d")
+
+    def test_checkpoint_path__lives_in_checkpoint_dir_keyed_by_hash(self) -> None:
+        key = checkpoint_key("ws", "proj", "ds")
+        path = checkpoint_path(key)
+        assert path.parent == checkpoint_module.checkpoint_dir()
+        assert path.name == f"opik-migrate-checkpoint-{key}.json"
+
+    def test_load_or_create__no_file__starts_fresh(self, tmp_path: Path) -> None:
+        cp = load_or_create(
+            workspace="ws",
+            project="proj",
+            dataset="ds",
+        )
+        assert cp.completed_count == 0
+        assert cp.in_flight is None
+        assert not cp.path.exists()
+
+    def test_flush_then_load__round_trips_completed_and_in_flight(
+        self, tmp_path: Path
+    ) -> None:
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        cp.total_experiments = 3
+        cp.mark_in_flight(
+            "src-exp-2", experiment_name="exp-2", dest_dataset_id="dest-ds"
+        )
+        cp.record_dest_trace_ids(["dest-trace-a", "dest-trace-b"])
+        cp.mark_completed("src-exp-1")
+        cp.flush()
+
+        reloaded = load_or_create(workspace="ws", project="proj", dataset="ds")
+        assert reloaded.total_experiments == 3
+        assert reloaded.completed_experiment_ids == {"src-exp-1"}
+        # ``mark_completed`` clears in_flight, so only the completed set survives.
+        assert reloaded.in_flight is None
+
+    def test_flush_preserves_in_flight_when_not_completed(self, tmp_path: Path) -> None:
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        cp.mark_in_flight(
+            "src-exp-2", experiment_name="exp-2", dest_dataset_id="dest-ds"
+        )
+        cp.record_dest_trace_ids(["dest-trace-a"])
+        cp.flush()
+
+        reloaded = load_or_create(workspace="ws", project="proj", dataset="ds")
+        assert reloaded.in_flight is not None
+        assert reloaded.in_flight.source_experiment_id == "src-exp-2"
+        assert reloaded.in_flight.experiment_name == "exp-2"
+        assert reloaded.in_flight.dest_dataset_id == "dest-ds"
+        assert reloaded.in_flight.dest_trace_ids == ["dest-trace-a"]
+
+    def test_load_or_create__corrupt_file__starts_fresh(self, tmp_path: Path) -> None:
+        # A truncated / unparseable checkpoint (e.g. an interrupted write on a
+        # filesystem without atomic replace) must be treated as "no progress",
+        # not crash the migration.
+        key = checkpoint_key("ws", "proj", "ds")
+        checkpoint_path(key).write_text('{"schema_version": 1, "comp')
+
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        assert cp.completed_count == 0
+        assert cp.in_flight is None
+
+    def test_load_or_create__non_bool_dataset_phase_done__starts_fresh(
+        self, tmp_path: Path
+    ) -> None:
+        # ``dataset_phase_done`` gates resume-vs-full-plan, so a truthy non-bool
+        # (the string "false", a non-empty list) must NOT be coerced to True and
+        # force a resume that skips rename/create/replay. A non-bool is
+        # malformed -> fresh.
+        key = checkpoint_key("ws", "proj", "ds")
+        for bad in ('"false"', "[1]", "1"):
+            # Use the CURRENT schema so the load reaches the non-bool guard
+            # rather than short-circuiting on a schema-version mismatch.
+            checkpoint_path(key).write_text(
+                f'{{"schema_version": {SCHEMA_VERSION}, "dataset_phase_done": {bad}}}'
+            )
+            cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+            assert cp.dataset_phase_done is False
+
+    def test_load_or_create__malformed_in_flight__starts_fresh(
+        self, tmp_path: Path
+    ) -> None:
+        # A checkpoint with the right schema_version but a wrong-shaped
+        # ``in_flight`` (here a bare string, or a dict missing the required
+        # ``source_experiment_id``) must fall back to fresh rather than crash
+        # the CLI with a TypeError/KeyError.
+        key = checkpoint_key("ws", "proj", "ds")
+        for bad_in_flight in ('"not-a-dict"', '{"experiment_name": "x"}'):
+            checkpoint_path(key).write_text(
+                '{"schema_version": 1, "completed_experiment_ids": ["e1"], '
+                f'"in_flight": {bad_in_flight}}}'
+            )
+            cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+            assert cp.completed_count == 0
+            assert cp.in_flight is None
+
+    def test_load_or_create__corrupt_completed_ids__starts_fresh(
+        self, tmp_path: Path
+    ) -> None:
+        # ``set("abc")`` / ``set([1, 2])`` don't raise, so a corrupt
+        # ``completed_experiment_ids`` (a bare string, or a list with non-string
+        # entries) would silently seed the wrong completed set and make the
+        # cascade skip/re-run the wrong experiments. It must fall back to fresh.
+        key = checkpoint_key("ws", "proj", "ds")
+        for bad in ('"abc"', "[1, 2, 3]", "{}"):
+            checkpoint_path(key).write_text(
+                f'{{"schema_version": 1, "completed_experiment_ids": {bad}}}'
+            )
+            cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+            assert cp.completed_experiment_ids == set()
+
+    def test_load_or_create__corrupt_dest_trace_ids__starts_fresh(
+        self, tmp_path: Path
+    ) -> None:
+        # Same silent-corruption guard for the in-flight trace-id list.
+        key = checkpoint_key("ws", "proj", "ds")
+        checkpoint_path(key).write_text(
+            '{"schema_version": 1, "in_flight": '
+            '{"source_experiment_id": "e2", "dest_trace_ids": "trace-1"}}'
+        )
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        assert cp.in_flight is None
+        assert cp.completed_count == 0
+
+    def test_load_or_create__non_string_resume_names__starts_fresh(
+        self, tmp_path: Path
+    ) -> None:
+        # OPIK-7162: source_dataset_id / source_name / temp_dest_name are the
+        # resume handles that flow unvalidated into client.get_dataset(name=...)
+        # and stream_dataset_items(dataset_name=...). A hand-edited non-string
+        # would otherwise reach the API; it must fall back to fresh, matching
+        # the same corrupt-recovery contract as the id collections. Uses the
+        # CURRENT schema so the load reaches the field guard (not the schema
+        # short-circuit).
+        key = checkpoint_key("ws", "proj", "ds")
+        for field in ("source_dataset_id", "source_name", "temp_dest_name"):
+            for bad in ("[1]", "{}", "5"):
+                checkpoint_path(key).write_text(
+                    f'{{"schema_version": {SCHEMA_VERSION}, '
+                    f'"dataset_phase_done": true, "{field}": {bad}}}'
+                )
+                cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+                # Fell back to fresh: no phase-done carried over from the bad file.
+                assert cp.dataset_phase_done is False
+                assert getattr(cp, field) is None
+
+    def test_load_or_create__unresolvable_home__returns_none(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        # A homeless environment (Path.home() raising, as in some CI/containers)
+        # must NOT crash the migration before it starts -- load_or_create
+        # returns None and the caller runs without resume support.
+        def _boom() -> Path:
+            raise RuntimeError("Could not determine home directory")
+
+        monkeypatch.setattr(checkpoint_module, "checkpoint_dir", _boom, raising=True)
+        assert load_or_create(workspace="ws", project="proj", dataset="ds") is None
+
+    def test_flush__write_failure__swallowed_not_raised(
+        self, monkeypatch: pytest.MonkeyPatch
+    ) -> None:
+        # A checkpoint is only a resume aid: a read-only / full disk on flush
+        # must be logged and swallowed, never abort an otherwise-healthy
+        # migration.
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        assert cp is not None
+
+        def _readonly(*_a: Any, **_k: Any) -> None:
+            raise OSError("Read-only file system")
+
+        monkeypatch.setattr(Path, "mkdir", _readonly, raising=True)
+        cp.flush()  # must not raise
+
+    def test_flush_then_load__round_trips_dest_experiment_id(
+        self, tmp_path: Path
+    ) -> None:
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        cp.mark_in_flight(
+            "src-exp-2", experiment_name="exp-2", dest_dataset_id="dest-ds"
+        )
+        cp.record_dest_experiment_id("dest-exp-99")
+        cp.flush()
+
+        reloaded = load_or_create(workspace="ws", project="proj", dataset="ds")
+        assert reloaded.in_flight is not None
+        assert reloaded.in_flight.dest_experiment_id == "dest-exp-99"
+
+    def test_load_or_create__foreign_schema_version__starts_fresh(
+        self, tmp_path: Path
+    ) -> None:
+        key = checkpoint_key("ws", "proj", "ds")
+        checkpoint_path(key).write_text(
+            '{"schema_version": 999, "completed_experiment_ids": ["src-exp-1"]}'
+        )
+
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        # A newer schema we can't interpret is ignored -> fresh start.
+        assert cp.completed_count == 0
+
+    def test_delete__removes_file_and_is_idempotent(self, tmp_path: Path) -> None:
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        cp.flush()
+        assert cp.path.exists()
+        cp.delete()
+        assert not cp.path.exists()
+        # Idempotent: deleting again is a no-op, not an error.
+        cp.delete()
+
+    def test_flush__atomic__no_stray_tmp_file_left(self, tmp_path: Path) -> None:
+        cp = load_or_create(workspace="ws", project="proj", dataset="ds")
+        cp.flush()
+        # The temp file used for the atomic write must have been renamed away.
+        tmp_files = list(tmp_path.glob("*.tmp"))
+        assert tmp_files == []
+
+
+# ---------------------------------------------------------------------------
+# Cascade resume behaviour
+# ---------------------------------------------------------------------------
+
+
+def _two_experiment_rig() -> Tuple[Any, Any]:
+    """Two source experiments, each with one item/trace, wired end-to-end.
+
+    Returns ``(rest_client, client)`` ready to pass to ``cascade_experiments``.
+    """
+    exp1 = _Experiment(id="src-exp-1", name="exp-1", dataset_version_id="src-v-1")
+    exp2 = _Experiment(id="src-exp-2", name="exp-2", dataset_version_id="src-v-1")
+    item1 = _ExperimentItem(
+        id="i1",
+        experiment_id="src-exp-1",
+        trace_id="src-trace-1",
+        dataset_item_id="src-ds-item-1",
+    )
+    item2 = _ExperimentItem(
+        id="i2",
+        experiment_id="src-exp-2",
+        trace_id="src-trace-2",
+        dataset_item_id="src-ds-item-1",
+    )
+    rest_client = _cascade_rest_client(
+        experiments_by_dataset={"src-dataset-1": [exp1, exp2]},
+        items_by_experiment={"exp-1": [item1], "exp-2": [item2]},
+        traces_by_id={
+            "src-trace-1": _Trace(id="src-trace-1"),
+            "src-trace-2": _Trace(id="src-trace-2"),
+        },
+        spans_by_trace={"src-trace-1": [], "src-trace-2": []},
+    )
+    client = _client_with_recreate_capture(rest_client)
+    return rest_client, client
+
+
+def _run_cascade(
+    rest_client: Any, client: Any, checkpoint: MigrationCheckpoint, **overrides: Any
+) -> Tuple[Any, List[Tuple[int, int, str]]]:
+    """Run ``cascade_experiments`` with a progress-capturing callback."""
+    progress_calls: List[Tuple[int, int, str]] = []
+
+    kwargs: dict = dict(
+        source_dataset_id="src-dataset-1",
+        target_dataset_name="MyDataset",
+        target_project_name="DestProject",
+        target_dataset_id="dest-dataset-1",
+        version_remap={"src-v-1": "dest-v-1"},
+        item_id_remap={"src-ds-item-1": "dest-ds-item-1"},
+        audit=_audit(),
+        checkpoint=checkpoint,
+        progress_callback=lambda c, t, label: progress_calls.append((c, t, label)),
+    )
+    kwargs.update(overrides)
+    result = cascade_experiments(client, rest_client, **kwargs)
+    return result, progress_calls
+
+
+class TestCascadeResume:
+    def test_skip_completed__does_not_recreate_already_done_experiment(
+        self, tmp_path: Path
+    ) -> None:
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+        # Pretend exp-1 already migrated on a prior run.
+        cp.mark_completed("src-exp-1")
+
+        result, _ = _run_cascade(rest_client, client, cp)
+
+        # Only the not-yet-done experiment is recreated.
+        assert result.experiments_migrated == 1
+        assert client.create_experiment.call_count == 1
+        created_names = [
+            call.kwargs["name"] for call in client.create_experiment.call_args_list
+        ]
+        assert created_names == ["exp-2"]
+
+    def test_skip_completed__all_done__no_recreation(self, tmp_path: Path) -> None:
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+        cp.mark_completed("src-exp-1")
+        cp.mark_completed("src-exp-2")
+
+        result, _ = _run_cascade(rest_client, client, cp)
+
+        assert result.experiments_migrated == 0
+        client.create_experiment.assert_not_called()
+
+    def test_happyflow__marks_each_experiment_completed_and_flushes(
+        self, tmp_path: Path
+    ) -> None:
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+
+        result, _ = _run_cascade(rest_client, client, cp)
+
+        assert result.experiments_migrated == 2
+        assert cp.completed_experiment_ids == {"src-exp-1", "src-exp-2"}
+        # in_flight is cleared once the last experiment completes.
+        assert cp.in_flight is None
+        # Progress was flushed to disk (a re-run would see both done).
+        reloaded = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+        assert reloaded.completed_experiment_ids == {"src-exp-1", "src-exp-2"}
+
+    def test_re_migrate_incomplete__deletes_partial_dest_data_before_rerun(
+        self, tmp_path: Path
+    ) -> None:
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+        # exp-1 finished on the prior run; exp-2 was interrupted mid-flight with
+        # two destination traces and its destination experiment row already
+        # created (id recorded before creation).
+        cp.mark_completed("src-exp-1")
+        cp.mark_in_flight(
+            "src-exp-2", experiment_name="exp-2", dest_dataset_id="dest-dataset-1"
+        )
+        cp.record_dest_trace_ids(["stale-dest-trace-1", "stale-dest-trace-2"])
+        cp.record_dest_experiment_id("stale-dest-exp")
+
+        _run_cascade(rest_client, client, cp)
+
+        # Partial traces were deleted (spans cascade on the BE), so they don't
+        # duplicate on the re-run.
+        rest_client.traces.delete_traces.assert_called_once()
+        deleted_trace_ids = rest_client.traces.delete_traces.call_args.kwargs["ids"]
+        assert set(deleted_trace_ids) == {"stale-dest-trace-1", "stale-dest-trace-2"}
+        # The exact recorded destination experiment row was deleted BY ID --
+        # cleanup never does a name lookup, so no same-named peer can be hit.
+        rest_client.experiments.delete_experiments_by_id.assert_called_once()
+        deleted_exp_ids = (
+            rest_client.experiments.delete_experiments_by_id.call_args.kwargs["ids"]
+        )
+        assert deleted_exp_ids == ["stale-dest-exp"]
+        # in_flight is cleared after cleanup so a further crash won't re-delete.
+        assert cp.in_flight is None
+
+    def test_re_migrate_incomplete__no_experiment_row_yet__only_traces_deleted(
+        self, tmp_path: Path
+    ) -> None:
+        # Interruption landed after traces were flushed but before the
+        # experiment row was created: dest_experiment_id is None, so cleanup
+        # deletes the traces and skips the experiment delete entirely (no
+        # name lookup that could hit a peer).
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+        cp.mark_in_flight(
+            "src-exp-1", experiment_name="exp-1", dest_dataset_id="dest-dataset-1"
+        )
+        cp.record_dest_trace_ids(["stale-dest-trace-1"])
+        # no record_dest_experiment_id -> dest_experiment_id stays None
+
+        _run_cascade(rest_client, client, cp)
+
+        rest_client.traces.delete_traces.assert_called_once()
+        rest_client.experiments.delete_experiments_by_id.assert_not_called()
+
+    def test_re_migrate_incomplete__records_dest_ids_before_backend_write(
+        self, tmp_path: Path
+    ) -> None:
+        # OPIK-7168 (#533/#536): the destination trace ids and experiment id
+        # must be flushed to the checkpoint BEFORE the backend write that
+        # persists them, so a crash in that window still leaves them recorded
+        # for the next run's cleanup. We assert the checkpoint file on disk
+        # already carries this experiment's dest trace ids by the time the
+        # trace-flush happens.
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+
+        flushed_state: dict = {}
+        original_flush = client.flush
+
+        def _capture_on_first_flush() -> None:
+            # On the first client.flush() (the trace flush), read back the
+            # checkpoint from disk and snapshot its in-flight trace ids.
+            if "dest_trace_ids" not in flushed_state:
+                reloaded = load_or_create(
+                    workspace="ws",
+                    project="DestProject",
+                    dataset="MyDataset",
+                )
+                flushed_state["dest_trace_ids"] = (
+                    list(reloaded.in_flight.dest_trace_ids)
+                    if reloaded.in_flight
+                    else []
+                )
+            return original_flush()
+
+        client.flush = MagicMock(side_effect=_capture_on_first_flush)
+
+        _run_cascade(rest_client, client, cp)
+
+        # By the time the backend trace flush ran, the checkpoint on disk had
+        # already recorded this experiment's destination trace id.
+        assert len(flushed_state["dest_trace_ids"]) == 1
+
+    def test_progress_seeded__resumed_run_starts_at_completed_count(
+        self, tmp_path: Path
+    ) -> None:
+        rest_client, client = _two_experiment_rig()
+        cp = load_or_create(
+            workspace="ws",
+            project="DestProject",
+            dataset="MyDataset",
+        )
+        cp.mark_completed("src-exp-1")
+
+        _, progress_calls = _run_cascade(rest_client, client, cp)
+
+        # First progress tick reports 1 already-completed of 2 total (not 0),
+        # so the bar opens at 50% rather than restarting.
+        assert progress_calls[0] == (1, 2, "exp-2")
+        # Final tick snaps to total/total "done".
+        assert progress_calls[-1] == (2, 2, "done")
+
+    def test_no_checkpoint__cascade_unchanged(self) -> None:
+        # Without a checkpoint the cascade behaves exactly as before: both
+        # experiments migrate, no delete calls, progress counts from 0.
+        rest_client, client = _two_experiment_rig()
+        progress_calls: List[Tuple[int, int, str]] = []
+
+        result = cascade_experiments(
+            client,
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            version_remap={"src-v-1": "dest-v-1"},
+            item_id_remap={"src-ds-item-1": "dest-ds-item-1"},
+            audit=_audit(),
+            progress_callback=lambda c, t, label: progress_calls.append((c, t, label)),
+        )
+
+        assert result.experiments_migrated == 2
+        rest_client.traces.delete_traces.assert_not_called()
+        assert progress_calls[0] == (0, 2, "exp-1")

sdks/python/tests/unit/cli/test_migrate_dataset_optimizations_cascade.py

@@ -0,0 +1,496 @@
+"""Tests for ``opik migrate dataset`` Slice 5 — optimization cascade.
+
+Slice 5 sits between Slice 2 (``ReplayVersions``) and Slice 3
+(``CascadeExperiments``). It recreates every optimization referencing
+the source dataset under the destination project and populates
+``plan.optimization_id_remap`` so the experiment cascade can re-point
+each experiment's ``optimization_id`` FK.
+
+Scope this module covers (ticket AC for OPIK-6532):
+
+* Optimization fidelity round-trip (name / objective_name / status /
+  metadata / studio_config carried verbatim)
+* READ-ONLY aggregates (num_trials / feedback_scores / baseline_* /
+  best_* / total_optimization_cost) not forwarded
+* Fresh client-side UUID minted per source optimization; ``id_remap``
+  populated with source -> destination
+* Empty optimization (zero trials) round-trips and produces a remap entry
+* Zero-optimization datasets are a no-op (returns empty result, no
+  REST call)
+* Per-optimization ``migrate_optimization`` audit record emitted
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+from unittest.mock import MagicMock
+
+import pytest
+
+from opik.cli.migrate.audit import AuditLog
+from opik.cli.migrate.datasets import optimizations as optimizations_module
+from opik.cli.migrate.datasets.optimizations import cascade_optimizations
+
+
+# ---------------------------------------------------------------------------
+# Stand-in wire shapes for the cascade's reads.
+#
+# The cascade reads ``OptimizationPublic`` from ``find_optimizations`` and
+# writes via ``create_optimization``. Constructing the real pydantic
+# models would impose ~20 fields per row; the cascade only inspects a
+# handful, so we use plain objects with the fields it actually reads.
+# ---------------------------------------------------------------------------
+
+
+class _Optimization:
+    def __init__(
+        self,
+        *,
+        id: str,
+        name: Optional[str] = "opt",
+        objective_name: str = "accuracy",
+        status: str = "completed",
+        dataset_name: str = "MyDataset",
+        dataset_id: Optional[str] = "src-dataset-1",
+        metadata: Optional[Any] = None,
+        studio_config: Optional[Any] = None,
+        last_updated_at: Optional[Any] = None,
+        # READ-ONLY aggregates -- attached so we can assert the cascade
+        # does NOT forward them on the create call.
+        num_trials: Optional[int] = None,
+        feedback_scores: Optional[List[Any]] = None,
+        experiment_scores: Optional[List[Any]] = None,
+        baseline_objective_score: Optional[float] = None,
+        best_objective_score: Optional[float] = None,
+        baseline_duration: Optional[float] = None,
+        best_duration: Optional[float] = None,
+        baseline_cost: Optional[float] = None,
+        best_cost: Optional[float] = None,
+        total_optimization_cost: Optional[float] = None,
+    ) -> None:
+        self.id = id
+        self.name = name
+        self.objective_name = objective_name
+        self.status = status
+        self.dataset_name = dataset_name
+        self.dataset_id = dataset_id
+        self.metadata = metadata
+        self.studio_config = studio_config
+        self.last_updated_at = last_updated_at
+        self.num_trials = num_trials
+        self.feedback_scores = feedback_scores
+        self.experiment_scores = experiment_scores
+        self.baseline_objective_score = baseline_objective_score
+        self.best_objective_score = best_objective_score
+        self.baseline_duration = baseline_duration
+        self.best_duration = best_duration
+        self.baseline_cost = baseline_cost
+        self.best_cost = best_cost
+        self.total_optimization_cost = total_optimization_cost
+
+
+class _Page:
+    def __init__(self, content: List[_Optimization]) -> None:
+        self.content = content
+
+
+def _audit() -> AuditLog:
+    """Real AuditLog so the cascade's per-optimization ``record`` calls
+    land in a list we can assert on. The audit log carries no side
+    effects until ``write`` is called, so using the real type in tests
+    is cheaper than building a coupled mock."""
+    return AuditLog(command="migrate dataset", args={})
+
+
+def _cascade_rest_client(
+    pages: List[_Page],
+    *,
+    create_side_effect: Optional[Any] = None,
+) -> MagicMock:
+    """Build a rest_client mock for the optimization cascade.
+
+    ``find_optimizations`` is called once per page until a short page
+    breaks the loop. ``create_optimization`` is captured so individual
+    tests can assert on its kwargs.
+    """
+    rest_client = MagicMock()
+    rest_client.optimizations.find_optimizations.side_effect = pages
+    if create_side_effect is not None:
+        rest_client.optimizations.create_optimization.side_effect = create_side_effect
+    else:
+        rest_client.optimizations.create_optimization.return_value = None
+    return rest_client
+
+
+# ---------------------------------------------------------------------------
+# Cascade tests
+# ---------------------------------------------------------------------------
+
+
+class TestCascadeOptimizations:
+    def test_no_source_optimizations__noop_no_remap_no_create_calls(self) -> None:
+        # Empty result page -> no create_optimization calls, empty
+        # remap, counters at zero. Slice 5 always emits the action; the
+        # cascade handles the zero case quietly.
+        rest_client = _cascade_rest_client([_Page([])])
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        assert result.id_remap == {}
+        assert result.optimizations_total == 0
+        assert result.optimizations_migrated == 0
+        assert result.optimizations_skipped == 0
+        rest_client.optimizations.create_optimization.assert_not_called()
+
+    def test_one_optimization__fidelity_fields_forwarded_and_remap_populated(
+        self,
+    ) -> None:
+        # Full fidelity copy of the fields the BE round-trips:
+        # name, objective_name, status, metadata, studio_config,
+        # last_updated_at. Dest dataset_name + project_name are taken
+        # from the action, NOT from the source row (the source name was
+        # renamed to <name>_v1 by the time the cascade runs).
+        opt = _Optimization(
+            id="src-opt-1",
+            name="overnight-tune",
+            objective_name="f1",
+            status="completed",
+            metadata={"trial_count": 12, "user_note": "kept"},
+            studio_config=None,
+            num_trials=12,
+            baseline_objective_score=0.5,
+            best_objective_score=0.81,
+            total_optimization_cost=4.20,
+        )
+        rest_client = _cascade_rest_client([_Page([opt])])
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        rest_client.optimizations.create_optimization.assert_called_once()
+        kwargs = rest_client.optimizations.create_optimization.call_args.kwargs
+        # Fresh destination id -- not the source id.
+        assert kwargs["id"] != "src-opt-1"
+        assert kwargs["id"] == result.id_remap["src-opt-1"]
+        # Fidelity fields forwarded verbatim.
+        assert kwargs["name"] == "overnight-tune"
+        assert kwargs["objective_name"] == "f1"
+        assert kwargs["status"] == "completed"
+        assert kwargs["metadata"] == {"trial_count": 12, "user_note": "kept"}
+        # Destination scoping comes from the action, not the source row.
+        assert kwargs["dataset_name"] == "MyDataset"
+        assert kwargs["project_name"] == "DestProject"
+
+    def test_readonly_aggregates__not_forwarded_on_create(self) -> None:
+        # ``num_trials`` / ``feedback_scores`` / ``baseline_*`` / ``best_*``
+        # / ``total_optimization_cost`` are recomputed by the BE at read
+        # time from constituent experiments. The cascade must NOT forward
+        # them -- if it did, they'd be silently dropped on write today
+        # but might persist (with stale numbers) under a future BE
+        # schema change.
+        opt = _Optimization(
+            id="src-opt-1",
+            num_trials=99,
+            feedback_scores=[MagicMock()],
+            experiment_scores=[MagicMock()],
+            baseline_objective_score=0.1,
+            best_objective_score=0.9,
+            baseline_duration=10.0,
+            best_duration=2.0,
+            baseline_cost=1.0,
+            best_cost=0.5,
+            total_optimization_cost=42.0,
+        )
+        rest_client = _cascade_rest_client([_Page([opt])])
+
+        cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        kwargs = rest_client.optimizations.create_optimization.call_args.kwargs
+        for forbidden in (
+            "num_trials",
+            "feedback_scores",
+            "experiment_scores",
+            "baseline_objective_score",
+            "best_objective_score",
+            "baseline_duration",
+            "best_duration",
+            "baseline_cost",
+            "best_cost",
+            "total_optimization_cost",
+        ):
+            assert forbidden not in kwargs, (
+                f"{forbidden} is BE-computed; cascade must not forward it"
+            )
+
+    def test_empty_optimization__zero_trials__still_migrated(self) -> None:
+        # Optimization with zero constituent experiments is still a
+        # user-visible row in the Optimization Studio UI; migrating it
+        # preserves the user's setup even if no trials ran yet.
+        opt = _Optimization(id="src-opt-1", num_trials=0, name="never-ran")
+        rest_client = _cascade_rest_client([_Page([opt])])
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        assert result.optimizations_migrated == 1
+        assert "src-opt-1" in result.id_remap
+        rest_client.optimizations.create_optimization.assert_called_once()
+
+    def test_multiple_optimizations__each_gets_distinct_destination_id(
+        self,
+    ) -> None:
+        opt1 = _Optimization(id="src-opt-1", name="run-a")
+        opt2 = _Optimization(id="src-opt-2", name="run-b")
+        opt3 = _Optimization(id="src-opt-3", name="run-c")
+        rest_client = _cascade_rest_client([_Page([opt1, opt2, opt3])])
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        assert result.optimizations_migrated == 3
+        assert set(result.id_remap.keys()) == {
+            "src-opt-1",
+            "src-opt-2",
+            "src-opt-3",
+        }
+        # Destination ids must be unique (fresh UUIDs).
+        assert len(set(result.id_remap.values())) == 3
+
+    def test_audit__per_optimization_record_emitted(self) -> None:
+        # Each successful migration emits a ``migrate_optimization`` audit
+        # record carrying source_id -> destination_id mapping plus
+        # name / objective_name / status so users can trace what moved.
+        opt = _Optimization(
+            id="src-opt-1",
+            name="overnight",
+            objective_name="accuracy",
+            status="completed",
+        )
+        rest_client = _cascade_rest_client([_Page([opt])])
+        audit = _audit()
+
+        cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=audit,
+        )
+
+        migrate_records = [
+            r for r in audit.actions if r["type"] == "migrate_optimization"
+        ]
+        assert len(migrate_records) == 1
+        record = migrate_records[0]
+        # Audit-level status conveys action success/failure -- must not
+        # be shadowed by the optimization's own status.
+        assert record["status"] == "ok"
+        assert record["source_id"] == "src-opt-1"
+        assert record["destination_id"] != "src-opt-1"
+        assert record["name"] == "overnight"
+        assert record["objective_name"] == "accuracy"
+        # Optimization's own status lives under a non-shadowing key.
+        assert record["optimization_status"] == "completed"
+
+    def test_pagination__exhausts_pages_until_short(self) -> None:
+        # find_optimizations is paginated; the cascade walks until a
+        # short page (< page_size). With page_size=100 in production,
+        # two pages of 100 and a final short page of 50 should fire
+        # three reads in order.
+        page_size = optimizations_module._OPTIMIZATION_PAGE_SIZE
+        page_a = _Page([_Optimization(id=f"src-{i}") for i in range(page_size)])
+        page_b = _Page(
+            [_Optimization(id=f"src-{page_size + i}") for i in range(page_size)]
+        )
+        page_c = _Page([_Optimization(id=f"src-{2 * page_size + i}") for i in range(3)])
+        rest_client = _cascade_rest_client([page_a, page_b, page_c])
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        assert result.optimizations_total == 2 * page_size + 3
+        # 3 reads + N writes.
+        assert rest_client.optimizations.find_optimizations.call_count == 3
+
+    def test_create_optimization_raises__exception_propagates(self) -> None:
+        # The BE returns 409 on id conflict (we mint client-side UUIDs
+        # so within-run collisions are infinitesimal but possible across
+        # retries). Surface as a cascade error rather than silently
+        # continuing; the audit log's surrounding ``failed`` entry will
+        # capture partial progress via the umbrella action.
+        opt = _Optimization(id="src-opt-1")
+        rest_client = _cascade_rest_client(
+            [_Page([opt])],
+            create_side_effect=RuntimeError("409 conflict"),
+        )
+
+        with pytest.raises(RuntimeError, match="409 conflict"):
+            cascade_optimizations(
+                rest_client,
+                source_dataset_id="src-dataset-1",
+                target_dataset_name="MyDataset",
+                target_project_name="DestProject",
+                audit=_audit(),
+            )
+
+    def test_source_missing_id__skipped_not_remapped(self) -> None:
+        # Defensive path: BE returns an optimization without an id.
+        # Skip rather than fail -- other optimizations are still
+        # recoverable, and a non-zero ``optimizations_skipped`` counter
+        # in the audit surfaces the malformed row for investigation.
+        ok = _Optimization(id="src-opt-1")
+        bad = _Optimization(id="")
+        rest_client = _cascade_rest_client([_Page([ok, bad])])
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-dataset-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+
+        assert result.optimizations_migrated == 1
+        assert result.optimizations_skipped == 1
+        assert list(result.id_remap.keys()) == ["src-opt-1"]
+
+
+class TestCascadeOptimizationsWithExperiments:
+    """End-to-end planner-shape assertion: the optimization remap built
+    by Slice 5 is observable to Slice 3's experiment cascade via the
+    shared ``MigrationPlan``. The actual experiment-cascade FK forwarding
+    is asserted in ``test_migrate_dataset_experiments_cascade``; this
+    test only pins the wiring (action ordering + remap propagation).
+    """
+
+    def test_plan_remap__populated_by_cascade__readable_by_subsequent_actions(
+        self,
+    ) -> None:
+        from opik.cli.migrate.datasets.planner import MigrationPlan
+        from opik.cli.migrate.datasets.resolver import ResolvedDataset
+
+        opt = _Optimization(id="src-opt-1", name="trained")
+        rest_client = _cascade_rest_client([_Page([opt])])
+        plan = MigrationPlan(
+            source=ResolvedDataset(
+                id="src-1",
+                name="MyDataset",
+                project_name=None,
+                description=None,
+                visibility=None,
+                tags=None,
+                type="dataset",
+            ),
+            target_name="MyDataset",
+            to_project="DestProject",
+        )
+
+        result = cascade_optimizations(
+            rest_client,
+            source_dataset_id="src-1",
+            target_dataset_name="MyDataset",
+            target_project_name="DestProject",
+            audit=_audit(),
+        )
+        plan.optimization_id_remap.update(result.id_remap)
+
+        # Subsequent CascadeExperiments would read plan.optimization_id_remap;
+        # the wiring works as long as result.id_remap propagates here.
+        assert plan.optimization_id_remap == result.id_remap
+        assert "src-opt-1" in plan.optimization_id_remap
+
+
+# ---------------------------------------------------------------------------
+# Studio config Read -> Write reconstruction
+# ---------------------------------------------------------------------------
+
+
+class TestStudioConfigReconstruction:
+    """The Fern-generated Read and Write variants of
+    ``OptimizationStudioConfig`` are structurally identical but
+    nominally distinct; the cascade has to reconstruct the Write
+    variant from the Read variant so the typed ``create_optimization``
+    parameter is satisfied. Pin the round-trip so any future Fern
+    regeneration that changes the shape on one side gets caught.
+    """
+
+    def test_none_in_none_out(self) -> None:
+        from opik.cli.migrate.datasets.optimizations import (
+            _studio_config_public_to_write,
+        )
+
+        assert _studio_config_public_to_write(None) is None
+
+    def test_round_trip__fields_preserved(self) -> None:
+        from opik.cli.migrate.datasets.optimizations import (
+            _studio_config_public_to_write,
+        )
+        from opik.rest_api.types.optimization_studio_config_public import (
+            OptimizationStudioConfigPublic,
+        )
+        from opik.rest_api.types.studio_evaluation_public import (
+            StudioEvaluationPublic,
+        )
+        from opik.rest_api.types.studio_llm_model_public import (
+            StudioLlmModelPublic,
+        )
+        from opik.rest_api.types.studio_message_public import StudioMessagePublic
+        from opik.rest_api.types.studio_metric_public import StudioMetricPublic
+        from opik.rest_api.types.studio_optimizer_public import (
+            StudioOptimizerPublic,
+        )
+        from opik.rest_api.types.studio_prompt_public import StudioPromptPublic
+
+        public = OptimizationStudioConfigPublic(
+            dataset_name="MyDataset",
+            prompt=StudioPromptPublic(
+                messages=[StudioMessagePublic(role="system", content="you are helpful")]
+            ),
+            llm_model=StudioLlmModelPublic(
+                model="gpt-4o", parameters={"temperature": 0.2}
+            ),
+            evaluation=StudioEvaluationPublic(
+                metrics=[StudioMetricPublic(type="accuracy")]
+            ),
+            optimizer=StudioOptimizerPublic(type="grid", parameters={"steps": 5}),
+        )
+        write = _studio_config_public_to_write(public)
+        assert write is not None
+        # Round-trip via model_dump -- the typed Write variant accepts
+        # extra fields so this also covers any forward-compatible BE
+        # additions on the Public variant.
+        write_dump: Dict[str, Any] = write.model_dump()
+        public_dump: Dict[str, Any] = public.model_dump()
+        assert write_dump == public_dump

sdks/python/tests/unit/cli/test_migrate_dataset_planner.py

@@ -0,0 +1,950 @@
+"""Planner + CLI-help tests for ``opik migrate dataset``.
+
+The planner cases (conflict, project-not-found, default flow ordering)
+live here; the meaty version-replay tests live in
+``test_migrate_dataset_version_replay.py`` and the cascade tests live in
+``test_migrate_dataset_experiments_cascade.py``.
+
+Shared helpers (``_DatasetRow``, ``_Page``, ``_planner_rest_client``)
+come from ``_migrate_helpers``.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+from click.testing import CliRunner
+
+from opik.cli import cli
+from opik.cli.migrate.audit import AuditLog
+from opik.cli.migrate.checkpoint import MigrationCheckpoint
+from opik.cli.migrate.datasets import planner as planner_module
+from opik.cli.migrate.datasets.planner import TEMP_MIGRATION_MARKER_TAG
+from opik.cli.migrate.datasets.resume import ReconstructedRemaps
+from opik.cli.migrate.errors import (
+    ConflictError,
+    DatasetNotFoundError,
+    MigrationError,
+    ProjectNotFoundError,
+)
+
+from ._migrate_helpers import (
+    _build_fake_client,
+    _DatasetRow,
+    _Page,
+    _planner_client,
+    _planner_rest_client,
+)
+
+
+# ---------------------------------------------------------------------------
+# Elapsed-time formatter
+# ---------------------------------------------------------------------------
+
+
+class TestFormatElapsed:
+    """Pin the wall-clock duration renderer used in the migrate success /
+    failure lines. Sub-minute → one decimal of seconds; past a minute →
+    integer ``Mm Ss`` or ``Hh Mm Ss`` (no fractional seconds).
+    """
+
+    def test_sub_minute__one_decimal_seconds(self) -> None:
+        from opik.cli.migrate.main import _format_elapsed
+
+        assert _format_elapsed(0.0) == "0.0s"
+        assert _format_elapsed(12.34) == "12.3s"
+        assert _format_elapsed(59.99) == "60.0s"
+
+    def test_minute_range__integer_m_s(self) -> None:
+        from opik.cli.migrate.main import _format_elapsed
+
+        assert _format_elapsed(60.0) == "1m 0s"
+        assert _format_elapsed(125.7) == "2m 5s"
+
+    def test_hour_range__integer_h_m_s(self) -> None:
+        from opik.cli.migrate.main import _format_elapsed
+
+        assert _format_elapsed(3600.0) == "1h 0m 0s"
+        assert _format_elapsed(3725.0) == "1h 2m 5s"
+
+
+# ---------------------------------------------------------------------------
+# OPIK-6599: loud-fail on skipped items
+#
+# When the cascade emits any ``skip`` audit record, the migrate must:
+#   1. Finalize the audit log to ``failed`` (not ``ok``)
+#   2. Print a SKIP_SUMMARY line to stderr (not stdout) so CI gates can
+#      grep without parsing the JSON
+#   3. Exit non-zero
+#
+# Tests below cover ``_finalize_with_skips_or_ok`` directly so they don't
+# need a fully-wired Opik client + REST server stub. The CLI is only the
+# wrapper around this helper.
+# ---------------------------------------------------------------------------
+
+
+class TestFinalizeWithSkipsOrOk:
+    def _make_audit_with_skips(self) -> AuditLog:
+        audit = AuditLog(command="opik migrate dataset", args={})
+        audit.record(
+            type="skip",
+            status="skipped",
+            details={
+                "reason": "items_missing_dataset_item_remap",
+                "experiment_id": "src-exp-1",
+                "experiment_name": "exp-1",
+                "count": 2500,
+                "sample_source_ids": ["src-ds-item-1", "src-ds-item-2"],
+            },
+        )
+        return audit
+
+    def test_skips_present__finalizes_failed_exits_1_stderr_summary(
+        self, tmp_path, capsys
+    ) -> None:
+        from opik.cli.migrate.main import _finalize_with_skips_or_ok
+
+        audit = self._make_audit_with_skips()
+        audit_path = tmp_path / "audit.json"
+
+        with pytest.raises(SystemExit) as exc:
+            _finalize_with_skips_or_ok(
+                audit,
+                audit_path,
+                name="MyDataset",
+                target_label="MyDataset",
+                target_project="DestProject",
+                elapsed_seconds=12.3,
+            )
+
+        # AC 1: non-zero exit code
+        assert exc.value.code == 1
+
+        captured = capsys.readouterr()
+        # AC 3: skip message on stderr (not stdout), with the
+        # machine-parseable SKIP_SUMMARY suffix
+        assert "SKIP_SUMMARY:" in captured.err
+        assert "items_skipped_missing_item=2500" in captured.err
+        assert "experiments_skipped=0" in captured.err
+        assert "items_skipped_missing_trace=0" in captured.err
+        assert "NOT rolled back" in captured.err
+        # Rollback hint names the entities the operator must remove,
+        # the destination project, and the rename-back step on the source.
+        assert "roll back manually" in captured.err
+        assert "DestProject" in captured.err
+        assert "MyDataset_v1" in captured.err
+
+        # AC 2: audit finalized to failed with skip record intact
+        on_disk = json.loads(audit_path.read_text())
+        assert on_disk["status"] == "failed"
+        assert any(a.get("status") == "skipped" for a in on_disk["actions"])
+
+    def test_no_skips__finalizes_ok_no_exit_no_stderr(self, tmp_path, capsys) -> None:
+        from opik.cli.migrate.main import _finalize_with_skips_or_ok
+
+        audit = AuditLog(command="opik migrate dataset", args={})
+        audit.record(type="rename_source", status="ok", details={})
+        audit_path = tmp_path / "audit.json"
+
+        # Happy path returns without raising; happy-path message goes to
+        # stdout, stderr stays clean.
+        _finalize_with_skips_or_ok(
+            audit,
+            audit_path,
+            name="MyDataset",
+            target_label="MyDataset",
+            target_project="DestProject",
+            elapsed_seconds=5.0,
+        )
+
+        captured = capsys.readouterr()
+        assert "SKIP_SUMMARY:" not in captured.err
+        on_disk = json.loads(audit_path.read_text())
+        assert on_disk["status"] == "ok"
+
+    def test_multiple_skip_records__totals_aggregated_by_reason(
+        self, tmp_path, capsys
+    ) -> None:
+        from opik.cli.migrate.main import _finalize_with_skips_or_ok
+
+        # Two experiments, each contributing skips for both reasons.
+        audit = AuditLog(command="opik migrate dataset", args={})
+        for exp_id in ("src-exp-1", "src-exp-2"):
+            audit.record(
+                type="skip",
+                status="skipped",
+                details={
+                    "reason": "items_missing_trace_remap",
+                    "experiment_id": exp_id,
+                    "experiment_name": exp_id,
+                    "count": 7,
+                    "sample_source_ids": [],
+                },
+            )
+            audit.record(
+                type="skip",
+                status="skipped",
+                details={
+                    "reason": "items_missing_dataset_item_remap",
+                    "experiment_id": exp_id,
+                    "experiment_name": exp_id,
+                    "count": 3,
+                    "sample_source_ids": [],
+                },
+            )
+
+        with pytest.raises(SystemExit):
+            _finalize_with_skips_or_ok(
+                audit,
+                tmp_path / "audit.json",
+                name="MyDataset",
+                target_label="MyDataset",
+                target_project="DestProject",
+                elapsed_seconds=1.0,
+            )
+
+        captured = capsys.readouterr()
+        # 7 + 7 = 14 trace skips, 3 + 3 = 6 dataset-item skips
+        assert "items_skipped_missing_trace=14" in captured.err
+        assert "items_skipped_missing_item=6" in captured.err
+
+
+# ---------------------------------------------------------------------------
+# Help text
+# ---------------------------------------------------------------------------
+
+
+class TestMigrateHelp:
+    def test_migrate_group__help_invoked__lists_subcommands(self) -> None:
+        runner = CliRunner()
+        result = runner.invoke(cli, ["migrate", "--help"])
+        assert result.exit_code == 0
+        assert "Migrate Opik entities" in result.output
+        assert "dataset" in result.output
+
+    def test_migrate_dataset__help_invoked__lists_required_flags(self) -> None:
+        runner = CliRunner()
+        result = runner.invoke(cli, ["migrate", "dataset", "--help"])
+        assert result.exit_code == 0
+        assert "--to-project" in result.output
+        assert "--from-project" in result.output
+        assert "--dry-run" in result.output
+
+    def test_migrate_dataset__help_invoked__lists_exclude_experiments(self) -> None:
+        # OPIK-7161 AC: the opt-out flag must be discoverable in --help.
+        runner = CliRunner()
+        result = runner.invoke(cli, ["migrate", "dataset", "--help"])
+        assert result.exit_code == 0
+        assert "--exclude-experiments" in result.output
+
+    def test_migrate_dataset__exclude_experiments__cli_run_skips_and_reports(
+        self, tmp_path
+    ) -> None:
+        # OPIK-7161: exercise the flag through the public Click entrypoint,
+        # not just the finalize helper, so the option -> build_dataset_plan
+        # -> finalize plumbing in migrate_dataset_command is covered end to
+        # end (per .agents/skills/python-sdk/testing.md: test the public API).
+        # The fake client mocks the whole rename/create/replay surface; with
+        # --exclude-experiments the plan carries no cascade actions, so the
+        # command reaches the success finalize with zero experiment work.
+        client, _, _ = _build_fake_client(
+            source_rows=[_DatasetRow(id="src-1", name="MyDataset")],
+            destination_rows=[],
+            items=[{"id": "item-a", "input": "hello"}],
+        )
+        audit_path = tmp_path / "audit.json"
+
+        runner = CliRunner()
+        with patch("opik.cli.migrate.main._build_client", return_value=client):
+            result = runner.invoke(
+                cli,
+                [
+                    "migrate",
+                    "dataset",
+                    "MyDataset",
+                    "--to-project",
+                    "B",
+                    "--exclude-experiments",
+                    "--audit-log",
+                    str(audit_path),
+                ],
+            )
+
+        assert result.exit_code == 0, result.output
+        # User-facing output makes the intentional skip clear.
+        assert "--exclude-experiments" in result.output
+        assert "skipped" in result.output
+        # No experiment cascade ran: the source dataset was never queried for
+        # experiments (find_experiments belongs only to the cascade path).
+        assert client.rest_client.experiments.find_experiments.call_count == 0
+        # Audit log finalized ok and recorded the flag in its args.
+        on_disk = json.loads(audit_path.read_text())
+        assert on_disk["status"] == "ok"
+        assert on_disk["args"]["exclude_experiments"] is True
+        cascade_types = {a.get("type") for a in on_disk["actions"]}
+        assert "cascade_experiments" not in cascade_types
+        assert "cascade_optimizations" not in cascade_types
+
+
+class TestTempDestRenameOnSuccess:
+    """OPIK-7162 acceptance criteria, exercised through the public Click
+    entrypoint: the source keeps its name until the copy succeeds, then the
+    handoff runs (source -> _v1, temp -> original). A mid-run failure leaves
+    the source name untouched, and a re-run after failure is safe/idempotent.
+
+    All cases use ``--exclude-experiments`` so the plan is the minimal
+    Create -> Replay -> Rename -> Promote shape and the assertions stay
+    focused on the handoff, not the cascade.
+    """
+
+    def _run(self, client, tmp_path, extra_args=()):
+        audit_path = tmp_path / "audit.json"
+        runner = CliRunner()
+        with patch("opik.cli.migrate.main._build_client", return_value=client):
+            result = runner.invoke(
+                cli,
+                [
+                    "migrate",
+                    "dataset",
+                    "MyDataset",
+                    "--to-project",
+                    "B",
+                    "--exclude-experiments",
+                    "--audit-log",
+                    str(audit_path),
+                    *extra_args,
+                ],
+            )
+        return result, audit_path
+
+    def test_success__source_renamed_to_v1_and_dest_promoted_to_original(
+        self, tmp_path
+    ) -> None:
+        # AC: a successful migration leaves the destination under the
+        # original name and the source under the _v1 suffix. The two renames
+        # happen via update_dataset PUTs; the destination is created under
+        # the temp name first.
+        client, _, _ = _build_fake_client(
+            source_rows=[_DatasetRow(id="src-1", name="MyDataset")],
+            destination_rows=[],
+            items=[{"id": "item-a", "input": "hello"}],
+        )
+        result, audit_path = self._run(client, tmp_path)
+
+        assert result.exit_code == 0, result.output
+        rest = client.rest_client
+        # Destination created under the temp name (not the final name), and
+        # stamped with the migration marker tag so a future re-run can prove
+        # it's ours before discarding it.
+        create_kwargs = rest.datasets.create_dataset.call_args.kwargs
+        assert create_kwargs["name"] == "MyDataset__migrating"
+        assert TEMP_MIGRATION_MARKER_TAG in (create_kwargs["tags"] or [])
+        # Two rename PUTs: source -> _v1, then temp -> original. The promote PUT
+        # re-passes the source's ORIGINAL tags (marker stripped).
+        rename_calls = [c.kwargs for c in rest.datasets.update_dataset.call_args_list]
+        source_rename = next(c for c in rename_calls if c["id"] == "src-1")
+        assert source_rename["name"] == "MyDataset_v1"
+        promote = next(c for c in rename_calls if c.get("name") == "MyDataset")
+        assert promote["name"] == "MyDataset"
+        # The promote PUT must pass tags as an EXPLICIT list (never None), so the
+        # BE actually overwrites and drops the marker. A live backend treats
+        # ``tags=None`` as "leave unchanged", which would strand the marker on a
+        # source with no tags — so assert the concrete list, not just marker
+        # absence. Source here has no tags -> promote clears with [].
+        assert promote["tags"] == []
+        assert TEMP_MIGRATION_MARKER_TAG not in promote["tags"]
+        # Audit ends ok and records the handoff actions in order.
+        on_disk = json.loads(audit_path.read_text())
+        assert on_disk["status"] == "ok"
+        ok_types = [a["type"] for a in on_disk["actions"] if a.get("status") == "ok"]
+        assert ok_types.index("rename_source") < ok_types.index("promote_destination")
+        assert ok_types.index("create_destination") < ok_types.index("rename_source")
+
+    def test_success__source_tags_preserved_marker_stripped(self, tmp_path) -> None:
+        # When the source has real tags, the temp create adds the marker
+        # alongside them, and the promote re-passes exactly the source's
+        # originals (marker dropped, real tags kept).
+        client, _, _ = _build_fake_client(
+            source_rows=[
+                _DatasetRow(id="src-1", name="MyDataset", tags=["team-a", "prod"])
+            ],
+            destination_rows=[],
+            items=[{"id": "item-a", "input": "hello"}],
+        )
+        result, _ = self._run(client, tmp_path)
+
+        assert result.exit_code == 0, result.output
+        rest = client.rest_client
+        create_tags = rest.datasets.create_dataset.call_args.kwargs["tags"]
+        assert set(create_tags) == {"team-a", "prod", TEMP_MIGRATION_MARKER_TAG}
+        promote = next(
+            c.kwargs
+            for c in rest.datasets.update_dataset.call_args_list
+            if c.kwargs.get("name") == "MyDataset"
+        )
+        assert promote["tags"] == ["team-a", "prod"]
+        assert TEMP_MIGRATION_MARKER_TAG not in promote["tags"]
+
+    def test_midrun_failure__source_name_untouched(self, tmp_path) -> None:
+        # AC: a migration interrupted mid-run leaves the source name
+        # untouched. Blow up the destination create (the first copy action);
+        # the source-rename PUT must never fire.
+        client, _, _ = _build_fake_client(
+            source_rows=[_DatasetRow(id="src-1", name="MyDataset")],
+            destination_rows=[],
+            items=[{"id": "item-a", "input": "hello"}],
+        )
+        client.rest_client.datasets.create_dataset.side_effect = RuntimeError(
+            "boom mid-copy"
+        )
+        result, audit_path = self._run(client, tmp_path)
+
+        assert result.exit_code == 1
+        # No update_dataset PUT touched the source id -> its name is intact.
+        source_touched = [
+            c
+            for c in client.rest_client.datasets.update_dataset.call_args_list
+            if c.kwargs.get("id") == "src-1"
+        ]
+        assert source_touched == []
+        on_disk = json.loads(audit_path.read_text())
+        assert on_disk["status"] == "failed"
+        # The handoff actions never reached ``ok``.
+        ok_types = {a["type"] for a in on_disk["actions"] if a.get("status") == "ok"}
+        assert "rename_source" not in ok_types
+        assert "promote_destination" not in ok_types
+
+    def test_rerun_after_failure__discards_stale_temp_then_completes(
+        self, tmp_path
+    ) -> None:
+        # AC: re-running after an interrupted run completes with no manual
+        # cleanup. A stale ``MyDataset__migrating`` from the prior failed run —
+        # carrying the migration marker tag that proves it's ours — is
+        # discovered and deleted before the destination is recreated.
+        client, _, _ = _build_fake_client(
+            source_rows=[_DatasetRow(id="src-1", name="MyDataset")],
+            destination_rows=[],
+            items=[{"id": "item-a", "input": "hello"}],
+            stale_temp_rows=[
+                _DatasetRow(
+                    id="stale-1",
+                    name="MyDataset__migrating",
+                    tags=[TEMP_MIGRATION_MARKER_TAG],
+                )
+            ],
+        )
+        result, audit_path = self._run(client, tmp_path)
+
+        assert result.exit_code == 0, result.output
+        # The stale temp was deleted by id before recreate.
+        client.rest_client.datasets.delete_dataset.assert_called_once_with(id="stale-1")
+        on_disk = json.loads(audit_path.read_text())
+        assert on_disk["status"] == "ok"
+        action_types = [
+            a["type"] for a in on_disk["actions"] if a.get("status") == "ok"
+        ]
+        assert action_types.index("discard_stale_temp") < action_types.index(
+            "create_destination"
+        )
+
+
+# ---------------------------------------------------------------------------
+# Planner unit tests (no Click invocation)
+# ---------------------------------------------------------------------------
+
+
+class TestPlanBuilding:
+    def test_build_dataset_plan__default_flow__orders_create_replay_cascades_then_handoff(
+        self,
+    ) -> None:
+        # OPIK-7162: the plan builds the destination under a temp name FIRST
+        # (source keeps its name), runs the copy + cascades, then does the
+        # name handoff LAST: rename source -> <name>_v1, promote temp ->
+        # <name>. The order is load-bearing on two axes:
+        #   * CascadeOptimizations before CascadeExperiments (opt-id remap).
+        #   * RenameSource before PromoteDestination (source-away then
+        #     destination-in, so <name> is never held by two rows at once).
+        # Three find_datasets pages: source resolve, _v1 collision check,
+        # __migrating stale-temp lookup.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset", description="d")]),
+                _Page([]),
+                _Page([]),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="B",
+        )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == [
+            "CreateDestination",
+            "ReplayVersions",
+            "CascadeOptimizations",
+            "CascadeExperiments",
+            "RenameSource",
+            "PromoteDestination",
+        ]
+        # Destination is created under the temp name, not the final name.
+        create = plan.actions[0]
+        assert create.name == "MyDataset__migrating"
+        replay = plan.actions[1]
+        assert replay.source_name == "MyDataset"
+        assert replay.dest_name == "MyDataset__migrating"
+        # Handoff: source away first, destination in second.
+        rename = plan.actions[4]
+        assert rename.from_name == "MyDataset"
+        assert rename.to_name == "MyDataset_v1"
+        promote = plan.actions[5]
+        assert promote.from_name == "MyDataset__migrating"
+        assert promote.to_name == "MyDataset"
+        assert plan.target_name == "MyDataset"
+        # New remap dict starts empty; _cascade_optimizations populates it.
+        assert plan.optimization_id_remap == {}
+
+    def test_build_dataset_plan__exclude_experiments__omits_both_cascades(
+        self,
+    ) -> None:
+        # OPIK-7161: --exclude-experiments drops the experiment stage AND
+        # the optimization stage (optimizations are containers for the
+        # skipped experiments). No cascade actions, but the name handoff
+        # (rename + promote) still runs after the dataset + versions copy.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset", description="d")]),
+                _Page([]),
+                _Page([]),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="B",
+            exclude_experiments=True,
+        )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == [
+            "CreateDestination",
+            "ReplayVersions",
+            "RenameSource",
+            "PromoteDestination",
+        ]
+        assert not any(
+            isinstance(a, planner_module.CascadeExperiments) for a in plan.actions
+        )
+        assert not any(
+            isinstance(a, planner_module.CascadeOptimizations) for a in plan.actions
+        )
+
+    def test_build_dataset_plan__exclude_experiments_default_false__keeps_cascades(
+        self,
+    ) -> None:
+        # Default (flag off) is unchanged: both cascades still emitted.
+        # Guards the opt-out default so a plain migrate never silently
+        # starts skipping experiments.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset")]),
+                _Page([]),
+                _Page([]),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="B",
+        )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == [
+            "CreateDestination",
+            "ReplayVersions",
+            "CascadeOptimizations",
+            "CascadeExperiments",
+            "RenameSource",
+            "PromoteDestination",
+        ]
+
+    def test_build_dataset_plan__test_suite__type_forwarded_to_destination(
+        self,
+    ) -> None:
+        # Test suites flow through the same plan shape as plain datasets;
+        # the only difference is ``CreateDestination.type`` being forwarded
+        # so the target accepts suite-level evaluators + execution_policy
+        # via ``ReplayVersions``.
+        rest_client = _planner_rest_client(
+            [
+                _Page(
+                    [_DatasetRow(id="src-1", name="MySuite", type="evaluation_suite")]
+                ),
+                _Page([]),
+                _Page([]),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MySuite",
+            to_project="B",
+        )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == [
+            "CreateDestination",
+            "ReplayVersions",
+            "CascadeOptimizations",
+            "CascadeExperiments",
+            "RenameSource",
+            "PromoteDestination",
+        ]
+        replay = plan.actions[1]
+        assert replay.is_test_suite is True
+
+    def test_build_dataset_plan__rename_target_collides_workspace_wide__raises_conflict(
+        self,
+    ) -> None:
+        # The eventual source-rename target "<source>_v1" collides with
+        # another dataset in the workspace — caught up-front so a doomed run
+        # never does any copy work.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset")]),
+                _Page([_DatasetRow(id="other-1", name="MyDataset_v1")]),
+                _Page([]),
+            ]
+        )
+
+        with pytest.raises(ConflictError) as exc_info:
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="B",
+            )
+        assert "MyDataset_v1" in str(exc_info.value)
+
+    def test_build_dataset_plan__rename_target_match_is_source_itself__no_conflict(
+        self,
+    ) -> None:
+        # When find_datasets returns the source itself for the _v1 check, we
+        # must not treat that as a collision — it's about to be renamed.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset")]),
+                _Page([_DatasetRow(id="src-1", name="MyDataset_v1")]),
+                _Page([]),
+            ]
+        )
+        # Should NOT raise: the only "match" is the source dataset itself.
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="B",
+        )
+        assert plan.target_name == "MyDataset"
+
+    def test_build_dataset_plan__marked_stale_temp__prepends_discard_action(
+        self,
+    ) -> None:
+        # OPIK-7162 safe re-run: a leftover "<name>__migrating" carrying the
+        # migration marker tag (proof it's ours) is detected and a
+        # DiscardStaleTemp action is prepended so the re-run starts clean.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset")]),
+                _Page([]),
+                _Page(
+                    [
+                        _DatasetRow(
+                            id="stale-1",
+                            name="MyDataset__migrating",
+                            tags=[TEMP_MIGRATION_MARKER_TAG],
+                        )
+                    ]
+                ),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="B",
+        )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == [
+            "DiscardStaleTemp",
+            "CreateDestination",
+            "ReplayVersions",
+            "CascadeOptimizations",
+            "CascadeExperiments",
+            "RenameSource",
+            "PromoteDestination",
+        ]
+        discard = plan.actions[0]
+        assert discard.temp_id == "stale-1"
+        assert discard.temp_name == "MyDataset__migrating"
+
+    def test_build_dataset_plan__unmarked_name_collision__raises_conflict(
+        self,
+    ) -> None:
+        # A dataset named "<name>__migrating" WITHOUT the migration marker is a
+        # real user dataset that merely shares the name — it must NOT be
+        # deleted. The planner aborts with ConflictError instead.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset")]),
+                _Page([]),
+                _Page([_DatasetRow(id="user-1", name="MyDataset__migrating")]),
+            ]
+        )
+
+        with pytest.raises(ConflictError) as exc_info:
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="B",
+            )
+        assert "MyDataset__migrating" in str(exc_info.value)
+        assert TEMP_MIGRATION_MARKER_TAG in str(exc_info.value)
+
+    def test_build_dataset_plan__same_from_and_to_project_flag__raises_conflict(
+        self,
+    ) -> None:
+        # Cheap early-out: user literally passed --from-project A --to-project A.
+        # Rejected before any lookup.
+        rest_client = _planner_rest_client(
+            [_Page([_DatasetRow(id="src-1", name="MyDataset")])]
+        )
+
+        with pytest.raises(ConflictError, match="same project"):
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="A",
+                from_project="A",
+            )
+
+    def test_build_dataset_plan__omitted_flag_source_in_dest_project__raises_conflict(
+        self,
+    ) -> None:
+        # The gap the flag-only check missed: --from-project is OMITTED, but the
+        # source actually lives in the destination project. resolve_source
+        # populates source.project_name from the row's project_id, so the
+        # authoritative post-resolve guard still catches it.
+        source_row = _DatasetRow(id="src-1", name="MyDataset", project_id="proj-A")
+        rest_client = _planner_rest_client([_Page([source_row])])
+        # project_name_for_row -> client.get_project(id="proj-A").name == "A".
+        proj = MagicMock()
+        proj.name = "A"
+        rest_client.projects.get_project_by_id.return_value = proj
+
+        with pytest.raises(ConflictError, match="same project"):
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="A",
+            )
+
+    def test_build_dataset_plan__workspace_scoped_source__no_same_project_abort(
+        self,
+    ) -> None:
+        # A workspace-scoped source (no project_id -> project_name is None) has
+        # no single project to collide with --to-project, so a workspace-scoped
+        # -> project migrate is legitimate and must NOT be blocked.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset", project_id=None)]),
+                _Page([]),
+                _Page([]),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="A",
+        )
+        assert plan.target_name == "MyDataset"
+
+    def test_build_dataset_plan__no_stale_temp__no_discard_action(self) -> None:
+        # The common case: no leftover temp, so no DiscardStaleTemp emitted.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_DatasetRow(id="src-1", name="MyDataset")]),
+                _Page([]),
+                _Page([]),
+            ]
+        )
+
+        plan = planner_module.build_dataset_plan(
+            client=_planner_client(rest_client),
+            name="MyDataset",
+            to_project="B",
+        )
+
+        assert not any(
+            isinstance(a, planner_module.DiscardStaleTemp) for a in plan.actions
+        )
+
+    def test_build_dataset_plan__source_name_not_found__raises_dataset_not_found(
+        self,
+    ) -> None:
+        rest_client = _planner_rest_client([_Page([])])
+
+        with pytest.raises(DatasetNotFoundError) as exc_info:
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="Missing",
+                to_project="B",
+            )
+        assert "Missing" in str(exc_info.value)
+
+    def test_build_dataset_plan__source_name_resolves_to_many__raises_conflict(
+        self,
+    ) -> None:
+        # Workspace uniqueness is enforced by the BE (UNIQUE
+        # (workspace_id, name)); if the BE invariant is somehow
+        # violated, surface it as ConflictError rather than silently
+        # picking a row.
+        rest_client = _planner_rest_client(
+            [
+                _Page(
+                    [
+                        _DatasetRow(id="a", name="MyDataset"),
+                        _DatasetRow(id="b", name="MyDataset"),
+                    ]
+                )
+            ]
+        )
+
+        with pytest.raises(ConflictError):
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="B",
+            )
+
+    def test_build_dataset_plan__destination_project_missing__raises_project_not_found(
+        self,
+    ) -> None:
+        rest_client = _planner_rest_client(
+            find_side_effects=[],
+            target_project_exists=False,
+        )
+
+        with pytest.raises(ProjectNotFoundError) as exc_info:
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="DoesNotExist",
+            )
+        assert "DoesNotExist" in str(exc_info.value)
+
+    def test_build_dataset_plan__destination_project_missing__suggests_similar_names(
+        self,
+    ) -> None:
+        rest_client = _planner_rest_client(
+            find_side_effects=[],
+            target_project_exists=False,
+            workspace_project_names=["production", "staging", "Beat", "Best"],
+        )
+
+        with pytest.raises(ProjectNotFoundError) as exc_info:
+            planner_module.build_dataset_plan(
+                client=_planner_client(rest_client),
+                name="MyDataset",
+                to_project="Beta",
+            )
+        message = str(exc_info.value)
+        assert "Beta" in message
+        # difflib should surface the close one-letter neighbours.
+        assert "Did you mean" in message
+        assert "Beat" in message or "Best" in message
+
+
+# ---------------------------------------------------------------------------
+# OPIK-7162 + OPIK-7168 integration: resume reuses the temp destination and
+# finishes the pending handoff (rename source -> _v1, promote temp -> original).
+# ---------------------------------------------------------------------------
+
+
+class TestBuildResumePlanTempDest:
+    def _checkpoint(self) -> MigrationCheckpoint:
+        return MigrationCheckpoint(
+            key="k",
+            workspace="ws",
+            project="B",
+            dataset="MyDataset",
+            path=Path("/tmp/does-not-matter.json"),
+            dataset_phase_done=True,
+            source_dataset_id="src-1",
+            source_name="MyDataset",
+            temp_dest_name="MyDataset__migrating",
+        )
+
+    def _resume_client(self, source_id: str = "src-1") -> MagicMock:
+        # resolve_source(MyDataset) -> the still-unrenamed source;
+        # get_dataset(MyDataset__migrating) -> the temp destination.
+        rest_client = _planner_rest_client(
+            [_Page([_DatasetRow(id=source_id, name="MyDataset")])]
+        )
+        client = _planner_client(rest_client)
+        dest = MagicMock()
+        dest.id = "temp-dest-1"
+        client.get_dataset = MagicMock(return_value=dest)
+        return client
+
+    def test_resume__reuses_temp_and_appends_cascade_then_handoff(self) -> None:
+        # A dataset_phase_done checkpoint means create-temp/replay/optimizations
+        # already ran into MyDataset__migrating and the source still holds its
+        # original name. The resume plan must NOT re-create or re-replay; it
+        # resolves the temp destination, then emits the pending tail:
+        # CascadeExperiments -> RenameSource -> PromoteDestination.
+        with patch.object(
+            planner_module, "reconstruct_remaps", return_value=ReconstructedRemaps()
+        ):
+            plan = planner_module.build_dataset_plan(
+                client=self._resume_client(),
+                name="MyDataset",
+                to_project="B",
+                resume_checkpoint=self._checkpoint(),
+            )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == ["CascadeExperiments", "RenameSource", "PromoteDestination"]
+        assert plan.is_resume is True
+        # The cascade + promote target the TEMP destination (promote hasn't run
+        # yet); the source rename moves the original name to _v1.
+        cascade = plan.actions[0]
+        assert cascade.dest_name == "MyDataset__migrating"
+        rename = plan.actions[1]
+        assert rename.from_name == "MyDataset"
+        assert rename.to_name == "MyDataset_v1"
+        promote = plan.actions[2]
+        assert promote.from_name == "MyDataset__migrating"
+        assert promote.to_name == "MyDataset"
+
+    def test_resume__source_id_mismatch__raises(self) -> None:
+        # If the user-supplied name now resolves to a DIFFERENT dataset than the
+        # interrupted run's source, resume must refuse rather than migrate the
+        # wrong dataset.
+        with patch.object(
+            planner_module, "reconstruct_remaps", return_value=ReconstructedRemaps()
+        ):
+            with pytest.raises(MigrationError, match="different dataset"):
+                planner_module.build_dataset_plan(
+                    client=self._resume_client(source_id="DIFFERENT-id"),
+                    name="MyDataset",
+                    to_project="B",
+                    resume_checkpoint=self._checkpoint(),
+                )

sdks/python/tests/unit/cli/test_migrate_dataset_resume_reconstruct.py

@@ -0,0 +1,241 @@
+"""Unit tests for read-only remap reconstruction on resume (OPIK-7168).
+
+``reconstruct_remaps`` rebuilds version_remap / item_id_remap /
+optimization_id_remap from an already-migrated destination without writing
+anything. These tests mock the REST reads (versions, per-version items,
+optimizations) and assert the pairing/identity/fallback logic and the
+fail-loud guard on a version-count mismatch.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+from unittest.mock import MagicMock
+
+import pytest
+
+from opik.cli.migrate.datasets.resume import (
+    ResumeReconstructionError,
+    reconstruct_remaps,
+)
+
+
+class _Version:
+    def __init__(self, id: str, version_hash: str) -> None:
+        self.id = id
+        self.version_hash = version_hash
+
+
+class _Item:
+    def __init__(
+        self,
+        id: str,
+        data: Optional[Dict[str, Any]] = None,
+        description: Optional[str] = None,
+        tags: Optional[List[str]] = None,
+        evaluators: Optional[List[Any]] = None,
+        execution_policy: Optional[Any] = None,
+        trace_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        source: Optional[str] = None,
+    ) -> None:
+        self.id = id
+        self.data = data or {"q": id}
+        self.description = description
+        self.tags = tags
+        self.evaluators = evaluators
+        self.execution_policy = execution_policy
+        self.trace_id = trace_id
+        self.span_id = span_id
+        self.source = source
+
+
+class _Optimization:
+    def __init__(self, id: str, name: str) -> None:
+        self.id = id
+        self.name = name
+
+
+def _page(content: List[Any]) -> MagicMock:
+    return MagicMock(content=content)
+
+
+def _rest_client(
+    *,
+    source_versions: List[_Version],
+    dest_versions: List[_Version],
+    items_by_version_hash: Dict[str, List[_Item]],
+    source_optimizations: Optional[List[_Optimization]] = None,
+    dest_optimizations: Optional[List[_Optimization]] = None,
+) -> MagicMock:
+    rest = MagicMock()
+
+    def _list_versions(id: str, page: int, size: int) -> MagicMock:
+        # reconstruct calls _iter_source_versions_chronological which lists
+        # NEWEST-first then reverses; return newest-first here.
+        if page != 1:
+            return _page([])
+        if id == "src-ds":
+            return _page(list(reversed(source_versions)))
+        return _page(list(reversed(dest_versions)))
+
+    rest.datasets.list_dataset_versions.side_effect = _list_versions
+
+    def _stream(dataset_name: str, steam_limit: int, dataset_version: str, **kw: Any):
+        # _stream_version_items_raw parses a byte stream via rest_stream_parser;
+        # tests patch that helper directly (see the patch_stream fixture), so
+        # this raw endpoint should never be hit.
+        raise AssertionError("stream_dataset_items should be patched in tests")
+
+    rest.datasets.stream_dataset_items.side_effect = _stream
+
+    def _find_opts(dataset_id: str, page: int, size: int) -> MagicMock:
+        if page != 1:
+            return _page([])
+        if dataset_id == "src-ds":
+            return _page(source_optimizations or [])
+        return _page(dest_optimizations or [])
+
+    rest.optimizations.find_optimizations.side_effect = _find_opts
+    return rest
+
+
+@pytest.fixture
+def patch_stream(monkeypatch: pytest.MonkeyPatch):
+    """Patch ``_stream_version_items_raw`` (byte-stream parser) with a direct
+    lookup keyed by version_hash, so tests supply plain ``_Item`` lists."""
+    registry: Dict[str, List[_Item]] = {}
+
+    def _fake_stream(
+        rest_client: Any,
+        *,
+        dataset_name: str,
+        project_name: Optional[str],
+        version_hash: Optional[str],
+    ) -> List[_Item]:
+        return registry.get(version_hash or "", [])
+
+    monkeypatch.setattr(
+        "opik.cli.migrate.datasets.resume._stream_version_items_raw", _fake_stream
+    )
+    return registry
+
+
+def test_reconstruct__paired_versions_and_identity_items__remaps_built(
+    patch_stream: Dict[str, List[_Item]],
+) -> None:
+    source_versions = [_Version("sv1", "h1"), _Version("sv2", "h2")]
+    # Destination hashes DIFFER from source (BE recomputes them) — pairing must
+    # be positional, not by hash.
+    dest_versions = [_Version("dv1", "dh1"), _Version("dv2", "dh2")]
+    # Identity: dest item ids equal source item ids (copy_from preserves them).
+    patch_stream["h1"] = [_Item("i1"), _Item("i2")]
+    patch_stream["dh1"] = [_Item("i1"), _Item("i2")]
+    patch_stream["h2"] = [_Item("i3")]
+    patch_stream["dh2"] = [_Item("i3")]
+
+    rest = _rest_client(
+        source_versions=source_versions,
+        dest_versions=dest_versions,
+        items_by_version_hash={},
+    )
+
+    remaps = reconstruct_remaps(
+        rest,
+        source_dataset_id="src-ds",
+        source_name="ds",
+        source_project_name=None,
+        dest_dataset_id="dst-ds",
+        dest_name="ds",
+        dest_project_name="proj",
+    )
+
+    assert remaps.version_remap == {"sv1": "dv1", "sv2": "dv2"}
+    assert remaps.item_id_remap == {"i1": "i1", "i2": "i2", "i3": "i3"}
+
+
+def test_reconstruct__dest_ids_differ__content_hash_fallback_maps_them(
+    patch_stream: Dict[str, List[_Item]],
+) -> None:
+    # First-version read-back case: dest item id was minted fresh, so identity
+    # match fails and we fall back to content-hash matching.
+    source_versions = [_Version("sv1", "h1")]
+    dest_versions = [_Version("dv1", "dh1")]
+    patch_stream["h1"] = [_Item("src-i1", data={"q": "same"})]
+    patch_stream["dh1"] = [_Item("dst-i1", data={"q": "same"})]
+
+    rest = _rest_client(
+        source_versions=source_versions,
+        dest_versions=dest_versions,
+        items_by_version_hash={},
+    )
+
+    remaps = reconstruct_remaps(
+        rest,
+        source_dataset_id="src-ds",
+        source_name="ds",
+        source_project_name=None,
+        dest_dataset_id="dst-ds",
+        dest_name="ds",
+        dest_project_name="proj",
+    )
+
+    assert remaps.item_id_remap == {"src-i1": "dst-i1"}
+
+
+def test_reconstruct__version_count_mismatch__raises_reconstruction_error(
+    patch_stream: Dict[str, List[_Item]],
+) -> None:
+    # Destination has fewer versions than source (e.g. crash mid-replay) — must
+    # fail loud rather than mis-pair.
+    rest = _rest_client(
+        source_versions=[_Version("sv1", "h1"), _Version("sv2", "h2")],
+        dest_versions=[_Version("dv1", "dh1")],
+        items_by_version_hash={},
+    )
+
+    with pytest.raises(ResumeReconstructionError, match="Delete the destination"):
+        reconstruct_remaps(
+            rest,
+            source_dataset_id="src-ds",
+            source_name="ds",
+            source_project_name=None,
+            dest_dataset_id="dst-ds",
+            dest_name="ds",
+            dest_project_name="proj",
+        )
+
+
+def test_reconstruct__optimizations__matched_by_name_into_remap(
+    patch_stream: Dict[str, List[_Item]],
+) -> None:
+    source_versions = [_Version("sv1", "h1")]
+    dest_versions = [_Version("dv1", "dh1")]
+    patch_stream["h1"] = []
+    patch_stream["dh1"] = []
+
+    rest = _rest_client(
+        source_versions=source_versions,
+        dest_versions=dest_versions,
+        items_by_version_hash={},
+        source_optimizations=[
+            _Optimization("so1", "opt-a"),
+            _Optimization("so2", "opt-b"),
+        ],
+        dest_optimizations=[
+            _Optimization("do1", "opt-a"),
+            _Optimization("do2", "opt-b"),
+        ],
+    )
+
+    remaps = reconstruct_remaps(
+        rest,
+        source_dataset_id="src-ds",
+        source_name="ds",
+        source_project_name=None,
+        dest_dataset_id="dst-ds",
+        dest_name="ds",
+        dest_project_name="proj",
+    )
+
+    assert remaps.optimization_id_remap == {"so1": "do1", "so2": "do2"}

sdks/python/tests/unit/cli/test_migrate_prompt_executor.py

@@ -0,0 +1,305 @@
+"""Executor tests for ``opik migrate prompt``.
+
+Drives ``execute_plan`` against fully-mocked rest_client surfaces and
+asserts that each action dispatches to the right Fern call with the
+right kwargs. The verbatim-commit replay path and per-version audit
+shape are exercised here (the dataset-side analogue lives in
+``test_migrate_dataset_version_replay.py``).
+"""
+
+from __future__ import annotations
+
+from typing import Any, List
+from unittest.mock import MagicMock
+
+from opik.cli.migrate.audit import AuditLog
+from opik.cli.migrate.prompts import executor as executor_module
+from opik.cli.migrate.prompts.planner import (
+    CreateDestination,
+    MigrationPlan,
+    RenameSource,
+    ReplayVersions,
+)
+from opik.cli.migrate.prompts.resolver import ResolvedPrompt
+from opik.rest_api.types.prompt_version_detail import PromptVersionDetail
+
+from ._migrate_prompt_helpers import _Page, _PromptVersionRow
+
+
+def _client_with_rest_mock(version_pages: List[_Page]) -> Any:
+    rest_client = MagicMock()
+    rest_client.prompts.update_prompt = MagicMock()
+    rest_client.prompts.create_prompt = MagicMock()
+    rest_client.prompts.create_prompt_version = MagicMock(
+        side_effect=lambda **kw: MagicMock(
+            id=f"dest-version-{kw['version'].commit}",
+            commit=kw["version"].commit,
+            environments=kw["version"].environments,
+        )
+    )
+    rest_client.prompts.get_prompt_versions.side_effect = version_pages
+
+    client = MagicMock()
+    client.rest_client = rest_client
+    return client, rest_client
+
+
+def _make_plan() -> MigrationPlan:
+    source = ResolvedPrompt(
+        id="src-1",
+        name="MyPrompt",
+        project_id=None,
+        project_name=None,
+        description="orig",
+        tags=["a"],
+        template_structure="text",
+    )
+    plan = MigrationPlan(source=source, target_name="MyPrompt", to_project="B")
+    plan.actions = [
+        RenameSource(
+            source_id="src-1",
+            from_name="MyPrompt",
+            to_name="MyPrompt_v1",
+            description="orig",
+            tags=["a"],
+        ),
+        CreateDestination(
+            name="MyPrompt",
+            project_name="B",
+            description="orig",
+            tags=["a"],
+            template_structure="text",
+        ),
+        ReplayVersions(
+            source_prompt_id="src-1",
+            source_name_after_rename="MyPrompt_v1",
+            source_project_name=None,
+            dest_name="MyPrompt",
+            dest_project_name="B",
+            template_structure="text",
+        ),
+    ]
+    return plan
+
+
+class TestExecuteActions:
+    def test_rename_source__re_passes_description_and_tags(self) -> None:
+        # The BE rename PUT wipes description (no COALESCE in the UPDATE);
+        # the executor must re-pass it so the source description survives.
+        client, rest_client = _client_with_rest_mock(
+            [_Page([])],  # no versions
+        )
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        rest_client.prompts.update_prompt.assert_called_once_with(
+            id="src-1",
+            name="MyPrompt_v1",
+            description="orig",
+            tags=["a"],
+        )
+
+    def test_create_destination__omits_template_so_be_does_not_auto_mint_v1(
+        self,
+    ) -> None:
+        # If we passed a template here, the BE would auto-create a v1
+        # with its own commit hash, defeating the verbatim-replay
+        # guarantee. The executor must NOT include template in the
+        # create_prompt call.
+        client, rest_client = _client_with_rest_mock([_Page([])])
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        rest_client.prompts.create_prompt.assert_called_once()
+        kwargs = rest_client.prompts.create_prompt.call_args.kwargs
+        assert kwargs["name"] == "MyPrompt"
+        assert kwargs["project_name"] == "B"
+        assert kwargs["description"] == "orig"
+        assert kwargs["tags"] == ["a"]
+        assert kwargs["template_structure"] == "text"
+        # Critical: no template, no type, no metadata at the container
+        # level — those are version-level fields handled by the replay.
+        assert "template" not in kwargs
+        assert "type" not in kwargs
+        assert "metadata" not in kwargs
+
+    def test_replay_versions__preserves_commit_hashes_verbatim(self) -> None:
+        # The verbatim-commit invariant is the architectural reason this
+        # slice exists: source versions are minted on the destination
+        # with the source's exact commit hash, made possible by the
+        # destination prompt's fresh id (so (workspace_id, prompt_id,
+        # commit) never collides).
+        v1 = _PromptVersionRow(
+            id="src-v-1",
+            prompt_id="src-1",
+            commit="aaaaaaaa",
+            template="hello {{name}}",
+            type="mustache",
+        )
+        v2 = _PromptVersionRow(
+            id="src-v-2",
+            prompt_id="src-1",
+            commit="bbbbbbbb",
+            template="hello {{name}} v2",
+            type="mustache",
+        )
+        # BE returns newest-first (pv.id DESC); the replay loop reverses
+        # to oldest-first. Pass them in newest-first order to mirror the
+        # BE response shape.
+        client, rest_client = _client_with_rest_mock([_Page([v2, v1])])
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        # Two POST /v1/private/prompts/versions calls, oldest-first.
+        assert rest_client.prompts.create_prompt_version.call_count == 2
+        first_call = rest_client.prompts.create_prompt_version.call_args_list[0]
+        second_call = rest_client.prompts.create_prompt_version.call_args_list[1]
+
+        first_version = first_call.kwargs["version"]
+        second_version = second_call.kwargs["version"]
+        assert isinstance(first_version, PromptVersionDetail)
+        assert isinstance(second_version, PromptVersionDetail)
+        assert first_version.commit == "aaaaaaaa"
+        assert second_version.commit == "bbbbbbbb"
+        assert first_version.template == "hello {{name}}"
+        assert second_version.template == "hello {{name}} v2"
+
+    def test_replay_versions__carries_environments_verbatim(self) -> None:
+        # Environment ownership is per-version data the BE accepts inline on
+        # create_prompt_version. The history mixes an env-less version, a
+        # non-latest version owning a single env, and a latest version owning
+        # multiple envs — proving single + multi env work and that a
+        # non-latest version's env is preserved (not just the newest one's).
+        v1 = _PromptVersionRow(
+            id="src-v-1",
+            prompt_id="src-1",
+            commit="aaaaaaaa",
+            template="hello",
+        )
+        v2 = _PromptVersionRow(
+            id="src-v-2",
+            prompt_id="src-1",
+            commit="bbbbbbbb",
+            template="hi",
+            environments=["development"],
+        )
+        v3 = _PromptVersionRow(
+            id="src-v-3",
+            prompt_id="src-1",
+            commit="cccccccc",
+            template="hey",
+            environments=["production", "staging"],
+        )
+        # BE returns newest-first; the loop reverses to oldest-first.
+        client, rest_client = _client_with_rest_mock([_Page([v3, v2, v1])])
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        calls = rest_client.prompts.create_prompt_version.call_args_list
+        assert calls[0].kwargs["version"].environments is None
+        assert calls[1].kwargs["version"].environments == ["development"]
+        assert calls[2].kwargs["version"].environments == ["production", "staging"]
+
+    def test_replay_versions__records_environments_in_audit(self) -> None:
+        v1 = _PromptVersionRow(
+            id="src-v-1",
+            prompt_id="src-1",
+            commit="aaaaaaaa",
+            template="hello",
+            environments=["staging", "production"],
+        )
+        client, _ = _client_with_rest_mock([_Page([v1])])
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        record = next(a for a in audit.actions if a["type"] == "replay_prompt_version")
+        # Environments are recorded verbatim (order is irrelevant), so compare
+        # order-insensitively.
+        assert sorted(record["source_environments"]) == ["production", "staging"]
+        assert sorted(record["target_environments"]) == ["production", "staging"]
+
+    def test_replay_versions__populates_prompt_version_id_remap(self) -> None:
+        # Slice 7 (OPIK-6575) reads ``plan.prompt_version_id_remap`` to
+        # remap experiment FK references. The executor must populate it.
+        v1 = _PromptVersionRow(
+            id="src-v-1",
+            prompt_id="src-1",
+            commit="aaaaaaaa",
+            template="hello",
+        )
+        v2 = _PromptVersionRow(
+            id="src-v-2",
+            prompt_id="src-1",
+            commit="bbbbbbbb",
+            template="hi",
+        )
+        client, rest_client = _client_with_rest_mock([_Page([v2, v1])])
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        assert plan.prompt_version_id_remap == {
+            "src-v-1": "dest-version-aaaaaaaa",
+            "src-v-2": "dest-version-bbbbbbbb",
+        }
+
+    def test_replay_versions__records_one_audit_entry_per_version(self) -> None:
+        v1 = _PromptVersionRow(
+            id="src-v-1",
+            prompt_id="src-1",
+            commit="aaaaaaaa",
+            template="hello",
+        )
+        v2 = _PromptVersionRow(
+            id="src-v-2",
+            prompt_id="src-1",
+            commit="bbbbbbbb",
+            template="hi",
+        )
+        client, _ = _client_with_rest_mock([_Page([v2, v1])])
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.execute_plan(client, plan, audit)
+
+        per_version_records = [
+            a for a in audit.actions if a["type"] == "replay_prompt_version"
+        ]
+        assert len(per_version_records) == 2
+        # All per-version records are recorded as ``ok`` (the per-version
+        # audit happens after a successful create_prompt_version).
+        assert all(a["status"] == "ok" for a in per_version_records)
+        assert {a["source_commit"] for a in per_version_records} == {
+            "aaaaaaaa",
+            "bbbbbbbb",
+        }
+
+
+class TestRecordPlanned:
+    def test_record_planned__emits_one_planned_entry_per_action(self) -> None:
+        # Dry-run path: every action gets a ``planned`` record, no REST
+        # calls fire.
+        plan = _make_plan()
+        audit = AuditLog(command="opik migrate prompt", args={})
+
+        executor_module.record_planned(plan, audit)
+
+        statuses = [a["status"] for a in audit.actions]
+        types = [a["type"] for a in audit.actions]
+        assert statuses == ["planned", "planned", "planned"]
+        assert types == [
+            "rename_source",
+            "create_destination",
+            "replay_versions",
+        ]

sdks/python/tests/unit/cli/test_migrate_prompt_planner.py

@@ -0,0 +1,244 @@
+"""Planner tests for ``opik migrate prompt``.
+
+The dataset planner tests are next door; this file tracks the same
+patterns (happy-path action ordering, conflict detection, project-not-
+found, source-not-found) for the prompt slice.
+"""
+
+from __future__ import annotations
+
+import pytest
+from click.testing import CliRunner
+
+from opik.cli import cli
+from opik.cli.migrate.errors import (
+    ConflictError,
+    ProjectNotFoundError,
+    PromptNotFoundError,
+)
+from opik.cli.migrate.prompts import planner as prompt_planner
+
+from ._migrate_prompt_helpers import (
+    _Page,
+    _PromptRow,
+    _planner_client,
+    _planner_rest_client,
+)
+
+
+class TestMigratePromptHelp:
+    def test_migrate_group_help__lists_prompt_subcommand(self) -> None:
+        runner = CliRunner()
+        result = runner.invoke(cli, ["migrate", "--help"])
+        assert result.exit_code == 0
+        assert "prompt" in result.output
+
+    def test_migrate_prompt_help__lists_required_flags(self) -> None:
+        runner = CliRunner()
+        result = runner.invoke(cli, ["migrate", "prompt", "--help"])
+        assert result.exit_code == 0
+        assert "--to-project" in result.output
+        assert "--from-project" in result.output
+        assert "--dry-run" in result.output
+
+
+class TestPlanBuilding:
+    def test_build_prompt_plan__default_flow__orders_rename_create_replay(
+        self,
+    ) -> None:
+        # Action ordering invariant: the rename frees the workspace-unique
+        # name BEFORE the destination claims it, and the replay loop runs
+        # AFTER the destination container exists.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_PromptRow(id="src-1", name="MyPrompt", description="d")]),
+                _Page([]),  # rename-collision preflight finds nothing
+            ]
+        )
+
+        plan = prompt_planner.build_prompt_plan(
+            client=_planner_client(rest_client),
+            name="MyPrompt",
+            to_project="B",
+        )
+
+        types = [type(a).__name__ for a in plan.actions]
+        assert types == ["RenameSource", "CreateDestination", "ReplayVersions"]
+
+    def test_build_prompt_plan__rename_re_passes_description_and_tags(
+        self,
+    ) -> None:
+        # The BE rename PUT wipes description (UPDATE has no COALESCE on
+        # description); the planner forwards the source's description and
+        # tags into the RenameSource record so the executor re-passes them.
+        rest_client = _planner_rest_client(
+            [
+                _Page(
+                    [
+                        _PromptRow(
+                            id="src-1",
+                            name="MyPrompt",
+                            description="orig description",
+                            tags=["a", "b"],
+                        )
+                    ]
+                ),
+                _Page([]),
+            ]
+        )
+
+        plan = prompt_planner.build_prompt_plan(
+            client=_planner_client(rest_client),
+            name="MyPrompt",
+            to_project="B",
+        )
+
+        rename = plan.actions[0]
+        assert isinstance(rename, prompt_planner.RenameSource)
+        assert rename.from_name == "MyPrompt"
+        assert rename.to_name == "MyPrompt_v1"
+        assert rename.description == "orig description"
+        assert rename.tags == ["a", "b"]
+
+    def test_build_prompt_plan__create_destination_carries_metadata(
+        self,
+    ) -> None:
+        # CreateDestination forwards container-level metadata so the
+        # destination inherits description / tags / template_structure
+        # from the source.
+        rest_client = _planner_rest_client(
+            [
+                _Page(
+                    [
+                        _PromptRow(
+                            id="src-1",
+                            name="MyPrompt",
+                            description="orig",
+                            tags=["a"],
+                            template_structure="chat",
+                        )
+                    ]
+                ),
+                _Page([]),
+            ]
+        )
+
+        plan = prompt_planner.build_prompt_plan(
+            client=_planner_client(rest_client),
+            name="MyPrompt",
+            to_project="B",
+        )
+
+        create = plan.actions[1]
+        assert isinstance(create, prompt_planner.CreateDestination)
+        assert create.name == "MyPrompt"
+        assert create.project_name == "B"
+        assert create.description == "orig"
+        assert create.tags == ["a"]
+        assert create.template_structure == "chat"
+
+    def test_build_prompt_plan__replay_versions_carries_source_id_and_template_structure(
+        self,
+    ) -> None:
+        rest_client = _planner_rest_client(
+            [
+                _Page(
+                    [
+                        _PromptRow(
+                            id="src-1",
+                            name="MyPrompt",
+                            template_structure="chat",
+                        )
+                    ]
+                ),
+                _Page([]),
+            ]
+        )
+
+        plan = prompt_planner.build_prompt_plan(
+            client=_planner_client(rest_client),
+            name="MyPrompt",
+            to_project="B",
+        )
+
+        replay = plan.actions[2]
+        assert isinstance(replay, prompt_planner.ReplayVersions)
+        assert replay.source_prompt_id == "src-1"
+        assert replay.source_name_after_rename == "MyPrompt_v1"
+        assert replay.dest_name == "MyPrompt"
+        assert replay.dest_project_name == "B"
+        assert replay.template_structure == "chat"
+
+
+class TestPreflightErrors:
+    def test_build_prompt_plan__destination_project_missing__raises(
+        self,
+    ) -> None:
+        rest_client = _planner_rest_client(
+            [],  # no prompt lookups happen if project preflight fails
+            target_project_exists=False,
+        )
+
+        with pytest.raises(ProjectNotFoundError) as excinfo:
+            prompt_planner.build_prompt_plan(
+                client=_planner_client(rest_client),
+                name="MyPrompt",
+                to_project="DoesNotExist",
+            )
+
+        assert "DoesNotExist" in str(excinfo.value)
+
+    def test_build_prompt_plan__source_prompt_missing__raises(self) -> None:
+        rest_client = _planner_rest_client(
+            [_Page([])]  # source lookup returns no matches
+        )
+
+        with pytest.raises(PromptNotFoundError) as excinfo:
+            prompt_planner.build_prompt_plan(
+                client=_planner_client(rest_client),
+                name="Missing",
+                to_project="B",
+            )
+
+        assert "Missing" in str(excinfo.value)
+        assert "workspace" in str(excinfo.value).lower()
+
+    def test_build_prompt_plan__from_project_set_and_missing__raises_with_project_scope(
+        self,
+    ) -> None:
+        # When --from-project is passed and the source prompt is not in that
+        # project, the not-found error should name the project (not "the
+        # workspace") so the user knows where we looked.
+        rest_client = _planner_rest_client(
+            [_Page([])]  # source lookup returns no matches in the project
+        )
+
+        with pytest.raises(PromptNotFoundError) as excinfo:
+            prompt_planner.build_prompt_plan(
+                client=_planner_client(rest_client),
+                name="MyPrompt",
+                to_project="B",
+                from_project="A",
+            )
+
+        assert "MyPrompt" in str(excinfo.value)
+        assert "project 'A'" in str(excinfo.value)
+
+    def test_build_prompt_plan__rename_collision__raises(self) -> None:
+        # Source resolves, but the rename target ``MyPrompt_v1`` is already
+        # taken by another prompt -> ConflictError.
+        rest_client = _planner_rest_client(
+            [
+                _Page([_PromptRow(id="src-1", name="MyPrompt")]),
+                _Page([_PromptRow(id="other-1", name="MyPrompt_v1")]),
+            ]
+        )
+
+        with pytest.raises(ConflictError) as excinfo:
+            prompt_planner.build_prompt_plan(
+                client=_planner_client(rest_client),
+                name="MyPrompt",
+                to_project="B",
+            )
+
+        assert "MyPrompt_v1" in str(excinfo.value)

sdks/python/tests/unit/cli/test_migration_manifest.py

@@ -0,0 +1,358 @@
+"""Tests for the migration manifest used by opik import for resumable migrations."""
+
+import sqlite3
+from pathlib import Path
+
+import pytest
+
+from opik.cli.migration_manifest import MigrationManifest, MANIFEST_FILENAME
+
+
+@pytest.fixture
+def tmp_base(tmp_path: Path) -> Path:
+    return tmp_path
+
+
+class TestMigrationManifestLifecycle:
+    def test_manifest__fresh_instance__status_is_not_started(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        assert manifest.status == "not_started"
+        assert not manifest.is_in_progress
+        assert not manifest.is_completed
+
+    def test_start__fresh_manifest__status_becomes_in_progress_and_file_written(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.start()
+        assert manifest.is_in_progress
+        assert (tmp_base / MANIFEST_FILENAME).exists()
+
+    def test_complete__after_start__status_becomes_completed(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.start()
+        manifest.complete()
+        assert manifest.is_completed
+        assert not manifest.is_in_progress
+
+    def test_reset__after_start_with_data__all_state_cleared(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.start()
+        trace_file = tmp_base / "projects" / "p1" / "trace_abc.json"
+        trace_file.parent.mkdir(parents=True)
+        trace_file.touch()
+        manifest.mark_file_completed(trace_file)
+        manifest.add_trace_mapping("src-1", "dest-1")
+
+        manifest.reset()
+
+        assert manifest.status == "not_started"
+        assert manifest.completed_count() == 0
+        assert manifest.get_trace_id_map() == {}
+
+    def test_exists__before_any_save__returns_false(self, tmp_base: Path) -> None:
+        assert not MigrationManifest.exists(tmp_base)
+
+    def test_exists__after_start__returns_true(self, tmp_base: Path) -> None:
+        MigrationManifest(tmp_base).start()
+        assert MigrationManifest.exists(tmp_base)
+
+    def test_start__called_twice__started_at_preserved(self, tmp_base: Path) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.start()
+
+        conn = sqlite3.connect(str(tmp_base / MANIFEST_FILENAME))
+        first_started_at = conn.execute(
+            "SELECT value FROM status WHERE key = 'started_at'"
+        ).fetchone()[0]
+        conn.close()
+
+        manifest.start()  # second call must not overwrite started_at
+
+        conn = sqlite3.connect(str(tmp_base / MANIFEST_FILENAME))
+        second_started_at = conn.execute(
+            "SELECT value FROM status WHERE key = 'started_at'"
+        ).fetchone()[0]
+        conn.close()
+
+        assert second_started_at == first_started_at
+
+
+class TestFileTracking:
+    def _make_file(self, base: Path, relative: str) -> Path:
+        p = base / relative
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.touch()
+        return p
+
+    def test_is_file_completed__fresh_manifest__returns_false(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        f = self._make_file(tmp_base, "datasets/dataset_foo.json")
+        assert not manifest.is_file_completed(f)
+
+    def test_mark_file_completed__new_file__recorded_and_count_incremented(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        f = self._make_file(tmp_base, "datasets/dataset_foo.json")
+        manifest.mark_file_completed(f)
+        assert manifest.is_file_completed(f)
+        assert manifest.completed_count() == 1
+
+    def test_mark_file_completed__same_file_twice__count_remains_one(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        f = self._make_file(tmp_base, "datasets/dataset_foo.json")
+        manifest.mark_file_completed(f)
+        manifest.mark_file_completed(f)
+        assert manifest.completed_count() == 1
+
+    def test_mark_file_failed__new_file__recorded_with_error(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        f = self._make_file(tmp_base, "projects/p1/trace_abc.json")
+        manifest.mark_file_failed(f, "timeout")
+        assert manifest.failed_count() == 1
+        assert manifest.get_failed_files()[manifest.relative_path(f)] == "timeout"
+
+    def test_mark_file_completed__previously_failed_file__removed_from_failed(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        f = self._make_file(tmp_base, "projects/p1/trace_abc.json")
+        manifest.mark_file_failed(f, "timeout")
+        manifest.mark_file_completed(f)
+        assert manifest.failed_count() == 0
+        assert manifest.is_file_completed(f)
+
+    def test_relative_path__nested_file__returns_path_relative_to_base(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        f = self._make_file(tmp_base, "projects/my-project/trace_xyz.json")
+        assert manifest.relative_path(f) == str(
+            Path("projects") / "my-project" / "trace_xyz.json"
+        )
+
+
+class TestTraceIdMapping:
+    def test_add_trace_mapping__single_entry__retrievable_via_get_map(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.add_trace_mapping("src-111", "dest-222")
+        manifest.save()
+        assert manifest.get_trace_id_map() == {"src-111": "dest-222"}
+
+    def test_trace_id_map__save_then_new_instance__mapping_survives(
+        self, tmp_base: Path
+    ) -> None:
+        m1 = MigrationManifest(tmp_base)
+        m1.start()
+        m1.add_trace_mapping("src-aaa", "dest-bbb")
+        m1.save()
+
+        m2 = MigrationManifest(tmp_base)
+        assert m2.get_trace_id_map() == {"src-aaa": "dest-bbb"}
+
+    def test_get_trace_id_map__mutate_returned_dict__original_unaffected(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.add_trace_mapping("src-1", "dest-1")
+        manifest.save()
+        copy = manifest.get_trace_id_map()
+        copy["src-2"] = "dest-2"  # mutate the copy
+        assert "src-2" not in manifest.get_trace_id_map()
+
+    def test_add_trace_mapping__same_src_id_twice__last_dest_wins(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.add_trace_mapping("src-1", "dest-old")
+        manifest.add_trace_mapping("src-1", "dest-new")
+        manifest.save()
+        assert manifest.get_trace_id_map()["src-1"] == "dest-new"
+
+
+class TestDatabaseIntegrity:
+    def test_manifest_file__after_start_and_save__is_valid_sqlite_with_expected_tables(
+        self, tmp_base: Path
+    ) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.start()
+        manifest.add_trace_mapping("a", "b")
+        manifest.save()
+
+        conn = sqlite3.connect(str(tmp_base / MANIFEST_FILENAME))
+        tables = {
+            row[0]
+            for row in conn.execute(
+                "SELECT name FROM sqlite_master WHERE type='table'"
+            ).fetchall()
+        }
+        conn.close()
+
+        assert {"status", "completed_files", "failed_files", "trace_id_map"} <= tables
+
+    def test_save__after_start__no_tmp_file_created(self, tmp_base: Path) -> None:
+        manifest = MigrationManifest(tmp_base)
+        manifest.start()
+        assert not (tmp_base / "migration_manifest.tmp").exists()
+
+    def test_mark_file_completed__duplicate_write__count_remains_one(
+        self, tmp_base: Path
+    ) -> None:
+        """INSERT OR IGNORE means re-flushing the same path never duplicates rows."""
+        manifest = MigrationManifest(tmp_base, batch_size=1)
+        f = tmp_base / "projects" / "p" / "trace_x.json"
+        f.parent.mkdir(parents=True)
+        f.touch()
+        manifest.mark_file_completed(f)
+        manifest.mark_file_completed(f)
+        assert manifest.completed_count() == 1
+
+
+class TestBatchingBehavior:
+    """Document and verify the crash trade-off introduced by batched writes."""
+
+    def _make_trace(self, base: Path, name: str) -> Path:
+        p = base / "projects" / "proj" / name
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.touch()
+        return p
+
+    def test_mark_file_completed__within_batch__visible_via_api_before_flush(
+        self, tmp_base: Path
+    ) -> None:
+        """Buffered completions are visible through the public API (which flushes
+        before querying) even before the batch threshold is reached."""
+        manifest = MigrationManifest(tmp_base, batch_size=50)
+        f = self._make_trace(tmp_base, "trace_001.json")
+        manifest.mark_file_completed(f)
+        # completed_count() flushes first, so the buffered write IS visible.
+        assert manifest.completed_count() == 1
+        assert manifest.is_file_completed(f)
+
+    def test_mark_file_completed__crash_before_flush__unflushed_data_lost(
+        self, tmp_base: Path
+    ) -> None:
+        """Simulates a crash (new instance, no save()) with batch_size > pending count.
+
+        This is the documented trade-off: up to batch_size-1 completions may be
+        absent from the manifest after a crash. Those files will simply be
+        re-imported on resume.
+        """
+        f = self._make_trace(tmp_base, "trace_001.json")
+
+        m1 = MigrationManifest(tmp_base, batch_size=50)
+        m1.start()  # always flushes immediately
+        m1.mark_file_completed(f)  # buffered — NOT yet on disk
+        # Simulate crash: m1 is abandoned without save() or complete().
+        # __del__ is NOT called here because m1 is still in scope when m2 is created.
+
+        m2 = MigrationManifest(tmp_base, batch_size=50)
+        # The completion is not on disk — resume will re-process this file.
+        assert not m2.is_file_completed(f)
+        assert m2.completed_count() == 0
+
+    def test_mark_file_completed__batch_size_one__each_write_immediately_durable(
+        self, tmp_base: Path
+    ) -> None:
+        """With batch_size=1 every completion is immediately committed to disk."""
+        f = self._make_trace(tmp_base, "trace_001.json")
+
+        m1 = MigrationManifest(tmp_base, batch_size=1)
+        m1.start()
+        m1.mark_file_completed(f)  # auto-flushes (batch_size=1)
+
+        # New instance reads from disk — must see the completion.
+        m2 = MigrationManifest(tmp_base, batch_size=1)
+        assert m2.is_file_completed(f)
+        assert m2.completed_count() == 1
+
+    def test_mark_file_completed__batch_threshold_reached__auto_flushed_to_disk(
+        self, tmp_base: Path
+    ) -> None:
+        """When pending count hits batch_size the buffer is auto-flushed."""
+        batch_size = 3
+        files = [self._make_trace(tmp_base, f"trace_{i:03d}.json") for i in range(3)]
+
+        m1 = MigrationManifest(tmp_base, batch_size=batch_size)
+        m1.start()
+        for f in files:
+            m1.mark_file_completed(f)  # third call triggers auto-flush
+
+        # New instance must see all three completions.
+        m2 = MigrationManifest(tmp_base, batch_size=batch_size)
+        assert m2.completed_count() == 3
+        for f in files:
+            assert m2.is_file_completed(f)
+
+
+class TestResumeScenario:
+    """Simulate a real interrupted + resumed migration."""
+
+    def test_resume__interrupted_import__completed_files_skipped_and_id_map_available(
+        self, tmp_base: Path
+    ) -> None:
+        """After an interrupted run the second run skips already-flushed files.
+
+        batch_size=1 so every mark_file_completed is immediately durable —
+        this models the boundary between two process invocations where only
+        flushed data survives.
+        """
+        trace_files = []
+        for i in range(3):
+            p = tmp_base / "projects" / "proj" / f"trace_{i:03d}.json"
+            p.parent.mkdir(parents=True, exist_ok=True)
+            p.touch()
+            trace_files.append(p)
+
+        m1 = MigrationManifest(tmp_base, batch_size=1)
+        m1.start()
+        m1.add_trace_mapping("src-0", "dest-0")
+        m1.mark_file_completed(trace_files[0])
+        m1.add_trace_mapping("src-1", "dest-1")
+        m1.mark_file_completed(trace_files[1])
+        # Process crashes before trace_files[2] — no complete() call.
+
+        m2 = MigrationManifest(tmp_base)
+        assert m2.is_in_progress
+        assert m2.completed_count() == 2
+        assert m2.is_file_completed(trace_files[0])
+        assert m2.is_file_completed(trace_files[1])
+        assert not m2.is_file_completed(trace_files[2])
+
+        id_map = m2.get_trace_id_map()
+        assert id_map["src-0"] == "dest-0"
+        assert id_map["src-1"] == "dest-1"
+        assert "src-2" not in id_map
+
+    def test_reset__after_completed_files__all_state_cleared(
+        self, tmp_base: Path
+    ) -> None:
+        trace_file = tmp_base / "projects" / "p" / "trace_abc.json"
+        trace_file.parent.mkdir(parents=True)
+        trace_file.touch()
+
+        m1 = MigrationManifest(tmp_base)
+        m1.start()
+        m1.mark_file_completed(trace_file)
+
+        m2 = MigrationManifest(tmp_base)
+        m2.reset()
+
+        assert m2.completed_count() == 0
+        assert not m2.is_file_completed(trace_file)
+        assert m2.status == "not_started"

sdks/python/tests/unit/cli/test_pairing.py

@@ -0,0 +1,723 @@
+import base64
+import uuid
+from unittest.mock import MagicMock, patch
+
+import click
+import pytest
+
+from opik.cli.local_runner.pairing import (
+    PairingResult,
+    RunnerType,
+    build_pairing_link,
+    generate_runner_name,
+    hkdf_sha256,
+    resolve_project_id,
+    run_headless,
+    run_pairing,
+    validate_runner_name,
+)
+from opik.rest_api.errors.not_found_error import NotFoundError
+
+
+# Helpers — fill in the keyword-only args that production callers always supply
+# so tests don't have to repeat them. Tests override what they actually care
+# about via **overrides.
+def _resolve(api, project_name, **overrides):
+    kwargs = dict(
+        create_if_missing=False,
+        workspace=None,
+        base_url=None,
+        config_file_exists=True,
+    )
+    kwargs.update(overrides)
+    return resolve_project_id(api, project_name, **kwargs)
+
+
+def _headless(api, **overrides):
+    kwargs = dict(
+        api=api,
+        project_name="my-proj",
+        runner_name="r-1",
+        runner_type=RunnerType.ENDPOINT,
+        create_if_missing=False,
+        workspace=None,
+        base_url=None,
+        config_file_exists=True,
+    )
+    kwargs.update(overrides)
+    return run_headless(**kwargs)
+
+
+def _link(**overrides):
+    kwargs = dict(
+        base_url="https://www.comet.com/opik/api/",
+        session_id="550e8400-e29b-41d4-a716-446655440000",
+        activation_key=b"\x00" * 32,
+        project_id="660e8400-e29b-41d4-a716-446655440000",
+        project_name="my-proj",
+        runner_name="r",
+        runner_type=RunnerType.CONNECT,
+        workspace=None,
+    )
+    kwargs.update(overrides)
+    return build_pairing_link(**kwargs)
+
+
+def _pair(api, **overrides):
+    kwargs = dict(
+        api=api,
+        project_name="my-proj",
+        runner_name="r-1",
+        runner_type=RunnerType.ENDPOINT,
+        base_url="http://localhost:5173/api/",
+        workspace=None,
+        tui=None,
+        create_if_missing=False,
+        config_file_exists=True,
+    )
+    kwargs.update(overrides)
+    return run_pairing(**kwargs)
+
+
+class TestHKDF:
+    def test_hkdf__pinned_vector__matches_expected(self):
+        ikm = bytes(range(32))
+        salt = uuid.UUID("00000000-0000-0000-0000-000000000001").bytes
+        info = b"opik-bridge-v1"
+
+        okm = hkdf_sha256(ikm=ikm, salt=salt, info=info)
+
+        assert (
+            okm.hex()
+            == "9647b7959765ecad68dfef02f31cfca7f7901a9076c15ebabb1f35d015f71198"
+        )
+
+    def test_hkdf__any_input__output_is_32_bytes(self):
+        okm = hkdf_sha256(ikm=b"\x00" * 32, salt=b"\x00" * 16, info=b"test")
+        assert len(okm) == 32
+
+    def test_hkdf__different_info__different_output(self):
+        a = hkdf_sha256(ikm=b"\x00" * 32, salt=b"\x00" * 16, info=b"a")
+        b = hkdf_sha256(ikm=b"\x00" * 32, salt=b"\x00" * 16, info=b"b")
+        assert a != b
+
+
+class TestUUIDBytesOrder:
+    def test_uuid_bytes__known_uuid__matches_java_big_endian(self):
+        u = uuid.UUID("550e8400-e29b-41d4-a716-446655440000")
+        assert u.bytes.hex() == "550e8400e29b41d4a716446655440000"
+
+
+class TestBuildPairingLink:
+    def test_build_pairing_link__valid_inputs__correct_payload_layout(self):
+        session_id = "550e8400-e29b-41d4-a716-446655440000"
+        project_id = "660e8400-e29b-41d4-a716-446655440000"
+        activation_key = b"\xaa" * 32
+        runner_name = "my-runner"
+
+        link = _link(
+            base_url="https://www.comet.com/opik/api/",
+            session_id=session_id,
+            activation_key=activation_key,
+            project_id=project_id,
+            runner_name=runner_name,
+        )
+
+        assert link.startswith("https://www.comet.com/opik/pair/v1?")
+
+        fragment = link.split("#", 1)[1]
+        padding_needed = (4 - len(fragment) % 4) % 4
+        payload = base64.urlsafe_b64decode(fragment + "=" * padding_needed)
+
+        assert payload[0:16] == uuid.UUID(session_id).bytes
+        assert payload[16:48] == activation_key
+        assert payload[48:64] == uuid.UUID(project_id).bytes
+        assert payload[64] == len(runner_name.encode("utf-8"))
+        name_end = 65 + payload[64]
+        assert payload[65:name_end] == runner_name.encode("utf-8")
+        # Default runner type is CONNECT (0x00)
+        assert payload[name_end] == 0x00
+
+    def test_build_pairing_link__endpoint_type__encodes_0x01(self):
+        link = _link(
+            base_url="http://localhost:5173/api/",
+            session_id="550e8400-e29b-41d4-a716-446655440000",
+            activation_key=b"\x00" * 32,
+            project_id="660e8400-e29b-41d4-a716-446655440000",
+            runner_name="r",
+            runner_type=RunnerType.ENDPOINT,
+        )
+        fragment = link.split("#", 1)[1]
+        padding_needed = (4 - len(fragment) % 4) % 4
+        payload = base64.urlsafe_b64decode(fragment + "=" * padding_needed)
+        # name_len=1, name="r", then type byte
+        assert payload[66] == 0x01
+
+    def test_build_pairing_link__cloud_url__no_double_opik_path(self):
+        link = _link(
+            base_url="https://www.comet.com/opik/api/",
+            session_id="550e8400-e29b-41d4-a716-446655440000",
+            activation_key=b"\x00" * 32,
+            project_id="660e8400-e29b-41d4-a716-446655440000",
+            runner_name="r",
+        )
+        assert "/opik/opik/" not in link
+
+    def test_build_pairing_link__localhost_url__correct_prefix(self):
+        link = _link(
+            base_url="http://localhost:5173/api/",
+            session_id="550e8400-e29b-41d4-a716-446655440000",
+            activation_key=b"\x00" * 32,
+            project_id="660e8400-e29b-41d4-a716-446655440000",
+            runner_name="r",
+        )
+        assert link.startswith("http://localhost:5173/opik/pair/v1?")
+
+    def test_build_pairing_link__always_includes_url_param(self):
+        link = _link(
+            base_url="http://localhost:5173/api/",
+            workspace=None,
+        )
+        # `/api` is stripped — we send the URL the user would visit in a
+        # browser, not the internal API endpoint.
+        assert "url=http%3A%2F%2Flocalhost%3A5173%2F" in link
+        assert "%2Fapi%2F" not in link
+
+    def test_build_pairing_link__cloud_api_url__strips_api_suffix(self):
+        link = _link(
+            base_url="https://www.comet.com/opik/api/",
+            workspace=None,
+        )
+        assert "url=https%3A%2F%2Fwww.comet.com%2Fopik%2F" in link
+
+    def test_build_pairing_link__url_param_precedes_workspace(self):
+        link = _link(
+            base_url="http://localhost:5173/api/",
+            workspace="team-a",
+        )
+        # All query params present; order is stable for FE parsing.
+        assert "url=" in link
+        assert "project=my-proj" in link
+        assert "workspace=team-a" in link
+
+    def test_build_pairing_link__project_name_with_special_chars__url_encoded(
+        self,
+    ):
+        link = _link(project_name="My Project / Demo")
+        # Spaces and `/` must be percent-encoded so the query stays parseable.
+        assert "project=My%20Project%20%2F%20Demo" in link
+
+    def test_build_pairing_link__workspace_with_special_chars__url_encoded(
+        self,
+    ):
+        # `&` and `=` in a workspace name would otherwise split the query and
+        # the FE's URLSearchParams would read a truncated value.
+        link = _link(workspace="weird&name=oops")
+        assert "workspace=weird%26name%3Doops" in link
+
+
+class TestResolveProjectId:
+    def test_resolve_project_id__project_exists__returns_id(self):
+        api = MagicMock()
+        project = MagicMock()
+        project.id = "proj-uuid-123"
+        api.projects.retrieve_project.return_value = project
+
+        result = _resolve(api, "my-project")
+        assert result == "proj-uuid-123"
+        api.projects.retrieve_project.assert_called_once_with(name="my-project")
+
+    def test_resolve_project_id__project_missing__uses_404_hint(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        # Project retrieve returns Opik's custom ErrorMessage shape.
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["Project not found"]}
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(api, "nonexistent")
+        message = exc_info.value.message
+        assert "Could not retrieve project 'nonexistent'" in message
+        assert "Project not found" in message
+        assert "check the project name and try again" in message
+
+    def test_resolve_project_id__unauthorized__surfaces_server_message_and_docs(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=401,
+            body={"code": 401, "message": "API key should be provided"},
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(api, "my-project")
+        message = exc_info.value.message
+        assert "Could not retrieve project 'my-project'" in message
+        assert "API key should be provided" in message
+        assert "Run: opik configure" in message
+        assert (
+            "https://www.comet.com/docs/opik/tracing/advanced/sdk_configuration"
+            in message
+        )
+
+    def test_resolve_project_id__forbidden__shares_auth_hint(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=403,
+            body={"code": 403, "message": "User is not allowed to access workspace"},
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(api, "my-project")
+        message = exc_info.value.message
+        assert "User is not allowed to access workspace" in message
+        assert "Run: opik configure" in message
+
+    def test_resolve_project_id__server_error__falls_back_to_generic_hint(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=500,
+            body="Internal Server Error",
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(api, "my-project")
+        message = exc_info.value.message
+        assert "Internal Server Error" in message
+        assert "verify your Opik configuration and connectivity" in message
+        assert (
+            "https://www.comet.com/docs/opik/tracing/advanced/sdk_configuration"
+            in message
+        )
+
+    def test_resolve_project_id__empty_body__falls_back_to_status_text(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(status_code=502, body=None)
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(api, "my-project")
+        assert "server returned HTTP 502" in exc_info.value.message
+
+    def test_resolve_project_id__missing_and_create_flag__creates_and_resolves(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        created = MagicMock()
+        created.id = "proj-uuid-new"
+        api.projects.retrieve_project.side_effect = [
+            ApiError(status_code=404, body={"errors": ["Project not found"]}),
+            created,
+        ]
+
+        result = _resolve(api, "new-project", create_if_missing=True)
+        assert result == "proj-uuid-new"
+        api.projects.create_project.assert_called_once_with(name="new-project")
+        assert api.projects.retrieve_project.call_count == 2
+
+    def test_resolve_project_id__known_missing__skips_initial_lookup(self):
+        api = MagicMock()
+        created = MagicMock()
+        created.id = "proj-uuid-new"
+        api.projects.retrieve_project.return_value = created
+
+        result = _resolve(
+            api,
+            "new-project",
+            create_if_missing=True,
+            project_known_missing=True,
+        )
+        assert result == "proj-uuid-new"
+        api.projects.create_project.assert_called_once_with(name="new-project")
+        # Interactive preflight already saw the 404 — only the post-create
+        # lookup should run.
+        assert api.projects.retrieve_project.call_count == 1
+
+    def test_resolve_project_id__missing_no_flag__does_not_create(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["Project not found"]}
+        )
+
+        with pytest.raises(click.ClickException):
+            _resolve(api, "missing", create_if_missing=False)
+        api.projects.create_project.assert_not_called()
+
+    def test_resolve_project_id__non_404_with_flag__does_not_create(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=401, body={"message": "unauthorized"}
+        )
+
+        with pytest.raises(click.ClickException):
+            _resolve(api, "x", create_if_missing=True)
+        api.projects.create_project.assert_not_called()
+
+    def test_resolve_project_id__create_fails__raises_clean_error(self):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["Project not found"]}
+        )
+        api.projects.create_project.side_effect = ApiError(
+            status_code=403, body={"message": "User cannot create projects"}
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(api, "blocked", create_if_missing=True)
+        message = exc_info.value.message
+        assert "Could not create project 'blocked'" in message
+        assert "User cannot create projects" in message
+
+    def test_resolve_project_id__retrieval_error_with_config_context__appends_workspace_and_url(
+        self,
+    ):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["Project not found"]}
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(
+                api,
+                "missing-proj",
+                workspace="team-a",
+                base_url="https://opik.example.com/api/",
+            )
+        message = exc_info.value.message
+        assert "team-a" in message
+        assert "https://opik.example.com/api/" in message
+
+    def test_resolve_project_id__retrieval_error_without_config_file__appends_opik_configure_hint(
+        self,
+    ):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=401, body={"message": "API key should be provided"}
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(
+                api,
+                "my-project",
+                workspace="default",
+                base_url="https://www.comet.com/opik/api/",
+                config_file_exists=False,
+            )
+        message = exc_info.value.message
+        assert "~/.opik.config" in message
+        assert "Run: opik configure" in message
+
+    def test_resolve_project_id__retrieval_error_with_config_file__omits_no_config_hint(
+        self,
+    ):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=500, body="Internal Server Error"
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(
+                api,
+                "p",
+                workspace="default",
+                base_url="https://www.comet.com/opik/api/",
+                config_file_exists=True,
+            )
+        assert "~/.opik.config" not in exc_info.value.message
+
+    def test_resolve_project_id__create_fails_with_config_context__includes_workspace_and_url(
+        self,
+    ):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["Project not found"]}
+        )
+        api.projects.create_project.side_effect = ApiError(
+            status_code=403, body={"message": "User cannot create projects"}
+        )
+
+        with pytest.raises(click.ClickException) as exc_info:
+            _resolve(
+                api,
+                "blocked",
+                create_if_missing=True,
+                workspace="team-a",
+                base_url="https://opik.example.com/api/",
+                config_file_exists=False,
+            )
+        message = exc_info.value.message
+        assert "Could not create project 'blocked'" in message
+        assert "team-a" in message
+        assert "https://opik.example.com/api/" in message
+        assert "Run: opik configure" in message
+
+
+class TestValidateRunnerName:
+    def test_validate_runner_name__valid_name__passes(self):
+        validate_runner_name("my-runner")
+
+    def test_validate_runner_name__empty__raises(self):
+        with pytest.raises(Exception, match="empty"):
+            validate_runner_name("")
+
+    def test_validate_runner_name__whitespace_only__raises(self):
+        with pytest.raises(Exception, match="empty"):
+            validate_runner_name("   ")
+
+    def test_validate_runner_name__over_128_chars__raises(self):
+        with pytest.raises(Exception, match="128 characters"):
+            validate_runner_name("x" * 129)
+
+    def test_validate_runner_name__over_255_utf8_bytes__raises(self):
+        name = "\U0001f600" * 64
+        with pytest.raises(Exception, match="255 UTF-8 bytes"):
+            validate_runner_name(name)
+
+
+class TestGenerateRunnerName:
+    def test_generate_runner_name__explicit_name__returns_it(self):
+        assert generate_runner_name("my-name") == "my-name"
+
+    def test_generate_runner_name__none__generates_random_hex(self):
+        name = generate_runner_name(None)
+        assert "-" in name
+        hex_part = name.rsplit("-", 1)[1]
+        assert len(hex_part) == 6
+        int(hex_part, 16)
+
+
+class TestRunPairing:
+    SESSION_ID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
+    RUNNER_ID = "11111111-2222-3333-4444-555555555555"
+    PROJECT_ID = "66666666-7777-8888-9999-aaaaaaaaaaaa"
+
+    def _make_api(self, project_id=None):
+        if project_id is None:
+            project_id = self.PROJECT_ID
+        api = MagicMock()
+        project = MagicMock()
+        project.id = project_id
+        api.projects.retrieve_project.return_value = project
+
+        session_resp = MagicMock()
+        session_resp.session_id = self.SESSION_ID
+        session_resp.runner_id = self.RUNNER_ID
+        api.pairing.create_pairing_session.return_value = session_resp
+
+        runner = MagicMock()
+        runner.status = "connected"
+        api.runners.get_runner.return_value = runner
+
+        return api
+
+    @patch("opik.cli.local_runner.pairing.time.sleep")
+    def test_run_pairing__happyflow__returns_result(self, mock_sleep):
+        api = self._make_api()
+        result = _pair(
+            api=api,
+            project_name="my-proj",
+            runner_name="test-runner",
+            runner_type=RunnerType.CONNECT,
+            base_url="http://localhost:5173/api/",
+            ttl_seconds=1,
+        )
+
+        assert isinstance(result, PairingResult)
+        assert result.runner_id == self.RUNNER_ID
+        assert result.project_name == "my-proj"
+        assert result.project_id == self.PROJECT_ID
+        assert len(result.bridge_key) == 32
+
+    @patch("opik.cli.local_runner.pairing.time.sleep")
+    def test_run_pairing__with_tui__calls_started_and_completed(self, mock_sleep):
+        api = self._make_api()
+        tui = MagicMock()
+
+        _pair(
+            api=api,
+            project_name="my-proj",
+            runner_name="test-runner",
+            runner_type=RunnerType.CONNECT,
+            base_url="http://localhost:5173/api/",
+            tui=tui,
+            ttl_seconds=1,
+        )
+
+        tui.pairing_started.assert_called_once()
+        tui.pairing_completed.assert_called_once()
+
+    @patch("opik.cli.local_runner.pairing.time.sleep")
+    def test_run_pairing__404_during_poll__retries_until_connected(self, mock_sleep):
+        api = self._make_api()
+        err = NotFoundError(body=None)
+        runner = MagicMock()
+        runner.status = "connected"
+        api.runners.get_runner.side_effect = [err, err, runner]
+
+        result = _pair(
+            api=api,
+            project_name="my-proj",
+            runner_name="test-runner",
+            runner_type=RunnerType.ENDPOINT,
+            base_url="http://localhost:5173/api/",
+            ttl_seconds=1,
+        )
+        assert result.runner_id == self.RUNNER_ID
+        assert api.runners.get_runner.call_count == 3
+
+    @patch("opik.cli.local_runner.pairing.time.monotonic")
+    @patch("opik.cli.local_runner.pairing.time.sleep")
+    def test_run_pairing__timeout__calls_tui_pairing_failed(
+        self, mock_sleep, mock_monotonic
+    ):
+        api = self._make_api()
+        runner = MagicMock()
+        runner.status = "pairing"
+        api.runners.get_runner.return_value = runner
+
+        # Strictly-increasing clock: each call returns previous + 1000s. This
+        # guarantees the while-loop exits on the call immediately after
+        # `deadline = monotonic() + ttl_seconds` is computed, regardless of
+        # how many other calls (from logging, coverage, plugins, etc.) land
+        # on the patched monotonic before or after the test's own calls.
+        clock = {"t": 0.0}
+
+        def advancing_monotonic():
+            clock["t"] += 1000.0
+            return clock["t"]
+
+        mock_monotonic.side_effect = advancing_monotonic
+
+        tui = MagicMock()
+        with pytest.raises(click.ClickException) as exc_info:
+            _pair(
+                api=api,
+                project_name="my-proj",
+                runner_name="test-runner",
+                runner_type=RunnerType.CONNECT,
+                base_url="http://localhost:5173/api/",
+                tui=tui,
+                ttl_seconds=300,
+            )
+        assert "timed out" in exc_info.value.message
+
+        tui.pairing_failed.assert_called_once_with("timed out")
+
+    @patch("opik.cli.local_runner.pairing.time.sleep")
+    def test_run_pairing__keyboard_interrupt__calls_tui_pairing_failed(
+        self, mock_sleep
+    ):
+        api = self._make_api()
+        api.runners.get_runner.side_effect = KeyboardInterrupt
+
+        tui = MagicMock()
+        with pytest.raises(KeyboardInterrupt):
+            _pair(
+                api=api,
+                project_name="my-proj",
+                runner_name="test-runner",
+                runner_type=RunnerType.CONNECT,
+                base_url="http://localhost:5173/api/",
+                tui=tui,
+                ttl_seconds=1,
+            )
+
+        tui.pairing_failed.assert_called_once_with("interrupted")
+
+    @patch("opik.cli.local_runner.pairing.time.sleep")
+    def test_run_pairing__non_404_api_error__propagates(self, mock_sleep):
+        from opik.rest_api.core.api_error import ApiError
+
+        api = self._make_api()
+        api.runners.get_runner.side_effect = ApiError(
+            status_code=429, body="rate limited"
+        )
+
+        with pytest.raises(ApiError) as exc_info:
+            _pair(
+                api=api,
+                project_name="my-proj",
+                runner_name="test-runner",
+                runner_type=RunnerType.CONNECT,
+                base_url="http://localhost:5173/api/",
+                ttl_seconds=1,
+            )
+        assert exc_info.value.status_code == 429
+
+
+class TestRunHeadless:
+    SESSION_ID = "aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee"
+    RUNNER_ID = "11111111-2222-3333-4444-555555555555"
+    PROJECT_ID = "66666666-7777-8888-9999-aaaaaaaaaaaa"
+
+    def _make_api(self):
+        api = MagicMock()
+        project = MagicMock()
+        project.id = self.PROJECT_ID
+        api.projects.retrieve_project.return_value = project
+
+        session_resp = MagicMock()
+        session_resp.session_id = self.SESSION_ID
+        session_resp.runner_id = self.RUNNER_ID
+        api.pairing.create_pairing_session.return_value = session_resp
+        api.pairing.activate_pairing_session.return_value = None
+
+        return api
+
+    def test_run_headless__creates_and_self_activates(self):
+        api = self._make_api()
+        result = _headless(
+            api=api,
+            project_name="my-proj",
+            runner_name="test-runner",
+            runner_type=RunnerType.ENDPOINT,
+        )
+
+        assert result.runner_id == self.RUNNER_ID
+        assert result.project_id == self.PROJECT_ID
+        assert result.bridge_key == b""
+        api.pairing.create_pairing_session.assert_called_once()
+        api.pairing.activate_pairing_session.assert_called_once()
+
+    def test_run_headless__no_polling(self):
+        api = self._make_api()
+        _headless(
+            api=api,
+            project_name="my-proj",
+            runner_name="test-runner",
+            runner_type=RunnerType.ENDPOINT,
+        )
+        # No get_runner polling — headless activates immediately
+        api.runners.get_runner.assert_not_called()
+
+    def test_run_headless__connect_type__raises(self):
+        api = self._make_api()
+        with pytest.raises(click.ClickException, match="not supported"):
+            _headless(
+                api=api,
+                project_name="my-proj",
+                runner_name="test-runner",
+                runner_type=RunnerType.CONNECT,
+            )

sdks/python/tests/unit/cli/test_run.py

@@ -0,0 +1,175 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from opik.cli.local_runner.preflight import (
+    maybe_auto_configure,
+    should_create_project,
+)
+from opik.rest_api.core.api_error import ApiError
+
+
+class TestShouldCreateProject:
+    def test__headless__returns_create_without_lookup(self):
+        api = MagicMock()
+        result = should_create_project(api, "any", workspace=None, headless=True)
+        # Headless skips preflight, so we want to create but don't yet know
+        # whether the project is actually missing — the resolver must look up.
+        assert result == (True, False)
+        api.projects.retrieve_project.assert_not_called()
+
+    def test__project_exists__returns_no_create(self):
+        api = MagicMock()
+        api.projects.retrieve_project.return_value = MagicMock(id="proj-1")
+        result = should_create_project(api, "exists", workspace="ws", headless=False)
+        assert result == (False, False)
+
+    def test__non_404_error__returns_no_create_to_let_downstream_format(self):
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=401, body={"message": "unauthorized"}
+        )
+        result = should_create_project(api, "x", workspace="ws", headless=False)
+        assert result == (False, False)
+
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__missing_and_not_tty__returns_no_create(self, mock_stdin):
+        mock_stdin.isatty.return_value = False
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["not found"]}
+        )
+        result = should_create_project(api, "missing", workspace="ws", headless=False)
+        assert result == (False, False)
+
+    @patch("opik.cli.local_runner.preflight.click.confirm", return_value=True)
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__missing_and_tty_user_confirms__returns_create_and_known_missing(
+        self, mock_stdin, mock_confirm
+    ):
+        mock_stdin.isatty.return_value = True
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["not found"]}
+        )
+        result = should_create_project(api, "missing", workspace="ws", headless=False)
+        # Interactive preflight saw 404 and user confirmed — downstream can
+        # skip the redundant lookup.
+        assert result == (True, True)
+        prompt_text = mock_confirm.call_args[0][0]
+        assert "missing" in prompt_text
+        assert "ws" in prompt_text
+
+    @patch("opik.cli.local_runner.preflight.click.confirm", return_value=False)
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__missing_and_tty_user_declines__returns_no_create(
+        self, mock_stdin, mock_confirm
+    ):
+        mock_stdin.isatty.return_value = True
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(
+            status_code=404, body={"errors": ["not found"]}
+        )
+        result = should_create_project(api, "missing", workspace="ws", headless=False)
+        assert result == (False, False)
+
+    @patch("opik.cli.local_runner.preflight.click.confirm", return_value=True)
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__no_workspace__prompt_omits_workspace_label(
+        self, mock_stdin, mock_confirm
+    ):
+        mock_stdin.isatty.return_value = True
+        api = MagicMock()
+        api.projects.retrieve_project.side_effect = ApiError(status_code=404, body={})
+        should_create_project(api, "missing", workspace=None, headless=False)
+        prompt_text = mock_confirm.call_args[0][0]
+        assert "in workspace" not in prompt_text
+
+
+class TestMaybeAutoConfigure:
+    def _patch_env(self, monkeypatch, opik_api_key=None):
+        if opik_api_key is None:
+            monkeypatch.delenv("OPIK_API_KEY", raising=False)
+        else:
+            monkeypatch.setenv("OPIK_API_KEY", opik_api_key)
+
+    def _config_probe(self, *, config_file_exists=False, api_key=None):
+        probe = MagicMock()
+        probe.config_file_exists = config_file_exists
+        probe.api_key = api_key
+        return probe
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    def test__non_interactive_flag__skips(self, mock_configure, monkeypatch):
+        self._patch_env(monkeypatch)
+        maybe_auto_configure(api_key_arg=None, non_interactive=True, headless=False)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    def test__headless__skips(self, mock_configure, monkeypatch):
+        self._patch_env(monkeypatch)
+        maybe_auto_configure(api_key_arg=None, non_interactive=False, headless=True)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    def test__api_key_arg__skips(self, mock_configure, monkeypatch):
+        self._patch_env(monkeypatch)
+        maybe_auto_configure(api_key_arg="abc", non_interactive=False, headless=False)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    def test__env_api_key__skips(self, mock_configure, monkeypatch):
+        self._patch_env(monkeypatch, opik_api_key="from-env")
+        maybe_auto_configure(api_key_arg=None, non_interactive=False, headless=False)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__stdin_not_tty__skips(self, mock_stdin, mock_configure, monkeypatch):
+        self._patch_env(monkeypatch)
+        mock_stdin.isatty.return_value = False
+        maybe_auto_configure(api_key_arg=None, non_interactive=False, headless=False)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    @patch("opik.cli.local_runner.preflight.OpikConfig")
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__config_file_already_exists__skips(
+        self, mock_stdin, mock_config_cls, mock_configure, monkeypatch
+    ):
+        self._patch_env(monkeypatch)
+        mock_stdin.isatty.return_value = True
+        mock_config_cls.return_value = self._config_probe(config_file_exists=True)
+        maybe_auto_configure(api_key_arg=None, non_interactive=False, headless=False)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    @patch("opik.cli.local_runner.preflight.OpikConfig")
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__probe_resolves_api_key_from_env_chain__skips(
+        self, mock_stdin, mock_config_cls, mock_configure, monkeypatch
+    ):
+        # The OpikConfig probe layer can pick up an API key from elsewhere
+        # (e.g. .env file) even when OPIK_API_KEY isn't in os.environ. We
+        # treat that as "already configured enough" and skip the prompt.
+        self._patch_env(monkeypatch)
+        mock_stdin.isatty.return_value = True
+        mock_config_cls.return_value = self._config_probe(api_key="from-dotenv")
+        maybe_auto_configure(api_key_arg=None, non_interactive=False, headless=False)
+        mock_configure.assert_not_called()
+
+    @patch("opik.cli.configure.run_interactive_configure")
+    @patch("opik.cli.local_runner.preflight.OpikConfig")
+    @patch("opik.cli.local_runner.preflight.sys.stdin")
+    def test__no_config_and_interactive__runs_configure(
+        self, mock_stdin, mock_config_cls, mock_configure, monkeypatch
+    ):
+        self._patch_env(monkeypatch)
+        mock_stdin.isatty.return_value = True
+        mock_config_cls.return_value = self._config_probe()
+        maybe_auto_configure(api_key_arg=None, non_interactive=False, headless=False)
+        mock_configure.assert_called_once()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

sdks/python/tests/unit/cli/test_tags_roundtrip.py

@@ -0,0 +1,440 @@
+"""Unit tests verifying that item tags round-trip through CLI export/import.
+
+Covers the regression where ``opik export`` / ``opik import`` silently dropped
+tags for prompts, datasets, and experiments because the exporters used
+hand-written field allowlists that omitted ``tags`` (and the importers never
+forwarded them). Traces and spans already round-tripped correctly.
+"""
+
+import json
+import sys
+from pathlib import Path
+from unittest.mock import Mock, MagicMock, patch
+
+# The experiment/prompt import modules pull in a prompt module that is awkward
+# to import in isolation; mirror the shim used by test_import_experiment.py.
+sys.modules.setdefault("opik.api_objects.prompt.prompt", MagicMock())
+
+from opik.cli.imports.prompt import import_prompts_from_directory  # noqa: E402
+from opik.cli.imports.dataset import import_datasets_from_directory  # noqa: E402
+from opik.cli.exports.prompt import export_single_prompt  # noqa: E402
+from opik.cli.imports.experiment import (  # noqa: E402
+    ExperimentData,
+    recreate_experiment,
+)
+from opik.cli.exports.utils import create_experiment_data_structure  # noqa: E402
+
+
+class TestPromptTagsImport:
+    def test_import_text_prompt__with_tags__forwards_tags(self, tmp_path: Path) -> None:
+        prompt_file = tmp_path / "prompt_p1.json"
+        prompt_file.write_text(
+            json.dumps(
+                {
+                    "name": "greeting",
+                    "current_version": {
+                        "prompt": "Hello {{name}}",
+                        "type": "mustache",
+                        "template_structure": "text",
+                        "tags": ["prod", "greeting"],
+                    },
+                    "history": [],
+                }
+            )
+        )
+
+        client = Mock()
+        client.create_prompt = Mock()
+
+        import_prompts_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.create_prompt.assert_called_once()
+        assert client.create_prompt.call_args.kwargs["tags"] == ["prod", "greeting"]
+
+    def test_import_chat_prompt__with_tags__forwards_tags(self, tmp_path: Path) -> None:
+        prompt_file = tmp_path / "prompt_c1.json"
+        prompt_file.write_text(
+            json.dumps(
+                {
+                    "name": "chat",
+                    "current_version": {
+                        "prompt": [{"role": "user", "content": "hi"}],
+                        "type": "mustache",
+                        "template_structure": "chat",
+                        "tags": ["chatty"],
+                    },
+                    "history": [],
+                }
+            )
+        )
+
+        client = Mock()
+        client.create_chat_prompt = Mock()
+
+        import_prompts_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.create_chat_prompt.assert_called_once()
+        assert client.create_chat_prompt.call_args.kwargs["tags"] == ["chatty"]
+
+
+class TestDatasetTagsImport:
+    def _write_dataset(self, tmp_path: Path, payload: dict) -> None:
+        (tmp_path / "dataset_d1.json").write_text(json.dumps(payload))
+
+    def _make_client(self) -> Mock:
+        client = Mock()
+        # Force the create path (get_dataset raises -> create_dataset used).
+        client.get_dataset = Mock(side_effect=Exception("not found"))
+        created = Mock()
+        created.id = "new-ds-id"
+        client.create_dataset = Mock(return_value=created)
+        client.rest_client = Mock()
+        client.rest_client.datasets = Mock()
+        client.rest_client.datasets.update_dataset = Mock()
+        return client
+
+    def test_import_dataset__flat_format_with_tags__applies_tags(
+        self, tmp_path: Path
+    ) -> None:
+        self._write_dataset(
+            tmp_path,
+            {"name": "ds", "tags": ["a", "b"], "items": []},
+        )
+        client = self._make_client()
+
+        import_datasets_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.rest_client.datasets.update_dataset.assert_called_once()
+        call = client.rest_client.datasets.update_dataset.call_args
+        assert call.args[0] == "new-ds-id"
+        assert call.kwargs["tags"] == ["a", "b"]
+
+    def test_import_dataset__nested_format_with_tags__applies_tags(
+        self, tmp_path: Path
+    ) -> None:
+        self._write_dataset(
+            tmp_path,
+            {"dataset": {"name": "ds", "id": "x", "tags": ["c"]}, "items": []},
+        )
+        client = self._make_client()
+
+        import_datasets_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.rest_client.datasets.update_dataset.assert_called_once()
+        assert client.rest_client.datasets.update_dataset.call_args.kwargs["tags"] == [
+            "c"
+        ]
+
+    def test_import_dataset__no_tags_field__skips_update(self, tmp_path: Path) -> None:
+        self._write_dataset(tmp_path, {"name": "ds", "items": []})
+        client = self._make_client()
+
+        import_datasets_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.rest_client.datasets.update_dataset.assert_not_called()
+
+    def test_import_dataset__empty_tags_list__clears_tags(self, tmp_path: Path) -> None:
+        # An explicit empty list must still call the REST update so a
+        # re-import can clear pre-existing tags on the destination.
+        self._write_dataset(tmp_path, {"name": "ds", "tags": [], "items": []})
+        client = self._make_client()
+
+        import_datasets_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.rest_client.datasets.update_dataset.assert_called_once()
+        assert client.rest_client.datasets.update_dataset.call_args.kwargs["tags"] == []
+
+    def test_import_dataset__existing_dataset_with_description__preserves_description(
+        self, tmp_path: Path
+    ) -> None:
+        # Tags-only import must not silently null out a description the
+        # destination dataset already had (the backend's update PUT doesn't
+        # COALESCE description against the existing row).
+        self._write_dataset(
+            tmp_path,
+            {"name": "ds", "tags": ["a"], "items": []},
+        )
+        client = self._make_client()
+        # Re-import onto a dataset that already exists (get_dataset succeeds).
+        existing_dataset = Mock()
+        existing_dataset.id = "existing-ds-id"
+        client.get_dataset = Mock(return_value=existing_dataset)
+        existing_remote = Mock()
+        existing_remote.description = "original description"
+        existing_remote.visibility = "private"
+        client.rest_client.datasets.get_dataset_by_id = Mock(
+            return_value=existing_remote
+        )
+
+        import_datasets_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        client.rest_client.datasets.update_dataset.assert_called_once()
+        call = client.rest_client.datasets.update_dataset.call_args
+        assert call.kwargs["description"] == "original description"
+        assert call.kwargs["visibility"] == "private"
+        assert call.kwargs["tags"] == ["a"]
+
+    def test_import_dataset__update_dataset_raises__continues_and_inserts_items(
+        self, tmp_path: Path
+    ) -> None:
+        # A tag-update failure must be swallowed with a warning so the rest of
+        # the import (item insertion, count, manifest) still proceeds.
+        self._write_dataset(
+            tmp_path,
+            {"name": "ds", "tags": ["a"], "items": [{"id": "1", "input": "x"}]},
+        )
+        client = self._make_client()
+        client.rest_client.datasets.update_dataset = Mock(side_effect=Exception("boom"))
+
+        result = import_datasets_from_directory(
+            client=client,
+            source_dir=tmp_path,
+            project_name="proj",
+            dry_run=False,
+            name_pattern=None,
+            debug=False,
+        )
+
+        # The import did not abort: items were still inserted and the dataset
+        # counted as imported.
+        client.create_dataset.return_value.insert.assert_called_once()
+        assert result["datasets"] == 1
+        assert result["datasets_errors"] == 0
+
+
+class TestPromptTagsExport:
+    """Prompt tags are container-level and are dropped by the direct
+    get_prompt() lookup; the exporter must recover them via search_prompts.
+    """
+
+    def test_export_single_prompt__object_has_tags__skips_search(
+        self, tmp_path: Path
+    ) -> None:
+        prompt = Mock()
+        prompt.id = "p1"
+        prompt.name = "greeting"
+        prompt.tags = ["prod"]  # direct lookup already carries the tags
+        client = Mock()
+        # No history available; _safe_prompt_history swallows the error.
+        client.get_prompt_history = Mock(side_effect=Exception("no history"))
+
+        count = export_single_prompt(
+            client=client,
+            prompt=prompt,
+            output_dir=tmp_path,
+            project_name="proj",
+            max_results=None,
+            force=True,
+            debug=False,
+            format="json",
+        )
+
+        assert count == 1
+        data = json.loads((tmp_path / "prompt_p1.json").read_text())
+        assert data["current_version"]["tags"] == ["prod"]
+        # Tags were already present, so no recovery search is needed.
+        client.search_prompts.assert_not_called()
+
+    def test_export_single_prompt__container_only_tags__recovered_from_search(
+        self, tmp_path: Path
+    ) -> None:
+        prompt = Mock()
+        prompt.id = "p1"
+        prompt.name = "greeting"
+        prompt.tags = []  # direct lookup dropped the container-level tags
+
+        candidate = Mock()
+        candidate.name = "greeting"
+        candidate.tags = ["prod", "greeting"]
+        client = Mock()
+        client.search_prompts = Mock(return_value=[candidate])
+        # No history available; _safe_prompt_history swallows the error.
+        client.get_prompt_history = Mock(side_effect=Exception("no history"))
+
+        count = export_single_prompt(
+            client=client,
+            prompt=prompt,
+            output_dir=tmp_path,
+            project_name="proj",
+            max_results=None,
+            force=True,
+            debug=False,
+            format="json",
+        )
+
+        assert count == 1
+        data = json.loads((tmp_path / "prompt_p1.json").read_text())
+        assert data["current_version"]["tags"] == ["prod", "greeting"]
+
+    def test_export_single_prompt__search_fails__still_exports_with_fallback(
+        self, tmp_path: Path
+    ) -> None:
+        # Tag resolution is best-effort: a search_prompts failure is logged but
+        # must not abort the export — the prompt still exports with whatever tags
+        # the object carried.
+        prompt = Mock()
+        prompt.id = "p1"
+        prompt.name = "greeting"
+        prompt.tags = []  # nothing on the object; recovery will be attempted
+        client = Mock()
+        client.search_prompts = Mock(side_effect=Exception("api down"))
+        client.get_prompt_history = Mock(side_effect=Exception("no history"))
+
+        count = export_single_prompt(
+            client=client,
+            prompt=prompt,
+            output_dir=tmp_path,
+            project_name="proj",
+            max_results=None,
+            force=True,
+            debug=False,
+            format="json",
+        )
+
+        assert count == 1
+        data = json.loads((tmp_path / "prompt_p1.json").read_text())
+        assert data["current_version"]["tags"] == []
+
+
+class TestExperimentTags:
+    def test_export_experiment_structure__with_tags__includes_tags(self) -> None:
+        experiment = Mock()
+        experiment.id = "exp-1"
+        experiment.name = "e"
+        experiment.dataset_name = "ds"
+        data_obj = Mock()
+        data_obj.tags = ["t1", "t2"]
+        experiment.get_experiment_data = Mock(return_value=data_obj)
+
+        structure = create_experiment_data_structure(experiment, [])
+
+        assert structure["experiment"]["tags"] == ["t1", "t2"]
+
+    def test_import_experiment_from_disk__with_tags__forwards_tags(self) -> None:
+        client = Mock()
+        client.flush = Mock(return_value=True)
+        client.get_or_create_dataset = Mock(return_value=Mock(name="ds"))
+        created_experiment = Mock()
+        created_experiment.id = "exp-123"
+        client.create_experiment = Mock(return_value=created_experiment)
+        client._rest_client = Mock()
+
+        experiment_data = ExperimentData(
+            experiment={
+                "id": "exp-123",
+                "name": "test-experiment",
+                "dataset_name": "test-dataset",
+                "type": "regular",
+                "tags": ["golden", "regression"],
+            },
+            items=[],
+        )
+
+        with patch("opik.cli.imports.experiment.id_helpers_module") as mock_id_helpers:
+            mock_id_helpers.generate_id = Mock(return_value="generated-id")
+
+            recreate_experiment(
+                client,
+                experiment_data,
+                "test-project",
+                {},  # trace_id_map
+                {},  # dataset_item_id_map
+                dry_run=False,
+                debug=False,
+                # target_project_name defaults to None -> import-from-disk path
+            )
+
+        client.create_experiment.assert_called_once()
+        assert client.create_experiment.call_args.kwargs["tags"] == [
+            "golden",
+            "regression",
+        ]
+
+    def test_import_experiment_from_disk__empty_tags_list__forwards_empty_tags(
+        self,
+    ) -> None:
+        # An explicit empty list must round-trip as tags=[], not collapse to
+        # tags=None (which would drop a deliberately cleared tag set).
+        client = Mock()
+        client.flush = Mock(return_value=True)
+        client.get_or_create_dataset = Mock(return_value=Mock(name="ds"))
+        created_experiment = Mock()
+        created_experiment.id = "exp-123"
+        client.create_experiment = Mock(return_value=created_experiment)
+        client._rest_client = Mock()
+
+        experiment_data = ExperimentData(
+            experiment={
+                "id": "exp-123",
+                "name": "test-experiment",
+                "dataset_name": "test-dataset",
+                "type": "regular",
+                "tags": [],
+            },
+            items=[],
+        )
+
+        with patch("opik.cli.imports.experiment.id_helpers_module") as mock_id_helpers:
+            mock_id_helpers.generate_id = Mock(return_value="generated-id")
+
+            recreate_experiment(
+                client,
+                experiment_data,
+                "test-project",
+                {},  # trace_id_map
+                {},  # dataset_item_id_map
+                dry_run=False,
+                debug=False,
+            )
+
+        client.create_experiment.assert_called_once()
+        assert client.create_experiment.call_args.kwargs["tags"] == []

sdks/python/tests/unit/cli/test_trace_filter.py

@@ -0,0 +1,557 @@
+"""Unit tests for matches_trace_filter in exports/trace_filter.py."""
+
+import sys
+from types import ModuleType
+from unittest.mock import MagicMock, patch
+
+
+from opik.cli.exports.trace_filter import matches_trace_filter
+
+
+def _make_oql_module(mock_oql: MagicMock) -> ModuleType:
+    """Return a fake opik.api_objects.opik_query_language module whose
+    OpikQueryLanguage.for_traces() returns *mock_oql*."""
+    mod = ModuleType("opik.api_objects.opik_query_language")
+    cls = MagicMock()
+    cls.for_traces.return_value = mock_oql
+    mod.OpikQueryLanguage = cls  # type: ignore[attr-defined]
+    return mod
+
+
+class TestMatchesTraceFilter:
+    def test_matches_trace_filter__date_time_gte__matches_trace_after_cutoff(self):
+        """A trace whose created_at is after the cutoff must pass the filter."""
+        trace = {
+            "created_at": "2024-06-01T12:00:00+00:00",
+        }
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "created_at",
+                "operator": ">=",
+                "value": "2024-01-01T00:00:00+00:00",
+                "type": "date_time",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'created_at >= "2024-01-01T00:00:00Z"')
+
+        assert result is True
+
+    def test_matches_trace_filter__date_time_gte__excludes_trace_before_cutoff(self):
+        """A trace whose created_at is before the cutoff must be excluded."""
+        trace = {
+            "created_at": "2023-06-01T12:00:00+00:00",
+        }
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "created_at",
+                "operator": ">=",
+                "value": "2024-01-01T00:00:00+00:00",
+                "type": "date_time",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'created_at >= "2024-01-01T00:00:00Z"')
+
+        assert result is False
+
+    def test_matches_trace_filter__invalid_filter__returns_true_with_warning(self):
+        """An unparseable filter string must return True and log a warning."""
+        import opik.cli.exports.trace_filter as trace_filter_module
+
+        trace = {"name": "my-trace"}
+
+        fake_mod = ModuleType("opik.api_objects.opik_query_language")
+        cls = MagicMock()
+        cls.for_traces.side_effect = ValueError("bad filter syntax")
+        fake_mod.OpikQueryLanguage = cls  # type: ignore[attr-defined]
+
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            with patch.object(trace_filter_module.logger, "warning") as mock_warn:
+                result = matches_trace_filter(trace, "THIS IS NOT VALID OQL")
+
+        assert result is True
+        mock_warn.assert_called_once()
+        # The warning message should contain both the filter string and the exception
+        warning_args = mock_warn.call_args
+        assert "bad filter syntax" in str(warning_args) or any(
+            "bad filter syntax" in str(a) for a in warning_args.args
+        )
+
+    def test_matches_trace_filter__string_contains__happyflow(self):
+        """A string 'contains' filter must match when the substring is present."""
+        trace = {"name": "evaluation-run-42"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "contains",
+                "value": "evaluation",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name contains "evaluation"')
+
+        assert result is True
+
+    def test_matches_trace_filter__no_expressions__returns_true(self):
+        """When the filter parses to zero expressions every trace must pass."""
+        trace = {"name": "any-trace"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = []
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, "")
+
+        assert result is True
+
+    # ------------------------------------------------------------------
+    # number type
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__number_gte__matches(self):
+        """A number >= filter must match when the field value satisfies it."""
+        trace = {"usage": {"total_tokens": 500}}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "usage.total_tokens",
+                "operator": ">=",
+                "value": "100",
+                "type": "number",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'usage.total_tokens >= "100"')
+
+        assert result is True
+
+    def test_matches_trace_filter__number_lt__excludes(self):
+        """A number < filter must exclude a trace that doesn't satisfy it."""
+        trace = {"usage": {"total_tokens": 50}}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "usage.total_tokens",
+                "operator": "<",
+                "value": "10",
+                "type": "number",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'usage.total_tokens < "10"')
+
+        assert result is False
+
+    # ------------------------------------------------------------------
+    # is_empty / is_not_empty
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__is_empty__none_field__matches(self):
+        """is_empty must match when the field is absent / None."""
+        trace = {}  # 'error' key is absent
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "error",
+                "operator": "is_empty",
+                "value": "",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, "error is_empty")
+
+        assert result is True
+
+    def test_matches_trace_filter__is_empty__non_empty_field__excludes(self):
+        """is_empty must exclude when the field has a value."""
+        trace = {"error": "something went wrong"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "error",
+                "operator": "is_empty",
+                "value": "",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, "error is_empty")
+
+        assert result is False
+
+    def test_matches_trace_filter__is_not_empty__non_empty_field__matches(self):
+        """is_not_empty must match when the field is present and non-empty."""
+        trace = {"name": "my-trace"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "is_not_empty",
+                "value": "",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, "name is_not_empty")
+
+        assert result is True
+
+    def test_matches_trace_filter__is_not_empty__none_field__excludes(self):
+        """is_not_empty must exclude when the field is absent."""
+        trace = {}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "is_not_empty",
+                "value": "",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, "name is_not_empty")
+
+        assert result is False
+
+    # ------------------------------------------------------------------
+    # Additional string operators
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__string_not_contains__excludes(self):
+        """not_contains must exclude when the substring is present."""
+        trace = {"name": "evaluation-run-42"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "not_contains",
+                "value": "evaluation",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name not_contains "evaluation"')
+
+        assert result is False
+
+    def test_matches_trace_filter__string_starts_with__matches(self):
+        """starts_with must match when the field begins with the prefix."""
+        trace = {"name": "evaluation-run-42"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "starts_with",
+                "value": "evaluation",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name starts_with "evaluation"')
+
+        assert result is True
+
+    def test_matches_trace_filter__string_ends_with__matches(self):
+        """ends_with must match when the field ends with the suffix."""
+        trace = {"name": "evaluation-run-42"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "ends_with",
+                "value": "42",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name ends_with "42"')
+
+        assert result is True
+
+    # ------------------------------------------------------------------
+    # Nested / dotted field access
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__dotted_field__matches(self):
+        """Dotted fields like 'usage.total_tokens' should be resolved correctly."""
+        trace = {"usage": {"total_tokens": 200}}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "usage.total_tokens",
+                "operator": "=",
+                "value": "200",
+                "type": "number",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'usage.total_tokens = "200"')
+
+        assert result is True
+
+    def test_matches_trace_filter__key_field_access__matches(self):
+        """Expressions with a non-empty 'key' resolve via trace_dict[field][key]."""
+        trace = {"metadata": {"env": "production"}}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "metadata",
+                "operator": "=",
+                "value": "production",
+                "type": "string",
+                "key": "env",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'metadata.env = "production"')
+
+        assert result is True
+
+    # ------------------------------------------------------------------
+    # None field value → False
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__missing_field__returns_false(self):
+        """When a non-empty-check expression references a missing field the trace must be excluded."""
+        trace = {}  # 'name' is absent
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "=",
+                "value": "something",
+                "type": "string",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name = "something"')
+
+        assert result is False
+
+    # ------------------------------------------------------------------
+    # Multiple expressions (AND semantics)
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__multiple_expressions__all_must_pass(self):
+        """When two expressions are given both must be satisfied (AND logic)."""
+        trace = {
+            "name": "evaluation-run-42",
+            "created_at": "2024-06-01T12:00:00+00:00",
+        }
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "contains",
+                "value": "evaluation",
+                "type": "string",
+                "key": "",
+            },
+            {
+                "field": "created_at",
+                "operator": ">=",
+                "value": "2024-01-01T00:00:00+00:00",
+                "type": "date_time",
+                "key": "",
+            },
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name contains "evaluation"')
+
+        assert result is True
+
+    def test_matches_trace_filter__multiple_expressions__first_fails__returns_false(
+        self,
+    ):
+        """When the first of two expressions fails the trace must be excluded."""
+        trace = {
+            "name": "other-run",
+            "created_at": "2024-06-01T12:00:00+00:00",
+        }
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "name",
+                "operator": "contains",
+                "value": "evaluation",
+                "type": "string",
+                "key": "",
+            },
+            {
+                "field": "created_at",
+                "operator": ">=",
+                "value": "2024-01-01T00:00:00+00:00",
+                "type": "date_time",
+                "key": "",
+            },
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'name contains "evaluation"')
+
+        assert result is False
+
+    # ------------------------------------------------------------------
+    # date_time: naive datetime field value (no tzinfo)
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__date_time_naive__treated_as_utc(self):
+        """A naive datetime field value must be compared as UTC."""
+        from datetime import datetime
+
+        trace = {"created_at": datetime(2024, 6, 1, 12, 0, 0)}  # no tzinfo
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "created_at",
+                "operator": ">=",
+                "value": "2024-01-01T00:00:00+00:00",
+                "type": "date_time",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'created_at >= "2024-01-01T00:00:00Z"')
+
+        assert result is True
+
+    # ------------------------------------------------------------------
+    # date_time: unparseable field value → result=False (not exception)
+    # ------------------------------------------------------------------
+
+    def test_matches_trace_filter__date_time_unparseable_value__returns_false(self):
+        """An unparseable date_time field value must exclude the trace, not raise."""
+        trace = {"created_at": "not-a-date"}
+
+        mock_oql = MagicMock()
+        mock_oql.get_filter_expressions.return_value = [
+            {
+                "field": "created_at",
+                "operator": ">=",
+                "value": "2024-01-01T00:00:00+00:00",
+                "type": "date_time",
+                "key": "",
+            }
+        ]
+
+        fake_mod = _make_oql_module(mock_oql)
+        with patch.dict(
+            sys.modules, {"opik.api_objects.opik_query_language": fake_mod}
+        ):
+            result = matches_trace_filter(trace, 'created_at >= "2024-01-01T00:00:00Z"')
+
+        assert result is False

sdks/python/tests/unit/configurator/mcp/test_detection.py

@@ -0,0 +1,137 @@
+import httpx
+
+from opik.configurator.mcp import detection
+
+BASE_URL = "https://dev.comet.com/"
+API_URL = "https://dev.comet.com/opik/api/"
+WELL_KNOWN_URL = "https://dev.comet.com/.well-known/oauth-authorization-server/opik"
+MCP_URL = "https://dev.comet.com/opik/api/v1/mcp"
+
+_VALID_METADATA = {
+    "issuer": "https://dev.comet.com/opik",
+    "authorization_endpoint": "https://dev.comet.com/opik/oauth/authorize",
+    "token_endpoint": "https://dev.comet.com/opik/oauth/token",
+}
+
+
+class _Response:
+    def __init__(self, status_code, json_data=None, json_error=False):
+        self.status_code = status_code
+        self._json_data = json_data
+        self._json_error = json_error
+
+    def json(self):
+        if self._json_error:
+            raise ValueError("Expecting value")
+        return self._json_data
+
+
+def _patch_client(monkeypatch, *, response=None, error=None):
+    """Wire ``httpx_client.get`` to a fake client, capturing the probe URL and the
+    client kwargs (so tests can assert the TLS flag is forwarded)."""
+    captured = {}
+
+    class _Client:
+        def __enter__(self):
+            return self
+
+        def __exit__(self, *exc):
+            return False
+
+        def get(self, url, timeout=None):
+            captured["url"] = url
+            captured["timeout"] = timeout
+            if error is not None:
+                raise error
+            return response
+
+    def fake_get(**kwargs):
+        captured["client_kwargs"] = kwargs
+        return _Client()
+
+    monkeypatch.setattr(detection.httpx_client, "get", fake_get)
+    return captured
+
+
+def test_detect__valid_oauth_metadata__returns_mcp_url(monkeypatch):
+    captured = _patch_client(
+        monkeypatch, response=_Response(200, json_data=_VALID_METADATA)
+    )
+
+    result = detection.detect_hosted_mcp_server(
+        base_url=BASE_URL, api_url=API_URL, check_tls_certificate=True
+    )
+
+    assert result == MCP_URL
+    assert captured["url"] == WELL_KNOWN_URL
+
+
+def test_detect__forwards_tls_flag_without_rereading_config(monkeypatch):
+    captured = _patch_client(
+        monkeypatch, response=_Response(200, json_data=_VALID_METADATA)
+    )
+
+    detection.detect_hosted_mcp_server(
+        base_url=BASE_URL, api_url=API_URL, check_tls_certificate=False
+    )
+
+    assert captured["client_kwargs"]["check_tls_certificate"] is False
+
+
+def test_detect__html_catch_all_200__returns_none(monkeypatch):
+    # The frontend SPA answers any unknown path with HTML and a 200 status.
+    _patch_client(monkeypatch, response=_Response(200, json_error=True))
+
+    result = detection.detect_hosted_mcp_server(
+        base_url=BASE_URL, api_url=API_URL, check_tls_certificate=True
+    )
+
+    assert result is None
+
+
+def test_detect__not_found__returns_none(monkeypatch):
+    _patch_client(
+        monkeypatch, response=_Response(404, json_data={"error": "not_found"})
+    )
+
+    assert (
+        detection.detect_hosted_mcp_server(
+            base_url=BASE_URL, api_url=API_URL, check_tls_certificate=True
+        )
+        is None
+    )
+
+
+def test_detect__json_without_authorization_endpoint__returns_none(monkeypatch):
+    _patch_client(
+        monkeypatch, response=_Response(200, json_data={"unrelated": "value"})
+    )
+
+    assert (
+        detection.detect_hosted_mcp_server(
+            base_url=BASE_URL, api_url=API_URL, check_tls_certificate=True
+        )
+        is None
+    )
+
+
+def test_detect__non_dict_json__returns_none(monkeypatch):
+    _patch_client(monkeypatch, response=_Response(200, json_data=["not", "a", "dict"]))
+
+    assert (
+        detection.detect_hosted_mcp_server(
+            base_url=BASE_URL, api_url=API_URL, check_tls_certificate=True
+        )
+        is None
+    )
+
+
+def test_detect__network_error__returns_none(monkeypatch):
+    _patch_client(monkeypatch, error=httpx.ConnectError("unreachable"))
+
+    assert (
+        detection.detect_hosted_mcp_server(
+            base_url=BASE_URL, api_url=API_URL, check_tls_certificate=True
+        )
+        is None
+    )

sdks/python/tests/unit/configurator/mcp/test_env.py

@@ -0,0 +1,59 @@
+from opik.configurator.mcp import env
+
+
+def test_build_mcp_env__cloud__api_key_and_workspace_only():
+    result = env.build_mcp_env(
+        api_key="some-key",
+        workspace="my-workspace",
+        base_url="https://www.comet.com/",
+        api_url="https://www.comet.com/opik/api/",
+        use_local=False,
+        self_hosted_comet=False,
+    )
+    assert result == {
+        "OPIK_API_KEY": "some-key",
+        "COMET_WORKSPACE": "my-workspace",
+    }
+
+
+def test_build_mcp_env__self_hosted_comet__url_override_is_base_url():
+    result = env.build_mcp_env(
+        api_key="some-key",
+        workspace="my-workspace",
+        base_url="https://opik.acme.com/",
+        api_url="https://opik.acme.com/opik/api/",
+        use_local=False,
+        self_hosted_comet=True,
+    )
+    assert result == {
+        "OPIK_API_KEY": "some-key",
+        "COMET_WORKSPACE": "my-workspace",
+        "COMET_URL_OVERRIDE": "https://opik.acme.com",
+    }
+
+
+def test_build_mcp_env__local__opik_url_is_api_url_without_api_key():
+    result = env.build_mcp_env(
+        api_key=None,
+        workspace="default",
+        base_url="http://localhost:5173/",
+        api_url="http://localhost:5173/api/",
+        use_local=True,
+        self_hosted_comet=False,
+    )
+    assert result == {
+        "COMET_WORKSPACE": "default",
+        "OPIK_URL": "http://localhost:5173/api/",
+    }
+
+
+def test_build_mcp_env__workspace_passed_verbatim__not_lowercased():
+    result = env.build_mcp_env(
+        api_key="some-key",
+        workspace="Mixed-Case-WS",
+        base_url="https://www.comet.com/",
+        api_url="https://www.comet.com/opik/api/",
+        use_local=False,
+        self_hosted_comet=False,
+    )
+    assert result["COMET_WORKSPACE"] == "Mixed-Case-WS"

sdks/python/tests/unit/configurator/mcp/test_install.py

@@ -0,0 +1,360 @@
+import pathlib
+import subprocess
+from unittest import mock
+
+import pytest
+
+from opik.configurator.mcp import install, spec, targets
+
+
+@pytest.fixture(autouse=True)
+def prefetch_run(monkeypatch):
+    """Stub the pre-fetch subprocess so tests never shell out to uv."""
+    run_mock = mock.Mock(
+        return_value=subprocess.CompletedProcess([], 0, stdout="", stderr="")
+    )
+    monkeypatch.setattr(install.subprocess, "run", run_mock)
+    return run_mock
+
+
+@pytest.fixture(autouse=True)
+def no_hosted_mcp(monkeypatch):
+    """Default every test to the uvx fallback: no hosted MCP server detected.
+
+    Tests that exercise the remote path override this within the test body.
+    """
+    monkeypatch.setattr(
+        install.mcp_detection, "detect_hosted_mcp_server", lambda **kwargs: None
+    )
+
+
+def _make_args(**overrides):
+    args = dict(
+        api_key="some-key",
+        workspace="ws",
+        base_url="https://www.comet.com/",
+        api_url="https://www.comet.com/opik/api/",
+        use_local=False,
+        self_hosted_comet=False,
+        check_tls_certificate=True,
+        force_local_server=False,
+    )
+    args.update(overrides)
+    return args
+
+
+def _target(key, detected, install_fn):
+    return targets.HostTarget(
+        key=key,
+        display_name=key,
+        config_path=lambda: pathlib.Path("/dev/null"),
+        top_level_key="mcpServers",
+        is_detected=lambda: detected,
+        install=install_fn,
+    )
+
+
+def test_setup_mcp_server__uvx_missing__does_not_install(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: None)
+    install_spy = mock.Mock()
+    monkeypatch.setattr(
+        targets, "HOST_TARGETS", [_target("claude-code", True, install_spy)]
+    )
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_not_called()
+
+
+def test_setup_mcp_server__no_host_detected__does_not_install(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock()
+    monkeypatch.setattr(
+        targets, "HOST_TARGETS", [_target("cursor", False, install_spy)]
+    )
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_not_called()
+
+
+def test_setup_mcp_server__no_host__manual_config_redacts_api_key(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    monkeypatch.setattr(
+        targets, "HOST_TARGETS", [_target("cursor", False, mock.Mock())]
+    )
+    logger_spy = mock.Mock()
+    monkeypatch.setattr(install, "LOGGER", logger_spy)
+
+    install.setup_mcp_server(**_make_args())
+
+    logged = " ".join(str(call) for call in logger_spy.info.call_args_list)
+    assert "some-key" not in logged
+    assert "***REDACTED***" in logged
+
+
+def test_setup_mcp_server__single_host_selected__installs(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")  # confirm the host
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_called_once()
+    spec = install_spy.call_args.args[0]
+    assert spec.command == "/usr/bin/uvx"
+    assert spec.args == ["opik-mcp"]
+    assert spec.env["OPIK_API_KEY"] == "some-key"
+
+
+def test_setup_mcp_server__menu_lists_detected_hosts(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    monkeypatch.setattr(
+        targets,
+        "HOST_TARGETS",
+        [
+            _target("Claude Code", True, mock.Mock()),
+            _target("Cursor", True, mock.Mock()),
+            _target("VS Code Copilot", False, mock.Mock()),
+        ],
+    )
+    prompts = []
+
+    def fake_input(message):
+        prompts.append(message)
+        return "4"  # Skip (2 hosts -> 1,2 hosts, 3 all, 4 skip)
+
+    monkeypatch.setattr("builtins.input", fake_input)
+
+    install.setup_mcp_server(**_make_args())
+
+    assert "Claude Code" in prompts[0]
+    assert "Cursor" in prompts[0]
+    assert "All of the above" in prompts[0]
+    assert "VS Code Copilot" not in prompts[0]
+
+
+def test_setup_mcp_server__select_all__installs_every_detected_host(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    claude_spy = mock.Mock(return_value=targets.InstallResult("Claude", True, "Added"))
+    cursor_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(
+        targets,
+        "HOST_TARGETS",
+        [_target("Claude Code", True, claude_spy), _target("Cursor", True, cursor_spy)],
+    )
+    monkeypatch.setattr("builtins.input", lambda message: "3")  # All of the above
+
+    install.setup_mcp_server(**_make_args())
+
+    claude_spy.assert_called_once()
+    cursor_spy.assert_called_once()
+
+
+def test_setup_mcp_server__comma_separated_selection__installs_each(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    claude_spy = mock.Mock(return_value=targets.InstallResult("Claude", True, "Added"))
+    cursor_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    vscode_spy = mock.Mock(return_value=targets.InstallResult("VS Code", True, "Added"))
+    monkeypatch.setattr(
+        targets,
+        "HOST_TARGETS",
+        [
+            _target("Claude Code", True, claude_spy),
+            _target("Cursor", True, cursor_spy),
+            _target("VS Code Copilot", True, vscode_spy),
+        ],
+    )
+    monkeypatch.setattr("builtins.input", lambda message: "1,3")  # Claude + VS Code
+
+    install.setup_mcp_server(**_make_args())
+
+    claude_spy.assert_called_once()
+    vscode_spy.assert_called_once()
+    cursor_spy.assert_not_called()
+
+
+def test_setup_mcp_server__invalid_menu_choice_then_valid__retries(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    claude_spy = mock.Mock(return_value=targets.InstallResult("Claude", True, "Added"))
+    cursor_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(
+        targets,
+        "HOST_TARGETS",
+        [_target("Claude Code", True, claude_spy), _target("Cursor", True, cursor_spy)],
+    )
+    # invalid (non-digit), out-of-range, then a valid single choice
+    monkeypatch.setattr("builtins.input", mock.Mock(side_effect=["x", "99", "2"]))
+
+    install.setup_mcp_server(**_make_args())
+
+    cursor_spy.assert_called_once()
+    claude_spy.assert_not_called()
+
+
+def test_setup_mcp_server__select_subset__installs_only_chosen(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    claude_spy = mock.Mock(return_value=targets.InstallResult("Claude", True, "Added"))
+    cursor_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(
+        targets,
+        "HOST_TARGETS",
+        [_target("Claude Code", True, claude_spy), _target("Cursor", True, cursor_spy)],
+    )
+    monkeypatch.setattr("builtins.input", lambda message: "2")  # only Cursor
+
+    install.setup_mcp_server(**_make_args())
+
+    claude_spy.assert_not_called()
+    cursor_spy.assert_called_once()
+
+
+def test_setup_mcp_server__user_skips__does_not_install(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock()
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "n")  # decline
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_not_called()
+
+
+def test_setup_mcp_server__prefetches_opik_mcp_before_install(
+    monkeypatch, prefetch_run
+):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args())
+
+    commands = [call.args[0] for call in prefetch_run.call_args_list]
+    assert any(cmd[1:] == ["tool", "install", "opik-mcp"] for cmd in commands)
+    install_spy.assert_called_once()
+
+
+def test_setup_mcp_server__prefetch_failure__is_non_fatal(monkeypatch, prefetch_run):
+    prefetch_run.return_value = subprocess.CompletedProcess(
+        [], 1, stdout="", stderr="network unreachable"
+    )
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_called_once()
+
+
+def test_setup_mcp_server__prefetch_raises_oserror__is_non_fatal(
+    monkeypatch, prefetch_run
+):
+    prefetch_run.side_effect = OSError("uv vanished mid-flight")
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args())  # must not raise
+
+    install_spy.assert_called_once()
+
+
+def test_setup_mcp_server__skip__does_not_prefetch(monkeypatch, prefetch_run):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, mock.Mock())])
+    monkeypatch.setattr("builtins.input", lambda message: "n")  # decline
+
+    install.setup_mcp_server(**_make_args())
+
+    prefetch_run.assert_not_called()
+
+
+def test_setup_mcp_server__install_failure__is_reported(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    install_spy = mock.Mock(
+        return_value=targets.InstallResult("Cursor", False, "could not write config")
+    )
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")  # confirm the host
+    logger_spy = mock.Mock()
+    monkeypatch.setattr(install, "LOGGER", logger_spy)
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_called_once()
+    logged = " ".join(str(call) for call in logger_spy.warning.call_args_list)
+    assert "could not write config" in logged
+
+
+def test_setup_mcp_server__hosted_detected__installs_remote_spec(monkeypatch):
+    monkeypatch.setattr(
+        install.mcp_detection,
+        "detect_hosted_mcp_server",
+        lambda **kwargs: "https://dev.comet.com/opik/api/v1/mcp",
+    )
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_called_once()
+    server_spec = install_spy.call_args.args[0]
+    assert isinstance(server_spec, spec.RemoteServerSpec)
+    assert server_spec.url == "https://dev.comet.com/opik/api/v1/mcp"
+
+
+def test_setup_mcp_server__hosted_detected__does_not_prefetch(
+    monkeypatch, prefetch_run
+):
+    monkeypatch.setattr(
+        install.mcp_detection,
+        "detect_hosted_mcp_server",
+        lambda **kwargs: "https://dev.comet.com/opik/api/v1/mcp",
+    )
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args())
+
+    prefetch_run.assert_not_called()
+    install_spy.assert_called_once()
+
+
+def test_setup_mcp_server__force_local__skips_probe_and_installs_uvx(monkeypatch):
+    monkeypatch.setattr(install.shutil, "which", lambda name: "/usr/bin/uvx")
+    detect_spy = mock.Mock(return_value="https://dev.comet.com/opik/api/v1/mcp")
+    monkeypatch.setattr(install.mcp_detection, "detect_hosted_mcp_server", detect_spy)
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args(force_local_server=True))
+
+    detect_spy.assert_not_called()
+    install_spy.assert_called_once()
+    server_spec = install_spy.call_args.args[0]
+    assert isinstance(server_spec, spec.StdioServerSpec)
+
+
+def test_setup_mcp_server__hosted_detected__no_uvx_still_installs(monkeypatch):
+    """The remote path has no `uvx` prerequisite, unlike the local fallback."""
+    monkeypatch.setattr(install.shutil, "which", lambda name: None)
+    monkeypatch.setattr(
+        install.mcp_detection,
+        "detect_hosted_mcp_server",
+        lambda **kwargs: "https://dev.comet.com/opik/api/v1/mcp",
+    )
+    install_spy = mock.Mock(return_value=targets.InstallResult("Cursor", True, "Added"))
+    monkeypatch.setattr(targets, "HOST_TARGETS", [_target("cursor", True, install_spy)])
+    monkeypatch.setattr("builtins.input", lambda message: "y")
+
+    install.setup_mcp_server(**_make_args())
+
+    install_spy.assert_called_once()

sdks/python/tests/unit/configurator/mcp/test_json_config.py

@@ -0,0 +1,117 @@
+import json
+
+import pytest
+
+from opik.configurator.mcp import json_config
+
+SERVER_BLOCK = {
+    "type": "stdio",
+    "command": "/usr/bin/uvx",
+    "args": ["opik-mcp"],
+    "env": {"OPIK_API_KEY": "some-key"},
+}
+
+
+def test_merge_server_into_json_file__new_file__creates_with_entry(tmp_path):
+    config_path = tmp_path / "nested" / "mcp.json"
+
+    was_new = json_config.merge_server_into_json_file(
+        config_path=config_path,
+        top_level_key="mcpServers",
+        server_name="opik-mcp",
+        server_block=SERVER_BLOCK,
+    )
+
+    assert was_new is True
+    written = json.loads(config_path.read_text(encoding="utf-8"))
+    assert written == {"mcpServers": {"opik-mcp": SERVER_BLOCK}}
+
+
+def test_merge_server_into_json_file__existing_servers__preserved(tmp_path):
+    config_path = tmp_path / "mcp.json"
+    config_path.write_text(
+        json.dumps(
+            {
+                "someOtherTopLevelKey": {"keep": "me"},
+                "mcpServers": {"other-server": {"command": "foo"}},
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    was_new = json_config.merge_server_into_json_file(
+        config_path=config_path,
+        top_level_key="mcpServers",
+        server_name="opik-mcp",
+        server_block=SERVER_BLOCK,
+    )
+
+    assert was_new is True
+    written = json.loads(config_path.read_text(encoding="utf-8"))
+    assert written["someOtherTopLevelKey"] == {"keep": "me"}
+    assert written["mcpServers"]["other-server"] == {"command": "foo"}
+    assert written["mcpServers"]["opik-mcp"] == SERVER_BLOCK
+
+
+def test_merge_server_into_json_file__rerun__updates_in_place(tmp_path):
+    config_path = tmp_path / "mcp.json"
+    json_config.merge_server_into_json_file(
+        config_path=config_path,
+        top_level_key="mcpServers",
+        server_name="opik-mcp",
+        server_block={"command": "old"},
+    )
+
+    was_new = json_config.merge_server_into_json_file(
+        config_path=config_path,
+        top_level_key="mcpServers",
+        server_name="opik-mcp",
+        server_block=SERVER_BLOCK,
+    )
+
+    assert was_new is False
+    written = json.loads(config_path.read_text(encoding="utf-8"))
+    assert list(written["mcpServers"].keys()) == ["opik-mcp"]
+    assert written["mcpServers"]["opik-mcp"] == SERVER_BLOCK
+
+
+def test_merge_server_into_json_file__invalid_json__raises(tmp_path):
+    config_path = tmp_path / "mcp.json"
+    config_path.write_text("{ // a JSONC comment\n}", encoding="utf-8")
+
+    with pytest.raises(json.JSONDecodeError):
+        json_config.merge_server_into_json_file(
+            config_path=config_path,
+            top_level_key="servers",
+            server_name="opik-mcp",
+            server_block=SERVER_BLOCK,
+        )
+
+
+def test_merge_server_into_json_file__non_object_root__raises_value_error(tmp_path):
+    config_path = tmp_path / "mcp.json"
+    config_path.write_text('["not", "an", "object"]', encoding="utf-8")
+
+    with pytest.raises(ValueError):
+        json_config.merge_server_into_json_file(
+            config_path=config_path,
+            top_level_key="mcpServers",
+            server_name="opik-mcp",
+            server_block=SERVER_BLOCK,
+        )
+
+
+def test_merge_server_into_json_file__empty_existing_file__treated_as_empty(tmp_path):
+    config_path = tmp_path / "mcp.json"
+    config_path.write_text("", encoding="utf-8")
+
+    was_new = json_config.merge_server_into_json_file(
+        config_path=config_path,
+        top_level_key="mcpServers",
+        server_name="opik-mcp",
+        server_block=SERVER_BLOCK,
+    )
+
+    assert was_new is True
+    written = json.loads(config_path.read_text(encoding="utf-8"))
+    assert written["mcpServers"]["opik-mcp"] == SERVER_BLOCK