Fix eager Celery task handling

Only close the session that the signal handlers opened; an eager task inside a
caller's `Session()` was closing it out from under them. Also read per-call
heartbeat options back out of the headers, which is where they stay when
`before_task_publish` doesn't run.

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

Author: snopoke

taskbadger/celery.py

@@ -27,6 +27,9 @@
 TB_TASK_ID = f"{KWARG_PREFIX}task_id"
 TB_HEARTBEAT_INTERVAL = f"{KWARG_PREFIX}heartbeat_interval"
 TB_STALE_TIMEOUT = f"{KWARG_PREFIX}stale_timeout"
+# Marks a request whose signal handlers opened the Task Badger session, so that
+# only they close it again.
+TB_OWNS_SESSION = f"{KWARG_PREFIX}owns_session"
 
 log = logging.getLogger("taskbadger")
 
@@ -268,15 +271,18 @@ def _maybe_create_task(signal_sender):
     if not manual_track and not auto_track:
         return
 
-    enter_session()
+    enter_session(signal_sender)
 
     delivery_info = getattr(signal_sender.request, "delivery_info", None) or {}
     queue = delivery_info.get("routing_key")
     external_id = signal_sender.request.id
     create_kwargs = {"status": StatusEnum.PENDING, "data": data, "queue": queue, "external_id": external_id}
+    # `before_task_publish` never ran for these, so per-call options are still
+    # sitting in the headers rather than resolved into the message.
+    header_kwargs = headers.get(TB_KWARGS_ARG) or {}
     heartbeat_interval, stale_timeout = resolve_heartbeat_options(
-        getattr(signal_sender, TB_HEARTBEAT_INTERVAL, None),
-        getattr(signal_sender, TB_STALE_TIMEOUT, None),
+        header_kwargs.get("heartbeat_interval", getattr(signal_sender, TB_HEARTBEAT_INTERVAL, None)),
+        header_kwargs.get("stale_timeout", getattr(signal_sender, TB_STALE_TIMEOUT, None)),
         celery_system,
     )
     if stale_timeout is not None:
@@ -341,7 +347,7 @@ def _update_task(signal_sender, status, einfo=None):
         # ignore tasks that have already been set to a terminal state (probably in the task body)
         return
 
-    enter_session()
+    enter_session(signal_sender)
 
     data = None
     if einfo:
@@ -370,22 +376,32 @@ def _start_heartbeat(signal_sender):
     heartbeat.start(task_id, interval)
 
 
-def enter_session():
+def enter_session(signal_sender):
     if not Badger.is_configured():
         return
     session = Badger.current.session()
-    if not session.client:
-        session.__enter__()
+    if session.client:
+        # Already open, e.g. an eager task running inside the caller's
+        # `taskbadger.Session()`. Not ours, so don't close it either.
+        return
+    session.__enter__()
+    signal_sender.request.update({TB_OWNS_SESSION: True})
 
 
 def exit_session(signal_sender):
+    request = signal_sender.request
     # not `request.headers`: eager tasks carry the ID on the request itself
-    task_id = _get_taskbadger_task_id(signal_sender.request)
+    # (the task id is checked first so that untracked tasks never touch `Badger`)
+    task_id = _get_taskbadger_task_id(request)
     if not task_id or not Badger.is_configured():
         return
 
     task_cache.unset(task_id)
 
+    if not request.get(TB_OWNS_SESSION):
+        return
+    request.update({TB_OWNS_SESSION: False})
+
     session = Badger.current.session()
     if session.client:
         session.__exit__()

tests/test_celery_heartbeat.py

@@ -9,6 +9,7 @@
 
 import pytest
 
+import taskbadger
 from taskbadger import StatusEnum
 from taskbadger.celery import Task
 from taskbadger.mug import Badger, Settings
@@ -154,6 +155,58 @@ def add_eager(a, b):
     heartbeat.stop.assert_called_once_with(tb_task.id)
 
 
+@pytest.mark.usefixtures("_bind_settings_with_heartbeat")
+def test_heartbeat_for_eager_task_from_apply_async(celery_session_app, heartbeat):
+    """``before_task_publish`` doesn't run for eager tasks, so the per-call
+    options have to be read back out of the headers."""
+
+    @celery_session_app.task(base=Task)
+    def add_eager_options(a, b):
+        return a + b
+
+    celery_session_app.conf.task_always_eager = True
+    try:
+        with (
+            mock.patch("taskbadger.celery.create_task_safe") as create,
+            mock.patch("taskbadger.celery.update_task_safe"),
+            mock.patch("taskbadger.sdk.get_task"),
+        ):
+            tb_task = task_for_test()
+            create.return_value = tb_task
+            result = add_eager_options.apply_async((2, 2), taskbadger_heartbeat_interval=3)
+            assert result.get(timeout=10, propagate=True) == 4
+    finally:
+        celery_session_app.conf.task_always_eager = False
+
+    assert create.call_args.kwargs["stale_timeout"] == 6
+    heartbeat.start.assert_called_once_with(tb_task.id, 3)
+
+
+@pytest.mark.usefixtures("_bind_settings_with_heartbeat")
+def test_eager_task_leaves_an_open_session_alone(celery_session_app, heartbeat):
+    """An eager task must not close a session opened by its caller."""
+
+    @celery_session_app.task(base=Task)
+    def add_in_session(a, b):
+        return a + b
+
+    celery_session_app.conf.task_always_eager = True
+    try:
+        with (
+            mock.patch("taskbadger.celery.create_task_safe", return_value=task_for_test()),
+            mock.patch("taskbadger.celery.update_task_safe", return_value=task_for_test()),
+            mock.patch("taskbadger.sdk.get_task"),
+            mock.patch("taskbadger.mug.AuthenticatedClient"),
+        ):
+            with taskbadger.Session():
+                session = Badger.current.session()
+                assert add_in_session.delay(2, 2).get(timeout=10, propagate=True) == 4
+                assert session.client is not None, "the caller's session was closed by the task"
+            assert session.client is None
+    finally:
+        celery_session_app.conf.task_always_eager = False
+
+
 @pytest.mark.usefixtures("_bind_settings")
 def test_no_heartbeat_by_default(celery_session_app, celery_session_worker, heartbeat):
     @celery_session_app.task(base=Task)