Add tracking for Celery canvas primitives (map/starmap/chunks)

Previously, tasks spawned via Celery canvas primitives like map, starmap,
and chunks were not tracked because these execute as built-in celery.map
or celery.starmap tasks, which were filtered out by the before_task_publish
signal handler.

This change adds a _maybe_create_task function called during task_prerun
that:
- Creates TaskBadger tasks for canvas primitives using the inner task's
  name (not "celery.map")
- Includes metadata about canvas_type and item_count
- Optionally includes celery_task_items when record_task_args is enabled
- Carefully orders checks to avoid thread-local Badger state pollution
- Respects existing tracking intent via taskbadger_track header

For example, when using task.chunks(items, 2).group(), each chunk now
creates a TaskBadger task named after the user's task with data showing
it was executed via celery.starmap and how many items it processed.

Author: snopoke

taskbadger/celery.py

@@ -212,8 +212,88 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
         ctask.update_state(task_id=headers["id"], state="PENDING", meta=meta)
 
 
+def _maybe_create_task(signal_sender):
+    """Create a TaskBadger task if one doesn't exist yet.
+
+    This handles cases where before_task_publish didn't fire or was skipped:
+    - Eager mode (before_task_publish doesn't fire)
+    - Canvas primitives like map/starmap/chunks (fire for celery.* tasks)
+    """
+    # Check if task was already created FIRST (before accessing Badger)
+    # This avoids initializing thread-local Badger state for tasks like celery.ping
+    task_id = _get_taskbadger_task_id(signal_sender.request)
+    if task_id:
+        return
+
+    task_name = signal_sender.name
+
+    # Skip built-in celery tasks that we don't track (like celery.ping)
+    # Only handle celery.map and celery.starmap specially
+    if task_name.startswith("celery.") and task_name not in ("celery.map", "celery.starmap"):
+        return
+
+    # For non-canvas tasks, only create if there was an explicit intent to track
+    # (indicated by taskbadger_track header). This prevents creating tasks when
+    # Badger wasn't configured at publish time but has stale config in worker.
+    headers = signal_sender.request.headers or {}
+    is_canvas_task = task_name in ("celery.map", "celery.starmap")
+    if not is_canvas_task and not headers.get("taskbadger_track"):
+        return
+
+    # NOW it's safe to check Badger configuration
+    if not Badger.is_configured():
+        return
+
+    celery_system = Badger.current.settings.get_system_by_id("celery")
+    data = None
+    inner_task = None
+
+    # Handle celery.map and celery.starmap - extract the inner task name
+    if task_name in ("celery.map", "celery.starmap"):
+        inner_task_info = signal_sender.request.kwargs.get("task")
+        if inner_task_info:
+            # inner_task_info can be a dict (serialized signature) or a Signature object
+            if isinstance(inner_task_info, dict):
+                task_name = inner_task_info.get("task", task_name)
+            elif hasattr(inner_task_info, "name"):
+                task_name = inner_task_info.name
+            # Get the actual task class to check if it uses Task base
+            inner_task = celery.current_app.tasks.get(task_name)
+        items = signal_sender.request.kwargs.get("it", [])
+        # Convert to list if needed for counting and potential recording
+        items_list = list(items) if not isinstance(items, (list, tuple)) else items
+        item_count = len(items_list)
+        data = {"canvas_type": signal_sender.name, "item_count": item_count}
+
+        # Include task items if record_task_args is enabled
+        if celery_system and celery_system.record_task_args:
+            try:
+                _, _, value = serialization.dumps({"items": items_list}, serializer="json")
+                items_data = json.loads(value)
+                data["celery_task_items"] = items_data["items"]
+            except Exception:
+                log.warning("Error serializing canvas items for task '%s'", task_name)
+
+    # Check if we should track this task
+    auto_track = celery_system and celery_system.track_task(task_name)
+    # Check if the task (or inner task for map/starmap) uses our Task base class
+    task_to_check = inner_task if inner_task else signal_sender
+    manual_track = isinstance(task_to_check, Task)
+    if not manual_track and not auto_track:
+        return
+
+    enter_session()
+
+    task = create_task_safe(task_name, status=StatusEnum.PENDING, data=data)
+    if task:
+        # Store the task ID in the request so _update_task can find it
+        signal_sender.request.update({TB_TASK_ID: task.id})
+        safe_get_task.cache.set((task.id,), task)
+
+
 @task_prerun.connect
 def task_prerun_handler(sender=None, **kwargs):
