From 6a47fba40d9453f0b63be3233cb16a3135bb8c49 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 13:55:47 +0000 Subject: [PATCH 1/4] feat(celery): mark retries with is_retry Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01WDSbaMzpzXZKif1HbDFaT3 --- taskbadger.yaml | 18 +++++++++ taskbadger/_current_task.py | 29 ++++++++++++--- taskbadger/_integrations.py | 5 +++ taskbadger/celery.py | 37 +++++++++++++++++-- .../internal/models/patched_task_request.py | 11 ++++++ taskbadger/internal/models/task.py | 11 ++++++ taskbadger/internal/models/task_request.py | 11 ++++++ taskbadger/sdk.py | 8 ++++ tests/test_celery.py | 6 +++ tests/test_parents.py | 18 +++++++++ 10 files changed, 146 insertions(+), 8 deletions(-) diff --git a/taskbadger.yaml b/taskbadger.yaml index fde0ad1..affd5bd 100644 --- a/taskbadger.yaml +++ b/taskbadger.yaml @@ -712,6 +712,12 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + is_retry: + type: boolean + description: Set when this task is a retry of its parent task. The first + 3 retries of a parent don't count towards the monthly task allocation. + Ignored unless 'parent' is also set (a retry with no parent is a new task) + and can only be set when the task is created. name: type: string minLength: 1 @@ -823,6 +829,12 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + is_retry: + type: boolean + description: Set when this task is a retry of its parent task. The first + 3 retries of a parent don't count towards the monthly task allocation. + Ignored unless 'parent' is also set (a retry with no parent is a new task) + and can only be set when the task is created. name: type: string description: Name of the task @@ -932,6 +944,12 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + is_retry: + type: boolean + description: Set when this task is a retry of its parent task. The first + 3 retries of a parent don't count towards the monthly task allocation. + Ignored unless 'parent' is also set (a retry with no parent is a new task) + and can only be set when the task is created. name: type: string minLength: 1 diff --git a/taskbadger/_current_task.py b/taskbadger/_current_task.py index 3613bc9..9bf6534 100644 --- a/taskbadger/_current_task.py +++ b/taskbadger/_current_task.py @@ -14,21 +14,22 @@ from contextvars import ContextVar -# (id of the running task, id of its parent or None if it is a root task) -_current: ContextVar[tuple[str, str | None] | None] = ContextVar("taskbadger_current_task", default=None) +# (id of the running task, id of its parent or None if it is a root task, whether it is a retry) +_current: ContextVar[tuple[str, str | None, bool] | None] = ContextVar("taskbadger_current_task", default=None) -def enter_task(task_id: str, parent: str = None): +def enter_task(task_id: str, parent: str = None, is_retry: bool = False): """Mark `task_id` as the task running in this context. Arguments: task_id: The running task. parent: The running task's own parent, if it has one. + is_retry: Whether the running task is itself a retry of its parent. Returns: A token to pass to `exit_task`. """ - return _current.set((task_id, parent)) + return _current.set((task_id, parent, is_retry)) def exit_task(token) -> None: @@ -49,5 +50,23 @@ def parent_id() -> str | None: current = _current.get() if current is None: return None - task_id, parent = current + task_id, parent, _ = current return parent or task_id + + +def retry_parent_id() -> str | None: + """The id to use as `parent` for a retry of the task running right now. + + A retry is nested under the first attempt, so this is the running task for the first attempt + and the first attempt again when the running task is itself a retry. `None` when there is no + suitable first attempt: outside a tracked task, or when the running task is a child of some + other task (a fan-out child), since retries, like children, can't be nested a level deeper. + The caller should then create the retry as an ordinary task. + """ + current = _current.get() + if current is None: + return None + task_id, parent, is_retry = current + if parent is None: + return task_id + return parent if is_retry else None diff --git a/taskbadger/_integrations.py b/taskbadger/_integrations.py index c9d52ff..add4b87 100644 --- a/taskbadger/_integrations.py +++ b/taskbadger/_integrations.py @@ -88,6 +88,11 @@ def parent_of(task) -> str | None: return parent or None +def is_retry_task(task) -> bool: + """Whether ``task`` was created as a retry of its parent (``False`` for ``None``/unset).""" + return getattr(task, "is_retry", False) is True + + def match_task_name(task_name: str, includes, excludes) -> bool: """Return True if ``task_name`` should be tracked under the given rules. diff --git a/taskbadger/celery.py b/taskbadger/celery.py index 26c4e67..62e686f 100644 --- a/taskbadger/celery.py +++ b/taskbadger/celery.py @@ -14,10 +14,17 @@ from kombu import serialization from . import sdk -from ._current_task import enter_task, exit_task, parent_id +from ._current_task import enter_task, exit_task, parent_id, retry_parent_id from ._error_context import capture_error_data, reset_error_context, start_error_context from ._heartbeat import heartbeat -from ._integrations import TERMINAL_STATES, parent_of, resolve_heartbeat_options, safe_get_task, task_cache +from ._integrations import ( + TERMINAL_STATES, + is_retry_task, + parent_of, + resolve_heartbeat_options, + safe_get_task, + task_cache, +) from .internal.models import StatusEnum from .mug import Badger from .safe_sdk import create_task_safe, update_task_safe @@ -199,6 +206,16 @@ def _is_workflow_successor(task_name): return any(name == task_name for sig in successors for name in _signature_names(sig)) +def _is_retry_of_running_task(headers): + """Whether the message being published is a retry of the task that is running. + + ``self.retry()`` republishes with the running task's own id, which separates it from an + unrelated task that happens to carry a ``retries`` header. + """ + request = getattr(celery.current_task, "request", None) + return request is not None and request.id == headers.get("id") + + @before_task_publish.connect def task_publish_handler(sender=None, headers=None, body=None, **kwargs): routing_key = kwargs.get("routing_key") @@ -236,6 +253,13 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs): # is itself a tracked task this nests the new task under it. enclosing_task = parent_id() if enclosing_task and not _is_workflow_successor(sender): + if "parent" not in kwargs and headers.get("retries") and _is_retry_of_running_task(headers): + # A Celery retry republishes the same task. It is marked so the API can leave + # the first few out of the monthly allocation, and nested under the first attempt. + retry_parent = retry_parent_id() + if retry_parent: + kwargs["parent"] = retry_parent + kwargs["is_retry"] = True kwargs.setdefault("parent", enclosing_task) name = kwargs.pop("name", headers["task"]) @@ -372,6 +396,12 @@ def _maybe_create_task(signal_sender): # explicit `None`, which asks for a root task. "parent": header_kwargs["parent"] if "parent" in header_kwargs else parent_id(), } + if "parent" not in header_kwargs and getattr(signal_sender.request, "retries", 0): + # An eager retry runs nested inside the attempt that retried it. + retry_parent = retry_parent_id() + if retry_parent: + create_kwargs["parent"] = retry_parent + create_kwargs["is_retry"] = True heartbeat_interval, stale_timeout = resolve_heartbeat_options( header_kwargs.get("heartbeat_interval", getattr(signal_sender, TB_HEARTBEAT_INTERVAL, None)), header_kwargs.get("stale_timeout", getattr(signal_sender, TB_STALE_TIMEOUT, None)), @@ -407,7 +437,8 @@ def task_prerun_handler(sender=None, **kwargs): # Mark this as the running task so anything it publishes is nested under # it. `_update_task` above has already put it in the cache, so this is a # lookup rather than a fetch. - current_token = enter_task(task_id, parent_of(safe_get_task(task_id))) + running = safe_get_task(task_id) + current_token = enter_task(task_id, parent_of(running), is_retry_task(running)) sender.request.update({TB_CURRENT_TASK_TOKEN: current_token}) diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index 519cf02..fefb4dc 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/taskbadger/internal/models/patched_task_request.py @@ -24,6 +24,9 @@ class PatchedTaskRequest: Attributes: parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + is_retry (bool | Unset): Set when this task is a retry of its parent task. The first 3 retries of a parent don't + count towards the monthly task allocation. Ignored unless 'parent' is also set (a retry with no parent is a new + task) and can only be set when the task is created. name (str | Unset): Name of the task queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with @@ -53,6 +56,7 @@ class PatchedTaskRequest: """ parent: None | str | Unset = UNSET + is_retry: bool | Unset = UNSET name: str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET @@ -77,6 +81,8 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + is_retry = self.is_retry + name = self.name queue = self.queue @@ -140,6 +146,8 @@ def to_dict(self) -> dict[str, Any]: field_dict.update({}) if parent is not UNSET: field_dict["parent"] = parent + if is_retry is not UNSET: + field_dict["is_retry"] = is_retry if name is not UNSET: field_dict["name"] = name if queue is not UNSET: @@ -184,6 +192,8 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + is_retry = d.pop("is_retry", UNSET) + name = d.pop("name", UNSET) queue = d.pop("queue", UNSET) @@ -280,6 +290,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: patched_task_request = cls( parent=parent, + is_retry=is_retry, name=name, queue=queue, external_id=external_id, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index 67a4812..451d64e 100644 --- a/taskbadger/internal/models/task.py +++ b/taskbadger/internal/models/task.py @@ -33,6 +33,9 @@ class Task: public_url (str): parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + is_retry (bool | Unset): Set when this task is a retry of its parent task. The first 3 retries of a parent don't + count towards the monthly task allocation. Ignored unless 'parent' is also set (a retry with no parent is a new + task) and can only be set when the task is created. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -70,6 +73,7 @@ class Task: url: str public_url: str parent: None | str | Unset = UNSET + is_retry: bool | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -112,6 +116,8 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + is_retry = self.is_retry + queue = self.queue external_id = self.external_id @@ -185,6 +191,8 @@ def to_dict(self) -> dict[str, Any]: ) if parent is not UNSET: field_dict["parent"] = parent + if is_retry is not UNSET: + field_dict["is_retry"] = is_retry if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -249,6 +257,8 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + is_retry = d.pop("is_retry", UNSET) + queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -352,6 +362,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: url=url, public_url=public_url, parent=parent, + is_retry=is_retry, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index fe22822..c82800c 100644 --- a/taskbadger/internal/models/task_request.py +++ b/taskbadger/internal/models/task_request.py @@ -25,6 +25,9 @@ class TaskRequest: name (str): Name of the task parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + is_retry (bool | Unset): Set when this task is a retry of its parent task. The first 3 retries of a parent don't + count towards the monthly task allocation. Ignored unless 'parent' is also set (a retry with no parent is a new + task) and can only be set when the task is created. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -54,6 +57,7 @@ class TaskRequest: name: str parent: None | str | Unset = UNSET + is_retry: bool | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -79,6 +83,8 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + is_retry = self.is_retry + queue = self.queue external_id = self.external_id @@ -144,6 +150,8 @@ def to_dict(self) -> dict[str, Any]: ) if parent is not UNSET: field_dict["parent"] = parent + if is_retry is not UNSET: + field_dict["is_retry"] = is_retry if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -187,6 +195,8 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + is_retry = d.pop("is_retry", UNSET) + queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -282,6 +292,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: task_request = cls( name=name, parent=parent, + is_retry=is_retry, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index 4564580..f2905fd 100644 --- a/taskbadger/sdk.py +++ b/taskbadger/sdk.py @@ -184,6 +184,7 @@ def create_task( queue: str = None, external_id: str = None, parent: str = None, + is_retry: bool = None, ) -> "Task": """Create a Task. @@ -203,6 +204,9 @@ def create_task( parent: ID of the parent task. Tasks nest a single level deep, so this must be the ID of a task that is not itself a child. The Celery and Procrastinate integrations set this automatically for tasks enqueued from within a tracked task. + is_retry: Mark the task as a retry of `parent`. The first 3 retries of a parent don't + count towards your monthly task allocation. Ignored without `parent`. The Celery + integration sets this automatically for retries. Returns: Task: The created Task object. @@ -213,6 +217,8 @@ def create_task( } if parent is not None: task_dict["parent"] = parent + if is_retry: + task_dict["is_retry"] = True if queue is not None: task_dict["queue"] = queue if external_id is not None: @@ -392,6 +398,7 @@ def create( queue: str = None, external_id: str = None, parent: str = None, + is_retry: bool = None, ) -> "Task": """Create a new task @@ -411,6 +418,7 @@ def create( queue=queue, external_id=external_id, parent=parent, + is_retry=is_retry, ) def __init__(self, task): diff --git a/tests/test_celery.py b/tests/test_celery.py index 56ca7bc..94f6ebb 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -466,12 +466,18 @@ def add_retry(self, a, b, is_retry=False): mock.patch("taskbadger.sdk.get_task") as get_task, ): create.return_value = task_for_test() + update.return_value = task_for_test() get_task.return_value = task_for_test() result = add_retry.delay(2, 2) assert result.get(timeout=10, propagate=True) == 4 _wait_for_call_count(update, 4) assert create.call_count == 2 + first, retry = create.call_args_list + # the original attempt is an ordinary task; the retry is marked and nested under it + assert "is_retry" not in first.kwargs + assert retry.kwargs["is_retry"] is True + assert retry.kwargs["parent"] is not None assert update.call_args_list == [ mock.call(mock.ANY, status=StatusEnum.PROCESSING, data=mock.ANY), mock.call(mock.ANY, status=StatusEnum.ERROR, data=mock.ANY), diff --git a/tests/test_parents.py b/tests/test_parents.py index f648278..ab60c06 100644 --- a/tests/test_parents.py +++ b/tests/test_parents.py @@ -495,3 +495,21 @@ def records_parent(): assert seen["parent"] == "job_tb" assert current_task_id() is None + + +def test_retry_parent_id(): + from taskbadger._current_task import enter_task, exit_task, retry_parent_id + + assert retry_parent_id() is None # outside a tracked task + + token = enter_task("first") # a root task: its retry nests under it + assert retry_parent_id() == "first" + exit_task(token) + + token = enter_task("retry", "first", is_retry=True) # retries nest under the first attempt + assert retry_parent_id() == "first" + exit_task(token) + + token = enter_task("child", "batch") # a fan-out child can't have nested retries + assert retry_parent_id() is None + exit_task(token) From 7d5a307796e614483d50131fb85f2492e32fc0ad Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 14:22:00 +0000 Subject: [PATCH 2/4] feat(celery): send retry_of so retries of fan-out children are free too Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01WDSbaMzpzXZKif1HbDFaT3 --- taskbadger.yaml | 47 ++++++++++++++----- taskbadger/_current_task.py | 30 ++++++------ taskbadger/_integrations.py | 7 +-- taskbadger/celery.py | 20 +++----- .../internal/models/patched_task_request.py | 28 +++++++++-- taskbadger/internal/models/task.py | 28 +++++++++-- taskbadger/internal/models/task_request.py | 28 +++++++++-- taskbadger/sdk.py | 16 +++---- tests/test_celery.py | 6 +-- tests/test_parents.py | 22 +++++---- 10 files changed, 159 insertions(+), 73 deletions(-) diff --git a/taskbadger.yaml b/taskbadger.yaml index affd5bd..9c8aa02 100644 --- a/taskbadger.yaml +++ b/taskbadger.yaml @@ -712,12 +712,20 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of: + type: string + minLength: 1 + nullable: true + description: 'ID of the first attempt of the task this task is a retry of. + The first 3 retries of a task don''t count towards the monthly task allocation. + A retry is nested like a child of the original: under it for a top-level + task, otherwise alongside it under the same parent. Retries of a retry + point at the original. Can only be set when the task is created.' is_retry: type: boolean - description: Set when this task is a retry of its parent task. The first - 3 retries of a parent don't count towards the monthly task allocation. - Ignored unless 'parent' is also set (a retry with no parent is a new task) - and can only be set when the task is created. + description: 'Whether this task is a retry. On create, deprecated in favour + of ''retry_of'': marks the task as a retry of its ''parent'', which must + be a top-level task, and is ignored without ''parent''.' name: type: string minLength: 1 @@ -829,12 +837,19 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of: + type: string + nullable: true + description: 'ID of the first attempt of the task this task is a retry of. + The first 3 retries of a task don''t count towards the monthly task allocation. + A retry is nested like a child of the original: under it for a top-level + task, otherwise alongside it under the same parent. Retries of a retry + point at the original. Can only be set when the task is created.' is_retry: type: boolean - description: Set when this task is a retry of its parent task. The first - 3 retries of a parent don't count towards the monthly task allocation. - Ignored unless 'parent' is also set (a retry with no parent is a new task) - and can only be set when the task is created. + description: 'Whether this task is a retry. On create, deprecated in favour + of ''retry_of'': marks the task as a retry of its ''parent'', which must + be a top-level task, and is ignored without ''parent''.' name: type: string description: Name of the task @@ -944,12 +959,20 @@ components: nullable: true description: ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. + retry_of: + type: string + minLength: 1 + nullable: true + description: 'ID of the first attempt of the task this task is a retry of. + The first 3 retries of a task don''t count towards the monthly task allocation. + A retry is nested like a child of the original: under it for a top-level + task, otherwise alongside it under the same parent. Retries of a retry + point at the original. Can only be set when the task is created.' is_retry: type: boolean - description: Set when this task is a retry of its parent task. The first - 3 retries of a parent don't count towards the monthly task allocation. - Ignored unless 'parent' is also set (a retry with no parent is a new task) - and can only be set when the task is created. + description: 'Whether this task is a retry. On create, deprecated in favour + of ''retry_of'': marks the task as a retry of its ''parent'', which must + be a top-level task, and is ignored without ''parent''.' name: type: string minLength: 1 diff --git a/taskbadger/_current_task.py b/taskbadger/_current_task.py index 9bf6534..ac5fb1c 100644 --- a/taskbadger/_current_task.py +++ b/taskbadger/_current_task.py @@ -14,22 +14,23 @@ from contextvars import ContextVar -# (id of the running task, id of its parent or None if it is a root task, whether it is a retry) -_current: ContextVar[tuple[str, str | None, bool] | None] = ContextVar("taskbadger_current_task", default=None) +# (id of the running task, id of its parent or None if it is a root task, +# id of the first attempt if the running task is a retry, else None) +_current: ContextVar[tuple[str, str | None, str | None] | None] = ContextVar("taskbadger_current_task", default=None) -def enter_task(task_id: str, parent: str = None, is_retry: bool = False): +def enter_task(task_id: str, parent: str = None, retry_of: str = None): """Mark `task_id` as the task running in this context. Arguments: task_id: The running task. parent: The running task's own parent, if it has one. - is_retry: Whether the running task is itself a retry of its parent. + retry_of: The first attempt, if the running task is itself a retry. Returns: A token to pass to `exit_task`. """ - return _current.set((task_id, parent, is_retry)) + return _current.set((task_id, parent, retry_of)) def exit_task(token) -> None: @@ -54,19 +55,16 @@ def parent_id() -> str | None: return parent or task_id -def retry_parent_id() -> str | None: - """The id to use as `parent` for a retry of the task running right now. +def retry_of_id() -> str | None: + """The id to use as `retry_of` for a retry of the task running right now. - A retry is nested under the first attempt, so this is the running task for the first attempt - and the first attempt again when the running task is itself a retry. `None` when there is no - suitable first attempt: outside a tracked task, or when the running task is a child of some - other task (a fan-out child), since retries, like children, can't be nested a level deeper. - The caller should then create the retry as an ordinary task. + The first attempt: the running task itself, or the task it is a retry of. Retries all point at + the first attempt, so the API can count them per task. `None` outside a tracked task. The + retry's `parent` is still `parent_id()`: under the first attempt for a top-level task, + alongside it for a child. """ current = _current.get() if current is None: return None - task_id, parent, is_retry = current - if parent is None: - return task_id - return parent if is_retry else None + task_id, _, retry_of = current + return retry_of or task_id diff --git a/taskbadger/_integrations.py b/taskbadger/_integrations.py index add4b87..c309324 100644 --- a/taskbadger/_integrations.py +++ b/taskbadger/_integrations.py @@ -88,9 +88,10 @@ def parent_of(task) -> str | None: return parent or None -def is_retry_task(task) -> bool: - """Whether ``task`` was created as a retry of its parent (``False`` for ``None``/unset).""" - return getattr(task, "is_retry", False) is True +def retry_of(task) -> str | None: + """The id of the first attempt ``task`` is a retry of, or ``None`` (also for ``None``/unset).""" + value = getattr(task, "retry_of", None) if task is not None else None + return value or None def match_task_name(task_name: str, includes, excludes) -> bool: diff --git a/taskbadger/celery.py b/taskbadger/celery.py index 62e686f..f397e75 100644 --- a/taskbadger/celery.py +++ b/taskbadger/celery.py @@ -14,14 +14,14 @@ from kombu import serialization from . import sdk -from ._current_task import enter_task, exit_task, parent_id, retry_parent_id +from ._current_task import enter_task, exit_task, parent_id, retry_of_id from ._error_context import capture_error_data, reset_error_context, start_error_context from ._heartbeat import heartbeat from ._integrations import ( TERMINAL_STATES, - is_retry_task, parent_of, resolve_heartbeat_options, + retry_of, safe_get_task, task_cache, ) @@ -254,12 +254,9 @@ def task_publish_handler(sender=None, headers=None, body=None, **kwargs): enclosing_task = parent_id() if enclosing_task and not _is_workflow_successor(sender): if "parent" not in kwargs and headers.get("retries") and _is_retry_of_running_task(headers): - # A Celery retry republishes the same task. It is marked so the API can leave - # the first few out of the monthly allocation, and nested under the first attempt. - retry_parent = retry_parent_id() - if retry_parent: - kwargs["parent"] = retry_parent - kwargs["is_retry"] = True + # A Celery retry republishes the same task. It names the first attempt so the API + # can leave the first few retries of it out of the monthly allocation. + kwargs["retry_of"] = retry_of_id() kwargs.setdefault("parent", enclosing_task) name = kwargs.pop("name", headers["task"]) @@ -398,10 +395,7 @@ def _maybe_create_task(signal_sender): } if "parent" not in header_kwargs and getattr(signal_sender.request, "retries", 0): # An eager retry runs nested inside the attempt that retried it. - retry_parent = retry_parent_id() - if retry_parent: - create_kwargs["parent"] = retry_parent - create_kwargs["is_retry"] = True + create_kwargs["retry_of"] = retry_of_id() heartbeat_interval, stale_timeout = resolve_heartbeat_options( header_kwargs.get("heartbeat_interval", getattr(signal_sender, TB_HEARTBEAT_INTERVAL, None)), header_kwargs.get("stale_timeout", getattr(signal_sender, TB_STALE_TIMEOUT, None)), @@ -438,7 +432,7 @@ def task_prerun_handler(sender=None, **kwargs): # it. `_update_task` above has already put it in the cache, so this is a # lookup rather than a fetch. running = safe_get_task(task_id) - current_token = enter_task(task_id, parent_of(running), is_retry_task(running)) + current_token = enter_task(task_id, parent_of(running), retry_of(running)) sender.request.update({TB_CURRENT_TASK_TOKEN: current_token}) diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index fefb4dc..ec56b19 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/taskbadger/internal/models/patched_task_request.py @@ -24,9 +24,12 @@ class PatchedTaskRequest: Attributes: parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. - is_retry (bool | Unset): Set when this task is a retry of its parent task. The first 3 retries of a parent don't - count towards the monthly task allocation. Ignored unless 'parent' is also set (a retry with no parent is a new - task) and can only be set when the task is created. + retry_of (None | str | Unset): ID of the first attempt of the task this task is a retry of. The first 3 retries + of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under + it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. + Can only be set when the task is created. + is_retry (bool | Unset): Whether this task is a retry. On create, deprecated in favour of 'retry_of': marks the + task as a retry of its 'parent', which must be a top-level task, and is ignored without 'parent'. name (str | Unset): Name of the task queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with @@ -56,6 +59,7 @@ class PatchedTaskRequest: """ parent: None | str | Unset = UNSET + retry_of: None | str | Unset = UNSET is_retry: bool | Unset = UNSET name: str | Unset = UNSET queue: str | Unset = UNSET @@ -81,6 +85,12 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + retry_of: None | str | Unset + if isinstance(self.retry_of, Unset): + retry_of = UNSET + else: + retry_of = self.retry_of + is_retry = self.is_retry name = self.name @@ -146,6 +156,8 @@ def to_dict(self) -> dict[str, Any]: field_dict.update({}) if parent is not UNSET: field_dict["parent"] = parent + if retry_of is not UNSET: + field_dict["retry_of"] = retry_of if is_retry is not UNSET: field_dict["is_retry"] = is_retry if name is not UNSET: @@ -192,6 +204,15 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + def _parse_retry_of(data: object) -> None | str | Unset: + if data is None: + return data + if isinstance(data, Unset): + return data + return cast(None | str | Unset, data) + + retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) + is_retry = d.pop("is_retry", UNSET) name = d.pop("name", UNSET) @@ -290,6 +311,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: patched_task_request = cls( parent=parent, + retry_of=retry_of, is_retry=is_retry, name=name, queue=queue, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index 451d64e..c7ce5a0 100644 --- a/taskbadger/internal/models/task.py +++ b/taskbadger/internal/models/task.py @@ -33,9 +33,12 @@ class Task: public_url (str): parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. - is_retry (bool | Unset): Set when this task is a retry of its parent task. The first 3 retries of a parent don't - count towards the monthly task allocation. Ignored unless 'parent' is also set (a retry with no parent is a new - task) and can only be set when the task is created. + retry_of (None | str | Unset): ID of the first attempt of the task this task is a retry of. The first 3 retries + of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under + it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. + Can only be set when the task is created. + is_retry (bool | Unset): Whether this task is a retry. On create, deprecated in favour of 'retry_of': marks the + task as a retry of its 'parent', which must be a top-level task, and is ignored without 'parent'. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -73,6 +76,7 @@ class Task: url: str public_url: str parent: None | str | Unset = UNSET + retry_of: None | str | Unset = UNSET is_retry: bool | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET @@ -116,6 +120,12 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + retry_of: None | str | Unset + if isinstance(self.retry_of, Unset): + retry_of = UNSET + else: + retry_of = self.retry_of + is_retry = self.is_retry queue = self.queue @@ -191,6 +201,8 @@ def to_dict(self) -> dict[str, Any]: ) if parent is not UNSET: field_dict["parent"] = parent + if retry_of is not UNSET: + field_dict["retry_of"] = retry_of if is_retry is not UNSET: field_dict["is_retry"] = is_retry if queue is not UNSET: @@ -257,6 +269,15 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + def _parse_retry_of(data: object) -> None | str | Unset: + if data is None: + return data + if isinstance(data, Unset): + return data + return cast(None | str | Unset, data) + + retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) + is_retry = d.pop("is_retry", UNSET) queue = d.pop("queue", UNSET) @@ -362,6 +383,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: url=url, public_url=public_url, parent=parent, + retry_of=retry_of, is_retry=is_retry, queue=queue, external_id=external_id, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index c82800c..902687d 100644 --- a/taskbadger/internal/models/task_request.py +++ b/taskbadger/internal/models/task_request.py @@ -25,9 +25,12 @@ class TaskRequest: name (str): Name of the task parent (None | str | Unset): ID of the task this task is part of. Tasks can only be nested one level deep, and a task's parent can not be changed once set. - is_retry (bool | Unset): Set when this task is a retry of its parent task. The first 3 retries of a parent don't - count towards the monthly task allocation. Ignored unless 'parent' is also set (a retry with no parent is a new - task) and can only be set when the task is created. + retry_of (None | str | Unset): ID of the first attempt of the task this task is a retry of. The first 3 retries + of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under + it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. + Can only be set when the task is created. + is_retry (bool | Unset): Whether this task is a retry. On create, deprecated in favour of 'retry_of': marks the + task as a retry of its 'parent', which must be a top-level task, and is ignored without 'parent'. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -57,6 +60,7 @@ class TaskRequest: name: str parent: None | str | Unset = UNSET + retry_of: None | str | Unset = UNSET is_retry: bool | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET @@ -83,6 +87,12 @@ def to_dict(self) -> dict[str, Any]: else: parent = self.parent + retry_of: None | str | Unset + if isinstance(self.retry_of, Unset): + retry_of = UNSET + else: + retry_of = self.retry_of + is_retry = self.is_retry queue = self.queue @@ -150,6 +160,8 @@ def to_dict(self) -> dict[str, Any]: ) if parent is not UNSET: field_dict["parent"] = parent + if retry_of is not UNSET: + field_dict["retry_of"] = retry_of if is_retry is not UNSET: field_dict["is_retry"] = is_retry if queue is not UNSET: @@ -195,6 +207,15 @@ def _parse_parent(data: object) -> None | str | Unset: parent = _parse_parent(d.pop("parent", UNSET)) + def _parse_retry_of(data: object) -> None | str | Unset: + if data is None: + return data + if isinstance(data, Unset): + return data + return cast(None | str | Unset, data) + + retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) + is_retry = d.pop("is_retry", UNSET) queue = d.pop("queue", UNSET) @@ -292,6 +313,7 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: task_request = cls( name=name, parent=parent, + retry_of=retry_of, is_retry=is_retry, queue=queue, external_id=external_id, diff --git a/taskbadger/sdk.py b/taskbadger/sdk.py index f2905fd..a9b651b 100644 --- a/taskbadger/sdk.py +++ b/taskbadger/sdk.py @@ -184,7 +184,7 @@ def create_task( queue: str = None, external_id: str = None, parent: str = None, - is_retry: bool = None, + retry_of: str = None, ) -> "Task": """Create a Task. @@ -204,9 +204,9 @@ def create_task( parent: ID of the parent task. Tasks nest a single level deep, so this must be the ID of a task that is not itself a child. The Celery and Procrastinate integrations set this automatically for tasks enqueued from within a tracked task. - is_retry: Mark the task as a retry of `parent`. The first 3 retries of a parent don't - count towards your monthly task allocation. Ignored without `parent`. The Celery - integration sets this automatically for retries. + retry_of: ID of the first attempt, to mark the task as a retry of it. The first 3 retries + of a task don't count towards your monthly task allocation. The Celery integration + sets this automatically for retries. Returns: Task: The created Task object. @@ -217,8 +217,8 @@ def create_task( } if parent is not None: task_dict["parent"] = parent - if is_retry: - task_dict["is_retry"] = True + if retry_of is not None: + task_dict["retry_of"] = retry_of if queue is not None: task_dict["queue"] = queue if external_id is not None: @@ -398,7 +398,7 @@ def create( queue: str = None, external_id: str = None, parent: str = None, - is_retry: bool = None, + retry_of: str = None, ) -> "Task": """Create a new task @@ -418,7 +418,7 @@ def create( queue=queue, external_id=external_id, parent=parent, - is_retry=is_retry, + retry_of=retry_of, ) def __init__(self, task): diff --git a/tests/test_celery.py b/tests/test_celery.py index 94f6ebb..23b09f7 100644 --- a/tests/test_celery.py +++ b/tests/test_celery.py @@ -475,9 +475,9 @@ def add_retry(self, a, b, is_retry=False): assert create.call_count == 2 first, retry = create.call_args_list # the original attempt is an ordinary task; the retry is marked and nested under it - assert "is_retry" not in first.kwargs - assert retry.kwargs["is_retry"] is True - assert retry.kwargs["parent"] is not None + assert "retry_of" not in first.kwargs + assert retry.kwargs["retry_of"] is not None + assert retry.kwargs["parent"] == retry.kwargs["retry_of"] # a top-level task's retry nests under it assert update.call_args_list == [ mock.call(mock.ANY, status=StatusEnum.PROCESSING, data=mock.ANY), mock.call(mock.ANY, status=StatusEnum.ERROR, data=mock.ANY), diff --git a/tests/test_parents.py b/tests/test_parents.py index ab60c06..6138b3f 100644 --- a/tests/test_parents.py +++ b/tests/test_parents.py @@ -497,19 +497,23 @@ def records_parent(): assert current_task_id() is None -def test_retry_parent_id(): - from taskbadger._current_task import enter_task, exit_task, retry_parent_id +def test_retry_of_id(): + from taskbadger._current_task import enter_task, exit_task, parent_id, retry_of_id - assert retry_parent_id() is None # outside a tracked task + assert retry_of_id() is None # outside a tracked task - token = enter_task("first") # a root task: its retry nests under it - assert retry_parent_id() == "first" + token = enter_task("first") # a top-level task: its retries point at it and nest under it + assert (retry_of_id(), parent_id()) == ("first", "first") exit_task(token) - token = enter_task("retry", "first", is_retry=True) # retries nest under the first attempt - assert retry_parent_id() == "first" + token = enter_task("retry", "first", retry_of="first") # retries of retries point at the first attempt + assert (retry_of_id(), parent_id()) == ("first", "first") exit_task(token) - token = enter_task("child", "batch") # a fan-out child can't have nested retries - assert retry_parent_id() is None + token = enter_task("child", "batch") # a fan-out child's retries sit beside it under the batch + assert (retry_of_id(), parent_id()) == ("child", "batch") + exit_task(token) + + token = enter_task("child-retry", "batch", retry_of="child") + assert (retry_of_id(), parent_id()) == ("child", "batch") exit_task(token) From 79284d2c7c9a7eacde3c1514776a1a6635423605 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 2 Nov 2026 15:20:01 +0000 Subject: [PATCH 3/4] chore(internal): regenerate client without is_retry Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01WDSbaMzpzXZKif1HbDFaT3 --- taskbadger.yaml | 15 --------------- .../internal/models/patched_task_request.py | 10 ---------- taskbadger/internal/models/task.py | 10 ---------- taskbadger/internal/models/task_request.py | 10 ---------- 4 files changed, 45 deletions(-) diff --git a/taskbadger.yaml b/taskbadger.yaml index 9c8aa02..e369e50 100644 --- a/taskbadger.yaml +++ b/taskbadger.yaml @@ -721,11 +721,6 @@ components: A retry is nested like a child of the original: under it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. Can only be set when the task is created.' - is_retry: - type: boolean - description: 'Whether this task is a retry. On create, deprecated in favour - of ''retry_of'': marks the task as a retry of its ''parent'', which must - be a top-level task, and is ignored without ''parent''.' name: type: string minLength: 1 @@ -845,11 +840,6 @@ components: A retry is nested like a child of the original: under it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. Can only be set when the task is created.' - is_retry: - type: boolean - description: 'Whether this task is a retry. On create, deprecated in favour - of ''retry_of'': marks the task as a retry of its ''parent'', which must - be a top-level task, and is ignored without ''parent''.' name: type: string description: Name of the task @@ -968,11 +958,6 @@ components: A retry is nested like a child of the original: under it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. Can only be set when the task is created.' - is_retry: - type: boolean - description: 'Whether this task is a retry. On create, deprecated in favour - of ''retry_of'': marks the task as a retry of its ''parent'', which must - be a top-level task, and is ignored without ''parent''.' name: type: string minLength: 1 diff --git a/taskbadger/internal/models/patched_task_request.py b/taskbadger/internal/models/patched_task_request.py index ec56b19..1319d2c 100644 --- a/taskbadger/internal/models/patched_task_request.py +++ b/taskbadger/internal/models/patched_task_request.py @@ -28,8 +28,6 @@ class PatchedTaskRequest: of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. Can only be set when the task is created. - is_retry (bool | Unset): Whether this task is a retry. On create, deprecated in favour of 'retry_of': marks the - task as a retry of its 'parent', which must be a top-level task, and is ignored without 'parent'. name (str | Unset): Name of the task queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with @@ -60,7 +58,6 @@ class PatchedTaskRequest: parent: None | str | Unset = UNSET retry_of: None | str | Unset = UNSET - is_retry: bool | Unset = UNSET name: str | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET @@ -91,8 +88,6 @@ def to_dict(self) -> dict[str, Any]: else: retry_of = self.retry_of - is_retry = self.is_retry - name = self.name queue = self.queue @@ -158,8 +153,6 @@ def to_dict(self) -> dict[str, Any]: field_dict["parent"] = parent if retry_of is not UNSET: field_dict["retry_of"] = retry_of - if is_retry is not UNSET: - field_dict["is_retry"] = is_retry if name is not UNSET: field_dict["name"] = name if queue is not UNSET: @@ -213,8 +206,6 @@ def _parse_retry_of(data: object) -> None | str | Unset: retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) - is_retry = d.pop("is_retry", UNSET) - name = d.pop("name", UNSET) queue = d.pop("queue", UNSET) @@ -312,7 +303,6 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: patched_task_request = cls( parent=parent, retry_of=retry_of, - is_retry=is_retry, name=name, queue=queue, external_id=external_id, diff --git a/taskbadger/internal/models/task.py b/taskbadger/internal/models/task.py index c7ce5a0..01250a2 100644 --- a/taskbadger/internal/models/task.py +++ b/taskbadger/internal/models/task.py @@ -37,8 +37,6 @@ class Task: of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. Can only be set when the task is created. - is_retry (bool | Unset): Whether this task is a retry. On create, deprecated in favour of 'retry_of': marks the - task as a retry of its 'parent', which must be a top-level task, and is ignored without 'parent'. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -77,7 +75,6 @@ class Task: public_url: str parent: None | str | Unset = UNSET retry_of: None | str | Unset = UNSET - is_retry: bool | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -126,8 +123,6 @@ def to_dict(self) -> dict[str, Any]: else: retry_of = self.retry_of - is_retry = self.is_retry - queue = self.queue external_id = self.external_id @@ -203,8 +198,6 @@ def to_dict(self) -> dict[str, Any]: field_dict["parent"] = parent if retry_of is not UNSET: field_dict["retry_of"] = retry_of - if is_retry is not UNSET: - field_dict["is_retry"] = is_retry if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -278,8 +271,6 @@ def _parse_retry_of(data: object) -> None | str | Unset: retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) - is_retry = d.pop("is_retry", UNSET) - queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -384,7 +375,6 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: public_url=public_url, parent=parent, retry_of=retry_of, - is_retry=is_retry, queue=queue, external_id=external_id, status=status, diff --git a/taskbadger/internal/models/task_request.py b/taskbadger/internal/models/task_request.py index 902687d..481f6c0 100644 --- a/taskbadger/internal/models/task_request.py +++ b/taskbadger/internal/models/task_request.py @@ -29,8 +29,6 @@ class TaskRequest: of a task don't count towards the monthly task allocation. A retry is nested like a child of the original: under it for a top-level task, otherwise alongside it under the same parent. Retries of a retry point at the original. Can only be set when the task is created. - is_retry (bool | Unset): Whether this task is a retry. On create, deprecated in favour of 'retry_of': marks the - task as a retry of its 'parent', which must be a top-level task, and is ignored without 'parent'. queue (str | Unset): Queue the task is from external_id (str | Unset): Identifier from the originating system (e.g. Celery task ID) for correlating with logs @@ -61,7 +59,6 @@ class TaskRequest: name: str parent: None | str | Unset = UNSET retry_of: None | str | Unset = UNSET - is_retry: bool | Unset = UNSET queue: str | Unset = UNSET external_id: str | Unset = UNSET status: StatusEnum | Unset = StatusEnum.PENDING @@ -93,8 +90,6 @@ def to_dict(self) -> dict[str, Any]: else: retry_of = self.retry_of - is_retry = self.is_retry - queue = self.queue external_id = self.external_id @@ -162,8 +157,6 @@ def to_dict(self) -> dict[str, Any]: field_dict["parent"] = parent if retry_of is not UNSET: field_dict["retry_of"] = retry_of - if is_retry is not UNSET: - field_dict["is_retry"] = is_retry if queue is not UNSET: field_dict["queue"] = queue if external_id is not UNSET: @@ -216,8 +209,6 @@ def _parse_retry_of(data: object) -> None | str | Unset: retry_of = _parse_retry_of(d.pop("retry_of", UNSET)) - is_retry = d.pop("is_retry", UNSET) - queue = d.pop("queue", UNSET) external_id = d.pop("external_id", UNSET) @@ -314,7 +305,6 @@ def _parse_stale_timeout(data: object) -> int | None | Unset: name=name, parent=parent, retry_of=retry_of, - is_retry=is_retry, queue=queue, external_id=external_id, status=status, From 66e90d6adf3eeaef7b00bfb675346d14e0be32a6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 3 Nov 2026 09:59:02 +0000 Subject: [PATCH 4/4] test(celery): cover retries of fan-out children and the eager retry path Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01WDSbaMzpzXZKif1HbDFaT3 --- tests/test_parents.py | 62 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/tests/test_parents.py b/tests/test_parents.py index 6138b3f..8d67a30 100644 --- a/tests/test_parents.py +++ b/tests/test_parents.py @@ -430,6 +430,68 @@ def test_celery_eager_explicit_none_parent_makes_a_root_task(): assert create.call_args.kwargs["parent"] is None +@pytest.mark.usefixtures("_bind_settings") +def test_celery_publish_of_a_retry_of_a_fan_out_child_sits_beside_it(): + """A retry is published with the running task's own id and a `retries` header. For a + child of another task it points at the child and stays under the same parent.""" + add = _celery_app() + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.sdk.get_task"), + mock.patch("taskbadger.celery._is_retry_of_running_task", return_value=True), + ): + create.return_value = task_for_test() + token = enter_task("child", "batch") + try: + add.apply_async((2, 2), retries=1) + finally: + exit_task(token) + + assert create.call_args.kwargs["retry_of"] == "child" + assert create.call_args.kwargs["parent"] == "batch" + + +@pytest.mark.usefixtures("_bind_settings") +def test_celery_publish_without_retries_header_is_not_a_retry(): + add = _celery_app() + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.sdk.get_task"), + mock.patch("taskbadger.celery._is_retry_of_running_task", return_value=True), + ): + create.return_value = task_for_test() + token = enter_task("root_id") + try: + add.apply_async((2, 2)) + finally: + exit_task(token) + + assert "retry_of" not in create.call_args.kwargs + assert create.call_args.kwargs["parent"] == "root_id" + + +@pytest.mark.usefixtures("_bind_settings") +def test_celery_eager_retry_points_at_the_first_attempt(): + add = _celery_app(task_always_eager=True, task_eager_propagates=True) + + with ( + mock.patch("taskbadger.celery.create_task_safe") as create, + mock.patch("taskbadger.celery.update_task_safe"), + mock.patch("taskbadger.sdk.get_task"), + ): + create.return_value = task_for_test() + token = enter_task("first") + try: + add.apply_async((2, 2), retries=1) + finally: + exit_task(token) + + assert create.call_args.kwargs["retry_of"] == "first" + assert create.call_args.kwargs["parent"] == "first" + + # --- Procrastinate ------------------------------------------------------------