Implement preserve_on_clone flag system to fix condition cloning behavior (#1133)

* Initial plan

* Update test to validate current condition cloning issue

Co-authored-by: jdkent <[email protected]>

* Add comprehensive tests for condition cloning behavior and initial fix attempt

Co-authored-by: jdkent <[email protected]>

* Implement preserve_on_clone flag for condition cloning fix

Co-authored-by: jdkent <[email protected]>

* Add comprehensive integration test demonstrating successful condition cloning fix

Co-authored-by: jdkent <[email protected]>

* Refactor tests and fix style issues: separate schema and API tests

Co-authored-by: jdkent <[email protected]>

* Refactor condition cloning tests for clarity

Renamed tests for clarity and removed unused test.

* Remove unused import of pytest in test_condition_cloning

Removed unused import statement for pytest.

* Changes before error encountered

Co-authored-by: jdkent <[email protected]>

* fix tests to compare the right id ("external_id")

* remove redundant test files

* fix style

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: jdkent <[email protected]>
Co-authored-by: James Kent <[email protected]>

Author: Copilot

store/backend/neurostore/schemas/data.py

@@ -20,6 +20,7 @@
 # nested: serialize nested objects (true or false)
 # flat: do not display any relationships
 # info: only display info fields
+# preserve_on_clone: field metadata to preserve original values when cloning
 
 
 class BooleanOrString(fields.Field):
@@ -67,13 +68,36 @@ def _modify_schema(self):
         schema = self.schema
         schema.context = self.context
         if self.context.get("clone"):
-            id_fields = [
-                field
-                for field, f_obj in schema._declared_fields.items()
-                if f_obj.metadata.get("id_field")
-            ]
-            for f in id_fields:
-                schema.exclude.add(f)
+            # Check if this schema has preserve_on_clone set for any id field
+            has_preserve_on_clone = any(
+                f_obj.metadata.get("id_field")
+                and f_obj.metadata.get("preserve_on_clone")
+                for f_obj in schema._declared_fields.values()
+            )
+
+            if has_preserve_on_clone:
+                # For schemas with preserve_on_clone,
+                # include ONLY id fields with preserve_on_clone=True
+                preserve_fields = [
+                    field
+                    for field, f_obj in schema._declared_fields.items()
+                    if f_obj.metadata.get("id_field")
+                    and f_obj.metadata.get("preserve_on_clone")
+                ]
+                if preserve_fields:
+                    schema.only = schema.set_class(preserve_fields)
+                    schema.exclude = schema.set_class()
+            else:
+                # Normal cloning behavior: exclude id fields without preserve_on_clone
+                id_fields = [
+                    field
+                    for field, f_obj in schema._declared_fields.items()
+                    if f_obj.metadata.get("id_field")
+                    and not f_obj.metadata.get("preserve_on_clone")
+                ]
+                for f in id_fields:
+                    schema.exclude.add(f)
+
         if self.context.get("info"):
             info_fields = [
                 field
@@ -132,15 +156,38 @@ class BaseSchema(Schema):
     def __init__(self, *args, **kwargs):
         exclude = kwargs.get("exclude") or self.opts.exclude
         context = kwargs.get("context", {})
-        # if cloning and not only id, exclude id fields
-        if context.get("clone") and kwargs.get("only") != ("id",):
-            id_fields = [
-                field
-                for field, f_obj in self._declared_fields.items()
-                if f_obj.metadata.get("id_field")
-            ]
-            for f in id_fields:
-                exclude += (f,)
+        only = kwargs.get("only")
+
+        # if cloning and not only id, exclude id fields (unless preserve_on_clone is True)
+        if context.get("clone") and only != ("id",):
+            # Check if this schema has preserve_on_clone set for any id field
+            has_preserve_on_clone = any(
+                f_obj.metadata.get("id_field")
+                and f_obj.metadata.get("preserve_on_clone")
+                for f_obj in self._declared_fields.values()
+            )
+
+            if has_preserve_on_clone:
+                # For schemas with preserve_on_clone,
+                # include ONLY id fields with preserve_on_clone=True
+                preserve_fields = [
+                    field
+                    for field, f_obj in self._declared_fields.items()
+                    if f_obj.metadata.get("id_field")
+                    and f_obj.metadata.get("preserve_on_clone")
+                ]
+                if preserve_fields and only is None:
+                    kwargs["only"] = tuple(preserve_fields)
+            else:
+                # Normal cloning behavior: exclude id fields
+                id_fields = [
+                    field
+                    for field, f_obj in self._declared_fields.items()
+                    if f_obj.metadata.get("id_field")
+                ]
+                for f in id_fields:
+                    exclude += (f,)
+
         if context.get("flat"):
             relationships = [
                 field
@@ -182,6 +229,11 @@ class Meta:
         additional = ("name", "description")
         allow_none = ("name", "description")
 
+    # Override the id field to preserve it during cloning
+    id = fields.String(
+        metadata={"info_field": True, "id_field": True, "preserve_on_clone": True}
+    )
+
 
 class EntitySchema(BaseDataSchema):
     analysis_id = fields.String(data_key="analysis", metadata={"id_field": True})