+    _maybe_create_task(sender)
     _update_task(sender, StatusEnum.PROCESSING)
 
 

tests/test_celery.py

@@ -393,32 +393,100 @@ def task_signature(self, a):
 
 @pytest.mark.usefixtures("_bind_settings")
 def test_task_map(celery_session_worker):
-    """Tasks executed in a map or starmap are not executed as tasks"""
+    """Tasks executed via map canvas primitive should be tracked.
+
+    Note: The individual function calls within map are not separate Celery tasks,
+    so we track the map operation itself with the inner task's name.
+    """
 
     @celery.shared_task(bind=True, base=Task)
-    def task_map(self, a):
-        assert self.taskbadger_task is None
-        assert Badger.current.session().client is None
+    def task_map_fn(self, a):
         return a * 2
 
     celery_session_worker.reload()
 
-    task_map = task_map.map(list(range(5)))
+    map_canvas = task_map_fn.map(list(range(5)))
 
     with (
         mock.patch("taskbadger.celery.create_task_safe") as create,
         mock.patch("taskbadger.celery.update_task_safe") as update,
-        mock.patch("taskbadger.celery.get_task") as get_task,
+        mock.patch("taskbadger.celery.get_task"),
     ):
-        result = task_map.delay()
+        tb_task = task_for_test()
+        create.return_value = tb_task
+        result = map_canvas.delay()
         assert result.get(timeout=10, propagate=True) == [0, 2, 4, 6, 8]
 
-    assert create.call_count == 0
-    assert get_task.call_count == 0
-    assert update.call_count == 0
+    # Map operation should create one TaskBadger task
+    assert create.call_count == 1
+    # Verify the task name uses the inner task's name, not "celery.map"
+    call_args = create.call_args
+    assert "task_map_fn" in call_args[0][0]
+    assert update.call_count == 2  # PROCESSING and SUCCESS
     assert Badger.current.session().client is None
 
 
+@pytest.mark.usefixtures("_bind_settings")
+def test_task_starmap(celery_session_worker):
+    """Tasks executed via starmap canvas primitive should be tracked."""
+
+    @celery.shared_task(bind=True, base=Task)
+    def task_starmap_fn(self, a, b):
+        return a + b
+
+    celery_session_worker.reload()
+
+    starmap_canvas = task_starmap_fn.starmap([(1, 2), (3, 4), (5, 6)])
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe") as update,
+        mock.patch("taskbadger.celery.get_task"),
+    ):
+        tb_task = task_for_test()
+        create.return_value = tb_task
+        result = starmap_canvas.delay()
+        assert result.get(timeout=10, propagate=True) == [3, 7, 11]
+
+    # Starmap operation should create one TaskBadger task
+    assert create.call_count == 1
+    # Verify the task name uses the inner task's name
+    call_args = create.call_args
+    assert "task_starmap_fn" in call_args[0][0]
+    assert update.call_count == 2  # PROCESSING and SUCCESS
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_task_chunks(celery_session_worker):
+    """Tasks executed via chunks canvas primitive should be tracked."""
+
+    @celery.shared_task(bind=True, base=Task)
+    def task_chunks_fn(self, a):
+        return a * 2
+
+    celery_session_worker.reload()
+
+    # chunks creates multiple starmap tasks
+    chunks_canvas = task_chunks_fn.chunks([(x,) for x in range(6)], 2).group()
+
+    with (
+        mock.patch("taskbadger.celery.create_task_safe") as create,
+        mock.patch("taskbadger.celery.update_task_safe") as update,
+        mock.patch("taskbadger.celery.get_task"),
+    ):
+        tb_task = task_for_test()
+        create.return_value = tb_task
+        result = chunks_canvas.delay()
+        assert result.get(timeout=10, propagate=True) == [[0, 2], [4, 6], [8, 10]]
+
+    # Each chunk should create a TaskBadger task (3 chunks of 2)
+    assert create.call_count == 3
+    # Verify the task names use the inner task's name
+    for call in create.call_args_list:
+        assert "task_chunks_fn" in call[0][0]
+    assert update.call_count == 6  # 3 tasks * 2 updates each
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_celery_task_already_in_terminal_state(celery_session_worker):
     @celery.shared_task(bind=True, base=Task)