chore: FIT-671: Add state transitions to the current application flows (#8534)

Co-authored-by: Marcel Canu <[email protected]>
Co-authored-by: bmartel <[email protected]>
Co-authored-by: Wesley Lima <[email protected]>
Co-authored-by: Andrew <[email protected]>
Co-authored-by: hlomzik <[email protected]>
Co-authored-by: Nick Skriabin <[email protected]>
Co-authored-by: robot-ci-heartex <[email protected]>
Co-authored-by: nick-skriabin <[email protected]>

Author: 6 people

label_studio/core/current_request.py

@@ -1,4 +1,5 @@
 from threading import local
+from typing import Any
 
 from django.core.signals import request_finished
 from django.dispatch import receiver
@@ -7,18 +8,104 @@
 _thread_locals = local()
 
 
+class CurrentContext:
+    @classmethod
+    def set(cls, key: str, value: Any, shared: bool = True) -> None:
+        if not hasattr(_thread_locals, 'data'):
+            _thread_locals.data = {}
+        if not hasattr(_thread_locals, 'job_data'):
+            _thread_locals.job_data = {}
+
+        if shared:
+            _thread_locals.job_data[key] = value
+        else:
+            _thread_locals.data[key] = value
+
+    @classmethod
+    def get(cls, key: str, default=None):
+        return getattr(_thread_locals, 'job_data', {}).get(key, getattr(_thread_locals, 'data', {}).get(key, default))
+
+    @classmethod
+    def set_request(cls, request):
+        _thread_locals.request = request
+        if request.user:
+            cls.set_user(request.user)
+
+    @classmethod
+    def get_organization_id(cls):
+        return cls.get('organization_id')
+
+    @classmethod
+    def set_organization_id(cls, organization_id: int):
+        cls.set('organization_id', organization_id)
+
+    @classmethod
+    def get_user(cls):
+        return cls.get('user')
+
+    @classmethod
+    def set_user(cls, user):
+        cls.set('user', user)
+        if getattr(user, 'active_organization_id', None):
+            cls.set_organization_id(user.active_organization_id)
+
+    @classmethod
+    def set_fsm_disabled(cls, disabled: bool):
+        """
+        Temporarily disable/enable FSM for the current thread.
+
+        This is useful for test cleanup and bulk operations where FSM state
+        tracking is not needed and would cause performance issues.
+
+        Args:
+            disabled: True to disable FSM, False to enable it
+        """
+        cls.set('fsm_disabled', disabled)
+
+    @classmethod
+    def is_fsm_disabled(cls) -> bool:
+        """
+        Check if FSM is disabled for the current thread.
+
+        Returns:
+            True if FSM is disabled, False otherwise
+        """
+        return cls.get('fsm_disabled', False)
+
+    @classmethod
+    def get_job_data(cls) -> dict:
+        """
+        This data will be shared to jobs spawned by the current thread.
+        """
+        return getattr(_thread_locals, 'job_data', {})
+
+    @classmethod
+    def clear(cls) -> None:
+        if hasattr(_thread_locals, 'data'):
+            delattr(_thread_locals, 'data')
+
+        if hasattr(_thread_locals, 'job_data'):
+            delattr(_thread_locals, 'job_data')
+
+        if hasattr(_thread_locals, 'request'):
+            del _thread_locals.request
+
+    @classmethod
+    def get_request(cls):
+        return getattr(_thread_locals, 'request', None)
+
+
 def get_current_request():
     """returns the request object for this thread"""
-    result = getattr(_thread_locals, 'request', None)
+    result = CurrentContext.get_request()
     return result
 
 
 class ThreadLocalMiddleware(CommonMiddleware):
     def process_request(self, request):
-        _thread_locals.request = request
+        CurrentContext.set_request(request)
 
 
 @receiver(request_finished)
 def clean_request(sender, **kwargs):
-    if hasattr(_thread_locals, 'request'):
-        del _thread_locals.request
+    CurrentContext.clear()

label_studio/core/redis.py

@@ -4,9 +4,11 @@
 import sys
 from datetime import timedelta
 from functools import partial
+from typing import Any
 
 import django_rq
 import redis
+from core.current_request import CurrentContext
 from django.conf import settings
 from django_rq import get_connection
 from rq.command import send_stop_job_command
@@ -80,6 +82,40 @@ def redis_connected():
     return redis_healthcheck()
 
 
+def _is_serializable(value: Any) -> bool:
+    """Check if a value can be serialized for job context."""
+    return isinstance(value, (str, int, float, bool, list, dict, type(None)))
+
+
+def _capture_context() -> dict:
+    """
+    Capture the current context for passing to a job.
+    Returns a dictionary of context data that can be serialized.
+    """
+    context_data = {}
+
+    # Get user information
+    if user := CurrentContext.get_user():
+        context_data['user_id'] = user.id
+
+    # Get organization if set separately
+    if org_id := CurrentContext.get_organization_id():
+        context_data['organization_id'] = org_id
+
+    # If organization_id is not set, try to get it from the user, this ensures that we have an organization_id for the job
+    # And it prefers the original requesting user's organization_id over the current active organization_id of the user which could change during async jobs
+    if not org_id and user and hasattr(user, 'active_organization_id') and user.active_organization_id:
+        context_data['organization_id'] = user.active_organization_id
+
+    # Get any custom context values (exclude non-serializable objects)
+    job_data = CurrentContext.get_job_data()
+    for key, value in job_data.items():
+        if key not in ['user', 'request'] and _is_serializable(value):
+            context_data[key] = value
+
+    return context_data
+
+
 def redis_get(key):
     if not redis_healthcheck():
         return
@@ -112,7 +148,9 @@ def redis_delete(key):
 
 def start_job_async_or_sync(job, *args, in_seconds=0, **kwargs):
     """
-    Start job async with redis or sync if redis is not connected
+    Start job async with redis or sync if redis is not connected.
+    Automatically preserves context for async jobs and clears it after completion.
+
     :param job: Job function
     :param args: Function arguments
     :param in_seconds: Job will be delayed for in_seconds
@@ -122,28 +160,29 @@ def start_job_async_or_sync(job, *args, in_seconds=0, **kwargs):
 
     redis = redis_connected() and kwargs.get('redis', True)
     queue_name = kwargs.get('queue_name', 'default')
+
     if 'queue_name' in kwargs:
         del kwargs['queue_name']
     if 'redis' in kwargs:
         del kwargs['redis']
+
     job_timeout = None
     if 'job_timeout' in kwargs:
         job_timeout = kwargs['job_timeout']
         del kwargs['job_timeout']
+
     if redis:
-        # Auto-capture request_id from thread local and pass it via job meta
+        # Async execution with Redis - wrap job for context management
         try:
-            from label_studio.core.current_request import _thread_locals
+            context_data = _capture_context()
 
-            request_id = getattr(_thread_locals, 'request_id', None)
-            if request_id:
-                # Store in job meta for worker access
+            if context_data:
                 meta = kwargs.get('meta', {})
-                meta['request_id'] = request_id
+                # Store context data in job meta for worker access
+                meta.update(context_data)
                 kwargs['meta'] = meta
         except Exception:
-            # Fail silently if no request context
-            pass
+            logger.info(f'Failed to capture context for job {job.__name__} on queue {queue_name}')
 
         try:
             args_info = _truncate_args_for_logging(args, kwargs)
@@ -154,6 +193,7 @@ def start_job_async_or_sync(job, *args, in_seconds=0, **kwargs):
         enqueue_method = queue.enqueue
         if in_seconds > 0:
             enqueue_method = partial(queue.enqueue_in, timedelta(seconds=in_seconds))
+
         job = enqueue_method(
             job,
             *args,
@@ -164,8 +204,10 @@ def start_job_async_or_sync(job, *args, in_seconds=0, **kwargs):
         return job
     else:
         on_failure = kwargs.pop('on_failure', None)
+
         try:
-            return job(*args, **kwargs)
+            result = job(*args, **kwargs)
+            return result
         except Exception:
             exc_info = sys.exc_info()
             if on_failure:

label_studio/core/settings/base.py

@@ -216,6 +216,7 @@
     'rest_framework.authtoken',
     'rest_framework_simplejwt.token_blacklist',
     'drf_generators',
+    'fsm',  # MUST be before apps that register FSM transitions (projects, tasks)
     'core',
     'users',
     'organizations',
@@ -232,7 +233,6 @@
     'ml_model_providers',
     'jwt_auth',
     'session_policy',
-    'fsm',
 ]
 
 MIDDLEWARE = [

label_studio/core/tests/test_models.py

@@ -18,6 +18,9 @@ def _assert_delete_and_restore_equal(self, drow, original):
         original_dict.pop('_state')
         original_created_at = original_dict.pop('created_at')
         original_updated_at = original_dict.pop('updated_at')
+        # Pop _original_values - this is an internal FSM field that's recreated on __init__
+        # and shouldn't be compared
+        original_dict.pop('_original_values', None)
         original.delete()
 
         for deserialized_object in serializers.deserialize('json', json.dumps([drow.data])):
@@ -28,6 +31,9 @@ def _assert_delete_and_restore_equal(self, drow, original):
         new_dict.pop('_state')
         new_created_at = new_dict.pop('created_at')
         new_updated_at = new_dict.pop('updated_at')
+        # Pop _original_values - this is an internal FSM field that's recreated on __init__
+        # and shouldn't be compared
+        new_dict.pop('_original_values', None)
 
         assert new_dict == original_dict
         # Datetime loses microsecond precision, so we can't compare them directly

label_studio/data_manager/actions/basic.py

@@ -99,7 +99,9 @@ def delete_tasks_annotations(project, queryset, **kwargs):
         drafts = drafts.filter(user=int(annotator_id))
     project.summary.remove_created_drafts_and_labels(drafts)
 
-    count, _ = annotations.delete()
+    # count before delete to return the number of deleted items, not including cascade deletions
+    count = annotations.count()
+    annotations.delete()
     drafts.delete()  # since task-level annotation drafts will not have been deleted by CASCADE
     emit_webhooks_for_instance(project.organization, project, WebhookAction.ANNOTATIONS_DELETED, annotations_ids)
     request = kwargs['request']

label_studio/fsm/README.md

@@ -56,7 +56,7 @@ class OrderStateChoices(models.TextChoices):
 ### 2. Create State Model
 
 ```python
-from fsm.models import BaseState
+from fsm.state_models import BaseState
 from fsm.registry import register_state_model
 
 @register_state_model('order')