feat: FIT-710: Enable FSM state fields in serializers and models (#8775)

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

Author: bmartel

label_studio/data_manager/managers.py

@@ -30,6 +30,7 @@
 )
 from django.db.models.fields.json import KeyTextTransform
 from django.db.models.functions import Cast, Coalesce, Concat
+from fsm.queryset_mixins import FSMStateQuerySetMixin
 from pydantic import BaseModel
 
 from label_studio.core.utils.common import load_func
@@ -488,7 +489,16 @@ def apply_filters(queryset, filters, project, request):
     return queryset
 
 
-class TaskQuerySet(models.QuerySet):
+class TaskQuerySet(FSMStateQuerySetMixin, models.QuerySet):
+    """
+    QuerySet for Task model with FSM state annotation support.
+
+    Extends Django's QuerySet with:
+    - FSM state annotation (via FSMStateQuerySetMixin)
+    - Data Manager filters and ordering
+    - Selected items handling
+    """
+
     def prepared(self, prepare_params=None):
         """Apply filters, ordering and selected items to queryset
 
@@ -700,6 +710,29 @@ def dummy(queryset):
     return queryset
 
 
+def annotate_state(queryset):
+    """
+    Annotate queryset with FSM state as 'state' field.
+
+    Uses FSMStateQuerySetMixin.annotate_fsm_state() to efficiently annotate
+    the current state without causing N+1 queries. Aliases 'current_state' to
+    'state' to match the Data Manager column name.
+
+    Note: Feature flag checks and user context validation are handled by
+    annotate_fsm_state() itself, so no additional checks are needed here.
+    """
+    # Use the mixin's annotate_fsm_state() method which creates 'current_state' annotation
+    # (includes feature flag and user context checks)
+    queryset = queryset.annotate_fsm_state()
+
+    # Alias 'current_state' to 'state' for Data Manager column compatibility
+    # Only add the alias if current_state was actually added (feature flags enabled)
+    if 'current_state' in queryset.query.annotations:
+        return queryset.annotate(state=F('current_state'))
+
+    return queryset
+
+
 settings.DATA_MANAGER_ANNOTATIONS_MAP = {
     'avg_lead_time': annotate_avg_lead_time,
     'completed_at': annotate_completed_at,
@@ -712,6 +745,7 @@ def dummy(queryset):
     'file_upload': file_upload,
     'draft_exists': annotate_draft_exists,
     'storage_filename': annotate_storage_filename,
+    'state': annotate_state,
 }
 
 
@@ -724,6 +758,19 @@ def update_annotation_map(obj):
 
 
 class PreparedTaskManager(models.Manager):
+    """
+    Manager for Task model with Data Manager annotations.
+
+    Provides:
+    - Advanced query annotations for Data Manager
+    - Filter and ordering support
+    - FSM state annotation support (via TaskQuerySet)
+
+    Note: Overrides the base get_queryset() to return TaskQuerySet. Also has
+    a custom get_queryset(fields_for_evaluation, prepare_params, ...) method
+    for Data Manager-specific functionality.
+    """
+
     @staticmethod
     def annotate_queryset(
         queryset, fields_for_evaluation=None, all_fields=False, excluded_fields_for_evaluation=None, request=None
@@ -754,13 +801,23 @@ def get_queryset(
         self, fields_for_evaluation=None, prepare_params=None, all_fields=False, excluded_fields_for_evaluation=None
     ):
         """
