Fb fit 671 thread context

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,85 @@
 _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 is_async_job(cls) -> bool:
+        return cls.get_request() is None
+
+    @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,15 +160,25 @@ 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:
+        # Async execution with Redis - wrap job for context management
+        context_data = _capture_context()
+
+        if context_data:
+            # Only wrap if we have context to preserve
+            kwargs['_context_data'] = context_data
+
         try:
             args_info = _truncate_args_for_logging(args, kwargs)
             logger.info(f'Start async job {job.__name__} on queue {queue_name} with {args_info}.')
@@ -140,6 +188,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,
@@ -149,6 +198,8 @@ def start_job_async_or_sync(job, *args, in_seconds=0, **kwargs):
         )
         return job
     else:
+        # Sync execution - context is already available from request thread
+        # No need to wrap or modify the job function
         on_failure = kwargs.pop('on_failure', None)
         try:
             return job(*args, **kwargs)

label_studio/data_manager/actions/basic.py

@@ -97,7 +97,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/apps.py

@@ -11,3 +11,42 @@ class FsmConfig(AppConfig):
     default_auto_field = 'django.db.models.UUIDField'
     name = 'fsm'
     verbose_name = 'Label Studio FSM'
+
+    def ready(self):
+        """Initialize FSM integration when the app is ready"""
+        # Check if this is Community edition - only register signals for Community
+        from django.conf import settings
+
+        version_edition = getattr(settings, 'VERSION_EDITION', 'Community')
+
+        if version_edition != 'Community':
+            logger.info(f'Label Studio FSM: Skipping initialization for {version_edition} edition')
+            return
+
+        # Additional safety check: if LSE FSM apps are present, don't initialize
+        try:
+            from django.apps import apps
+
+            if apps.is_installed('lse_fsm'):
+                logger.info('Label Studio FSM: LSE FSM detected, skipping LSO FSM initialization to avoid conflicts')
+                return
+        except Exception as e:
+            # Log but continue - this shouldn't prevent initialization in pure LSO environments
+            logger.debug(f'Label Studio FSM: Error checking for LSE apps: {e}')
+
+        logger.info('Label Studio FSM app ready, initializing core integrations for Community edition')
+
+        # Import signal handlers to register them (only for Community edition)
+        try:
+            from . import signals  # noqa: F401
+
+            logger.info('Label Studio FSM signal handlers registered successfully')
+        except Exception as e:
+            logger.error(
+                'FSM: Failed to register Label Studio FSM signal handlers',
+                extra={
+                    'event': 'fsm.signal_registration_error',
+                    'error': str(e),
+                },
+                exc_info=True,
+            )