feat: FIT-710: Enable FSM state fields in serializers and models

label_studio/data_manager/managers.py

@@ -7,6 +7,7 @@
 from typing import ClassVar
 
 import ujson as json
+from core.current_request import CurrentContext
 from core.feature_flags import flag_set
 from core.utils.db import fast_first
 from data_manager.prepare_params import ConjunctionEnum
@@ -30,6 +31,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 +490,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 +711,39 @@ 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.
+
+    Both feature flags must be enabled and have a current user in context:
+    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/UI
+    """
+
+    user = CurrentContext.get_user()
+    # If no user in context, return unmodified queryset
+    if user is None:
+        return queryset
+
+    # Only annotate if both FSM feature flags are enabled
+    # This prevents unnecessary DB queries when state shouldn't be visible
+    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)
+    ):
+        return queryset
+
+    # Use the mixin's annotate_fsm_state() method which creates 'current_state' annotation
+    queryset = queryset.annotate_fsm_state()
+
+    # Alias 'current_state' to 'state' for Data Manager column compatibility
+    return queryset.annotate(state=F('current_state'))
+
+
 settings.DATA_MANAGER_ANNOTATIONS_MAP = {
     'avg_lead_time': annotate_avg_lead_time,
     'completed_at': annotate_completed_at,
@@ -712,6 +756,7 @@ def dummy(queryset):
     'file_upload': file_upload,
     'draft_exists': annotate_draft_exists,
     'storage_filename': annotate_storage_filename,
+    'state': annotate_state,
 }
 
 
@@ -724,6 +769,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 +812,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 +849,31 @@ 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):
+        """
+        Convenience method to return queryset with FSM state annotated.
+
+        Example:
+            tasks = Task.objects.with_state().filter(project=project)
+            for task in tasks:
+                print(task.current_state)  # No N+1 queries!
+        """
+        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,14 @@ class PredictionsDMFieldSerializer(serializers.SerializerMethodField):
 
 
 class DataManagerTaskSerializer(TaskSerializer):
+    """
+    Data Manager Task Serializer with FSM state support.
+
+    Note: The 'state' field will be populated from the queryset annotation
+    if present, preventing N+1 queries. Ensure your queryset uses .with_state()
+    or .annotate_fsm_state() for optimal performance.
+    """
+
     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 +465,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 +482,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):