+        Get queryset with optional Data Manager annotations and filters.
+
+        When called without parameters (Django internal use), returns TaskQuerySet.
+        When called with parameters (Data Manager use), returns annotated and filtered queryset.
+
         :param fields_for_evaluation: list of annotated fields in task
         :param prepare_params: filters, ordering, selected items
         :param all_fields: evaluate all fields for task
         :param excluded_fields_for_evaluation: list of fields to exclude even when all_fields=True
         :param request: request for user extraction
         :return: task queryset with annotated fields
         """
+        # If called without parameters, return base TaskQuerySet (for Django internal use)
+        if prepare_params is None:
+            return TaskQuerySet(self.model, using=self._db)
+
+        # Otherwise, use Data Manager filtering and annotation
         queryset = self.only_filtered(prepare_params=prepare_params)
         # Expose view data to annotation functions for column-specific configuration
         queryset.view_data = getattr(prepare_params, 'data', None)
@@ -781,5 +838,24 @@ def only_filtered(self, prepare_params=None):
 
 
 class TaskManager(models.Manager):
+    """
+    Default manager for Task model.
+
+    Provides:
+    - User-scoped filtering
+    - Custom QuerySet with FSM state support
+
+    Note: Overrides get_queryset() to return TaskQuerySet, which includes
+    FSMStateQuerySetMixin for state annotation support.
+    """
+
+    def get_queryset(self):
+        """Return TaskQuerySet which includes FSM state annotation support"""
+        return TaskQuerySet(self.model, using=self._db)
+
     def for_user(self, user):
-        return self.filter(project__organization=user.active_organization)
+        return self.get_queryset().filter(project__organization=user.active_organization)
+
+    def with_state(self):
+        """Return queryset with FSM state annotated."""
+        return self.get_queryset().annotate_fsm_state()

label_studio/data_manager/serializers.py

@@ -3,10 +3,13 @@
 import os
 
 import ujson as json
+from core.current_request import CurrentContext
+from core.feature_flags import flag_set
 from data_manager.models import Filter, FilterGroup, View
 from django.conf import settings
 from django.db import transaction
 from drf_spectacular.utils import extend_schema_field
+from fsm.serializer_fields import FSMStateField
 from projects.models import Project
 from rest_framework import serializers
 from tasks.models import Task
@@ -434,6 +437,8 @@ class PredictionsDMFieldSerializer(serializers.SerializerMethodField):
 
 
 class DataManagerTaskSerializer(TaskSerializer):
+    """Data Manager Task Serializer with FSM state support."""
+
     predictions = PredictionsDMFieldSerializer(required=False, read_only=True)
     annotations = AnnotationsDMFieldSerializer(required=False, many=True, default=[], read_only=True)
     drafts = AnnotationDraftDMFieldSerializer(required=False, read_only=True)
@@ -454,6 +459,7 @@ class DataManagerTaskSerializer(TaskSerializer):
     avg_lead_time = serializers.FloatField(required=False)
     draft_exists = serializers.BooleanField(required=False)
     updated_by = UpdatedByDMFieldSerializer(required=False, read_only=True)
+    state = FSMStateField(read_only=True)  # FSM state - automatically uses annotation if present
 
     CHAR_LIMITS = 500
 
@@ -470,6 +476,13 @@ def to_representation(self, obj):
             ret.pop('annotations', None)
         if not self.context.get('predictions'):
             ret.pop('predictions', None)
+        # Remove state field if feature flags are disabled
+        user = CurrentContext.get_user()
+        if not (
+            flag_set('fflag_feat_fit_568_finite_state_management', user=user)
+            and flag_set('fflag_feat_fit_710_fsm_state_fields', user=user)
+        ):
+            ret.pop('state', None)
         return ret
 
     def _pretty_results(self, task, field, unique=False):

label_studio/fsm/queryset_mixins.py

@@ -0,0 +1,113 @@
+"""
+FSM QuerySet Mixins for annotating entities with their current state.
+
+Provides reusable Django QuerySet mixins that efficiently annotate entities
+with their current FSM state using optimized subqueries to prevent N+1 queries.
+
+Usage:
+    class TaskQuerySet(FSMStateQuerySetMixin, models.QuerySet):
+        pass
+
+    class TaskManager(models.Manager):
+        def get_queryset(self):
+            return TaskQuerySet(self.model, using=self._db).annotate_fsm_state()
+
+Note:
+    All state annotation functionality is guarded by TWO feature flags:
+    1. 'fflag_feat_fit_568_finite_state_management' - Controls FSM background calculations
+    2. 'fflag_feat_fit_710_fsm_state_fields' - Controls state field display in APIs
+
+    When disabled, no annotation is performed and there is zero performance impact.
+"""
+
+import logging
+
+from core.current_request import CurrentContext
+from core.feature_flags import flag_set
+from django.db.models import OuterRef, Subquery
+from fsm.registry import get_state_model
+
+logger = logging.getLogger(__name__)
+
+
+class FSMStateQuerySetMixin:
+    """
+    Mixin for Django QuerySets to efficiently annotate FSM state.
+
+    Provides the `annotate_fsm_state()` method that adds a `current_state`
+    annotation to the queryset using an optimized subquery.
+
+    This approach:
+    - Prevents N+1 queries by using a single JOIN/subquery
+    - Handles missing states gracefully (returns None)
+    - Uses UUID7 natural ordering for optimal performance
+    - Works with any FSM entity that has a registered state model
+
+    Example:
+        # In your model manager
+        class TaskManager(models.Manager):
+            def get_queryset(self):
+                return TaskQuerySet(self.model, using=self._db)
+
+            def with_state(self):
+                return self.get_queryset().annotate_fsm_state()
+
+        # Usage
+        tasks = Task.objects.with_state().filter(project=project)
+        for task in tasks:
+            print(f"Task {task.id}: {task.current_state}")  # No additional queries!
+    """
+
+    def annotate_fsm_state(self):
+        """
+        Annotate the queryset with the current FSM state.
+
+        Adds a `current_state` field to each object containing the current
+        state string value. This is done using an efficient subquery that
+        leverages UUID7 natural ordering.
+
+        Returns:
+            QuerySet: The annotated queryset with `current_state` field
+
+        Note:
+            - If FSM feature flag is disabled, returns queryset unchanged (zero impact)
+            - If no state exists for an entity, `current_state` will be None
+            - The state is read-only and should not be modified directly
+        """
+        # Check feature flag directly (works for both core and enterprise)
+        # Using flag_set directly instead of is_fsm_enabled to work in enterprise context
+        user = CurrentContext.get_user()
+        if not (
+            flag_set('fflag_feat_fit_568_finite_state_management', user=user)
+            and flag_set('fflag_feat_fit_710_fsm_state_fields', user=user)
+        ):
+            logger.debug('FSM feature flag disabled, skipping state annotation')
+            return self
+
+        # Get the entity name from the model
+        entity_name = self.model._meta.model_name
+
+        # Get the state model for this entity
+        state_model = get_state_model(entity_name)
+
+        if not state_model:
+            # No state model registered, return queryset as-is
+            logger.debug(f'No state model registered for {entity_name}, skipping annotation')
+            return self
+
+        # Get the foreign key field name on the state model
+        # e.g., 'task_id' for TaskState
+        entity_field_name = state_model._get_entity_field_name()
+        fk_field = f'{entity_field_name}_id'
+
+        # Create subquery to get current state using UUID7 natural ordering
+        # This is extremely efficient because:
+        # 1. UUID7 provides natural time ordering (latest = highest ID)
+        # 2. We only fetch the state column, not the entire record
+        # 3. Django optimizes this into a single JOIN or lateral subquery
+        current_state_subquery = Subquery(
+            state_model.objects.filter(**{fk_field: OuterRef('pk')}).order_by('-id').values('state')[:1]
+        )
+
+        # Annotate the queryset with the current state
+        return self.annotate(current_state=current_state_subquery)