Don't nest chain links and callbacks

Celery dispatches the next chain link and any `link` callbacks from inside
`trace_task` after the body returns but before `task_postrun`, so the task they
follow still looks current and they were being recorded as its children.
They're successors, not subtasks.

Nothing on the wire distinguishes them from a publish the body made itself, so
this matches against the running task's own `request.chain`/`request.callbacks`
rather than inspecting the stack.

Retries still nest — a retry republishes the same task, which isn't among its
own successors.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: snopoke

integration_tests/tasks.py

@@ -39,6 +39,16 @@ def spawns_child(self, x, y):
     return spawns_grandchild.delay(x, y).id
 
 
+@shared_task(bind=True, base=taskbadger.celery.Task)
+def chain_head(self):
+    return self.taskbadger_task_id
+
+
+@shared_task(bind=True, base=taskbadger.celery.Task)
+def chain_tail(self, head_tb_id):
+    return {"head_tb_id": head_tb_id, "own_tb_id": self.taskbadger_task_id}
+
+
 @shared_task(bind=True, base=taskbadger.celery.Task, taskbadger_heartbeat_interval=HEARTBEAT_INTERVAL)
 def slow_add(self, x, y):
     """Runs long enough to go stale without a heartbeat, and never updates itself."""

integration_tests/test_celery.py

@@ -3,11 +3,12 @@
 import time
 
 import pytest
+from celery import chain
 
 import taskbadger
 from taskbadger import StatusEnum
 
-from .tasks import HEARTBEAT_INTERVAL, add, add_auto_track, slow_add, spawns_child
+from .tasks import HEARTBEAT_INTERVAL, add, add_auto_track, chain_head, chain_tail, slow_add, spawns_child
 
 
 @pytest.fixture(autouse=True)
@@ -69,6 +70,18 @@ def test_celery_grandchild_is_flattened_onto_the_root(celery_session_app, celery
     assert grandchild.parent != child["own_tb_id"]
 
 
+def test_celery_chain_links_are_not_nested(celery_session_app, celery_session_worker):
+    """Celery dispatches the next chain link from inside the previous task's run,
+    so it would otherwise be nested under it. Links are successors, not subtasks.
+    """
+    ids = chain(chain_head.s(), chain_tail.s()).apply_async().get(timeout=20, propagate=True)
+
+    assert ids["head_tb_id"], "the first link should be tracked"
+    assert ids["own_tb_id"], "the second link should be tracked"
+    assert not taskbadger.get_task(ids["own_tb_id"]).parent
+    assert taskbadger.list_tasks(parent=ids["head_tb_id"]).results == []
+
+
 def test_celery_heartbeat(celery_session_app, celery_session_worker):
     """The worker pings the task while it runs, so it doesn't go stale."""
     a, b = random.randint(1, 1000), random.randint(1, 1000)

taskbadger/celery.py

@@ -141,6 +141,41 @@ def taskbadger_task(self):
         return task
 
 
+def _signature_names(signature):
+    """The task names ``signature`` will publish, expanding groups.
+
+    Signatures reach us either as ``Signature`` objects or as the dicts they
+    serialize to, depending on where in Celery they came from.
+    """
+    if isinstance(signature, dict):
+        name = signature.get("task")
+        nested = (signature.get("kwargs") or {}).get("tasks")
+    else:
+        name = getattr(signature, "task", None)
+        nested = getattr(signature, "tasks", None)
+    if nested:
+        for inner in nested:
+            yield from _signature_names(inner)
+    elif name:
+        yield name
+
+
+def _is_workflow_successor(task_name):
+    """Whether the task being published is the running task's next chain link or
+    one of its callbacks, rather than something its body chose to enqueue.
+
+    Celery dispatches both from inside ``trace_task`` after the task body
+    returns but before ``task_postrun``, so the running task still looks
+    current. They're the task's successors rather than its subtasks, so they
+    shouldn't be nested under it.
+    """
+    request = getattr(celery.current_task, "request", None)
+    if request is None:
+        return False
+    successors = list(request.chain or []) + list(request.callbacks or [])
+    return any(name == task_name for sig in successors for name in _signature_names(sig))
+
+
 @before_task_publish.connect
 def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     routing_key = kwargs.get("routing_key")
@@ -174,7 +209,7 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs):
     # `before_task_publish` fires in the process doing the publishing, so if that
     # is itself a tracked task this nests the new task under it.
     enclosing_task = parent_id()
-    if enclosing_task:
+    if enclosing_task and not _is_workflow_successor(sender):
         kwargs.setdefault("parent", enclosing_task)
     name = kwargs.pop("name", headers["task"])
 

tests/test_parents.py

@@ -263,6 +263,51 @@ def test_celery_publish_has_no_parent_outside_a_task():
     assert "parent" not in create.call_args.kwargs
 
 
+def _running_task(chain=None, callbacks=None):
+    """Pretend a task with these successors is mid-run."""
+    request = mock.Mock(chain=chain, callbacks=callbacks)
+    return mock.patch("celery.current_task", mock.Mock(request=request))
+
+
+@pytest.mark.usefixtures("_bind_settings")
+@pytest.mark.parametrize(
+    "successors",
+    [
+        {"chain": [{"task": "child.task"}]},
+        {"callbacks": [{"task": "child.task"}]},
+        # a group of callbacks arrives as one signature wrapping the real tasks
+        {"callbacks": [{"task": "celery.group", "kwargs": {"tasks": [{"task": "child.task"}]}}]},
+    ],
+    ids=["chain", "callback", "callback_group"],
+)
+def test_celery_workflow_successors_are_not_nested(successors):
+    """Chain links and callbacks are dispatched while the task that precedes
+    them still looks current, but they follow it rather than belong to it."""
+    with mock.patch("taskbadger.celery.create_task_safe") as create, _running_task(**successors):
+        create.return_value = task_for_test()
+        token = enter_task("root_id")
+        try:
+            _publish()
+        finally:
+            exit_task(token)
+
+    assert "parent" not in create.call_args.kwargs
+
+
+@pytest.mark.usefixtures("_bind_settings")
+def test_celery_publish_from_a_task_body_still_nests_while_a_chain_is_pending():
+    """Only the successor itself is exempt — a task the body enqueues is not."""
+    with mock.patch("taskbadger.celery.create_task_safe") as create, _running_task(chain=[{"task": "next.task"}]):
+        create.return_value = task_for_test()
+        token = enter_task("root_id")
+        try:
+            _publish()
+        finally:
+            exit_task(token)
+
+    assert create.call_args.kwargs["parent"] == "root_id"
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_celery_publish_explicit_parent_wins():
     with mock.patch("taskbadger.celery.create_task_safe") as create: