Repository: celery/celery

Issue:
Custom header lost after using autoretry
<!--
Please fill this template entirely and do not erase parts of it.
We reserve the right to close without a response
bug reports which are incomplete.
-->
# Checklist
<!--
To check an item on the list replace [ ] with [x].
-->
- [x] I have verified that the issue exists against the `master` branch of Celery.
- [ ] This has already been asked to the [discussions forum](https://github.com/celery/celery/discussions) first.
- [x] I have read the relevant section in the
  [contribution guide](https://docs.celeryq.dev/en/master/contributing.html#other-bugs)
  on reporting bugs.
- [x] I have checked the [issues list](https://github.com/celery/celery/issues?q=is%3Aissue+label%3A%22Issue+Type%3A+Bug+Report%22+-label%3A%22Category%3A+Documentation%22)
  for similar or identical bug reports.
- [x] I have checked the [pull requests list](https://github.com/celery/celery/pulls?q=is%3Apr+label%3A%22PR+Type%3A+Bugfix%22+-label%3A%22Category%3A+Documentation%22)
  for existing proposed fixes.
- [ ] I have checked the [commit log](https://github.com/celery/celery/commits/master)
  to find out if the bug was already fixed in the master branch.
- [x] I have included all related issues and possible duplicate issues
  in this issue (If there are none, check this box anyway).

## Mandatory Debugging Information

- [x] I have included the output of ``celery -A proj report`` in the issue.
    (if you are not able to do this, then at least specify the Celery
     version affected).
- [x] I have verified that the issue exists against the `master` branch of Celery.
- [ ] I have included the contents of ``pip freeze`` in the issue.
- [x] I have included all the versions of all the external dependencies required
  to reproduce this bug.

## Optional Debugging Information
<!--
Try some of the below if you think they are relevant.
It will help us figure out the scope of the bug and how many users it affects.
-->
- [ ] I have tried reproducing the issue on more than one Python version
  and/or implementation.
- [ ] I have tried reproducing the issue on more than one message broker and/or
  result backend.
- [ ] I have tried reproducing the issue on more than one version of the message
  broker and/or result backend.
- [ ] I have tried reproducing the issue on more than one operating system.
- [ ] I have tried reproducing the issue on more than one workers pool.
- [ ] I have tried reproducing the issue with autoscaling, retries,
  ETA/Countdown & rate limits disabled.
- [ ] I have tried reproducing the issue after downgrading
  and/or upgrading Celery and its dependencies.

## Related Issues and Possible Duplicates
<!--
Please make sure to search and mention any related issues
or possible duplicates to this issue as requested by the checklist above.

This may or may not include issues in other repositories that the Celery project
maintains or other repositories that are dependencies of Celery.

If you don't know how to mention issues, please refer to Github's documentation
on the subject: https://help.github.com/en/articles/autolinked-references-and-urls#issues-and-pull-requests
-->

#### Related Issues

- None

#### Possible Duplicates

- None

## Environment & Settings
<!-- Include the contents of celery --version below -->
**Celery version**: 5.2.6 (dawn-chorus)
<!-- Include the output of celery -A proj report below -->
<details>
<summary><b><code>celery report
</code> Output:
</b></summary>
<p>

```
software -> celery:5.2.6 (dawn-chorus) kombu:5.2.3 py:3.10.0
            billiard:3.6.4.0 py-amqp:5.0.9
platform -> system:Windows arch:64bit, WindowsPE
            kernel version:10 imp:CPython
loader   -> celery.loaders.app.AppLoader
settings -> transport:amqp results:redis://localhost:6379/0

broker_url: 'amqp://user:********@localhost:5672//'
task_create_missing_queues: True
task_acks_late: True
result_extended: True
result_expires: 3600
worker_send_task_events: True
result_backend: 'redis://localhost:6379/0'
deprecated_settings: None
```

</p>
</details>

# Steps to Reproduce

1. Create simplest docker compose with redis & rabbitMQ
2. run app with `celery --app=app worker --pool=solo --concurrency=1 -E --loglevel=INFO --queues=my_queue `
3. run `python call.py`
4. You see in retried task headers there in no `"custom_header": "yumy yum"`


## Required Dependencies
<!-- Please fill the required dependencies to reproduce this issue -->
* **Minimal Python Version**: 3.10.0
* **Minimal Celery Version**: N/A or Unknown
* **Minimal Kombu Version**: N/A or Unknown
* **Minimal Broker Version**: N/A or Unknown
* **Minimal Result Backend Version**: N/A or Unknown
* **Minimal OS and/or Kernel Version**: N/A or Unknown
* **Minimal Broker Client Version**: N/A or Unknown
* **Minimal Result Backend Client Version**: N/A or Unknown

### Python Packages
<!-- Please fill the contents of pip freeze below -->
<details>
<summary><b><code>pip freeze</code> Output:</b></summary>
<p>

```
```

</p>
</details>

### Other Dependencies
<!--
Please provide system dependencies, configuration files
and other dependency information if applicable
-->
<details>
<p>
N/A
</p>
</details>

## Minimally Reproducible Test Case
<!--
Please provide a reproducible test case.
Refer to the Reporting Bugs section in our contribution guide.

We prefer submitting test cases in the form of a PR to our integration test suite.
If you can provide one, please mention the PR number below.
If not, please attach the most minimal code example required to reproduce the issue below.
If the test case is too large, please include a link to a gist or a repository below.
-->

<details>
<p>

Application app.py:
```python
from functools import partial
from typing import Final

from celery import Task, Celery

BASE_CELERY_BROKER_URL = "amqp://user:pass@localhost:5672"
CELERY_RESULT_URL = "redis://localhost:6379/0"


def make_celery_worker_app_stub() -> partial[Celery]:
    return partial(
        Celery,
        task_create_missing_queues=True,
        task_acks_late=True,
        result_extended=True,
        result_expires=3600,
        worker_send_task_events=True,
    )


app: Final = make_celery_worker_app_stub()(
    main=__name__,
    result_backend=CELERY_RESULT_URL,
    broker="%s" % (BASE_CELERY_BROKER_URL),

    set_as_current=False,
)


@app.task(
    bind=True,
    queue="my_queue",
    name="my_name:my_function",
    acks_late=True,
    task_reject_on_worker_lost=True,
    autoretry_for=(Exception,),
    retry_backoff=True,
    max_retries=3,
    task_time_limit=8,
)
def failing_function(
        self: Task,
        data
) -> None:
    print(data)  # prints hello
    print(self.request.properties["application_headers"])

    assert False == True

```
Caller call.py:
```python
from celery import Celery

BASE_CELERY_BROKER_URL = "amqp://user:pass@localhost:5672"
CELERY_RESULT_URL = "redis://localhost:6379/0"


def celery_app(
) -> Celery:
    return Celery(
        __name__,
        result_backend=CELERY_RESULT_URL,
        broker="%s" % (BASE_CELERY_BROKER_URL),
        task_create_missing_queues=True,
        celery_queue_ha_policy="all",
        task_acks_late=True,
        set_as_current=False,
        result_extended=True,
    )


client = celery_app()

options = {"headers": {
    "custom_header": "yumy yum"
}}

sign = client.signature(
    "my_name:my_function",
    queue="my_queue",
    args=("hello",),
    options=options
)
sign.apply_async()

```


Log:
```
[2022-05-12 22:37:20,703: INFO/MainProcess] Connected to amqp://user:**@127.0.0.1:5672//
[2022-05-12 22:37:20,736: INFO/MainProcess] mingle: searching for neighbors
[2022-05-12 22:37:21,764: INFO/MainProcess] mingle: all alone
[2022-05-12 22:37:21,799: INFO/MainProcess] celery@DESKTOP-51C03JP ready.
[2022-05-12 22:38:47,431: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] received
[2022-05-12 22:38:47,432: WARNING/MainProcess] hello
[2022-05-12 22:38:47,432: WARNING/MainProcess] {'lang': 'py', 'task': 'my_name:my_function', 'id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'shadow': None, 'eta': None, 'expires': None, 'group': None, 'group_index': None, 'retries': 0, 'timelimit': [None, None], 'root_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'parent_id': None, 'argsrepr': "('hello',)", 'kwargsrepr': '{}', 'origin': 'gen5572@DESKTOP-51C03JP', 'ignore_result': False, 'custom_header': 'yumy yum'}
[2022-05-12 22:38:47,468: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] retry: Retry in 1s: AssertionError()
[2022-05-12 22:38:47,469: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] received
[2022-05-12 22:38:48,446: WARNING/MainProcess] hello
[2022-05-12 22:38:48,447: WARNING/MainProcess] {'lang': 'py', 'task': 'my_name:my_function', 'id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'shadow': None, 'eta': '2022-05-12T20:38:48.433839+00:00', 'expires': None, 'group': None, 'group_index': None, 'retries': 1, 'timelimit': [None, None], 'root_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'parent_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'argsrepr': "('hello',)", 'kwargsrepr': '{}', 'origin': 'gen24336@DESKTOP-51C03JP', 'ignore_result': False}
[2022-05-12 22:38:48,470: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] received
[2022-05-12 22:38:48,473: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] retry: Retry in 1s: AssertionError()
[2022-05-12 22:38:49,464: WARNING/MainProcess] hello
[2022-05-12 22:38:49,465: WARNING/MainProcess] {'lang': 'py', 'task': 'my_name:my_function', 'id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'shadow': None, 'eta': '2022-05-12T20:38:49.447913+00:00', 'expires': None, 'group': None, 'group_index': None, 'retries': 2, 'timelimit': [None, None], 'root_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'parent_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'argsrepr': "('hello',)", 'kwargsrepr': '{}', 'origin': 'gen24336@DESKTOP-51C03JP', 'ignore_result': False}
[2022-05-12 22:38:49,468: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] received
[2022-05-12 22:38:49,469: WARNING/MainProcess] hello
[2022-05-12 22:38:49,469: INFO/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] retry: Retry in 0s: AssertionError()
[2022-05-12 22:38:49,469: WARNING/MainProcess] {'lang': 'py', 'task': 'my_name:my_function', 'id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'shadow': None, 'eta': None, 'expires': None, 'group': None, 'group_index': None, 'retries': 3, 'timelimit': [None, None], 'root_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'parent_id': 'bb7a8629-00d4-434d-b9d2-72162c2edbbf', 'argsrepr': "('hello',)", 'kwargsrepr': '{}', 'origin': 'gen24336@DESKTOP-51C03JP', 'ignore_result': False}
[2022-05-12 22:38:49,473: ERROR/MainProcess] Task my_name:my_function[bb7a8629-00d4-434d-b9d2-72162c2edbbf] raised unexpected: AssertionError()
Traceback (most recent call last):
  File "C:\Users\liman\Anaconda3\envs\dimensions\lib\site-packages\celery\app\trace.py", line 451, in trace_task
    R = retval = fun(*args, **kwargs)
  File "C:\Users\liman\Anaconda3\envs\dimensions\lib\site-packages\celery\app\trace.py", line 734, in __protected_call__
    return self.run(*args, **kwargs)
  File "C:\Users\liman\Anaconda3\envs\dimensions\lib\site-packages\celery\app\autoretry.py", line 54, in run
    ret = task.retry(exc=exc, **retry_kwargs)
  File "C:\Users\liman\Anaconda3\envs\dimensions\lib\site-packages\celery\app\task.py", line 717, in retry
    raise_with_context(exc)
  File "C:\Users\liman\Anaconda3\envs\dimensions\lib\site-packages\celery\app\autoretry.py", line 34, in run
    return task._orig_run(*args, **kwargs)
  File "C:\Users\liman\Github\Tests\randoms\api.py", line 48, in failing_function
    assert False == True
AssertionError
```

</p>
</details>

# Expected Behavior
Expecting to keep custom headers when retrying the task

# Actual Behavior
Not keeping custom task headers 



Discussion:
Hey @Leem0sh :wave:,
Thank you for opening an issue. We will get back to you as soon as we can.
Also, check out our [Open Collective](https://opencollective.com/celery) and consider backing us - every little helps!

We also offer priority support for our sponsors.
If you require immediate assistance please consider sponsoring us.


Resolving patch:
diff --git a/celery/app/task.py b/celery/app/task.py
--- a/celery/app/task.py
+++ b/celery/app/task.py
@@ -96,6 +96,18 @@ class Context:
 
     def __init__(self, *args, **kwargs):
         self.update(*args, **kwargs)
+        if self.headers is None:
+            self.headers = self._get_custom_headers(*args, **kwargs)
+
+    def _get_custom_headers(self, *args, **kwargs):
+        headers = {}
+        headers.update(*args, **kwargs)
+        celery_keys = {*Context.__dict__.keys(), 'lang', 'task', 'argsrepr', 'kwargsrepr'}
+        for key in celery_keys:
+            headers.pop(key, None)
+        if not headers:
+            return None
+        return headers
 
     def update(self, *args, **kwargs):
         return self.__dict__.update(*args, **kwargs)
diff --git a/t/integration/tasks.py b/t/integration/tasks.py
--- a/t/integration/tasks.py
+++ b/t/integration/tasks.py
@@ -217,6 +217,16 @@ def retry_once_priority(self, *args, expires=60.0, max_retries=1,
                      max_retries=max_retries)
 
 
+@shared_task(bind=True, max_retries=1)
+def retry_once_headers(self, *args, max_retries=1,
+                       countdown=0.1):
+    """Task that fails and is retried. Returns headers."""
+    if self.request.retries:
+        return self.request.headers
+    raise self.retry(countdown=countdown,
+                     max_retries=max_retries)
+
+
 @shared_task
 def redis_echo(message, redis_key="redis-echo"):
     """Task that appends the message to a redis list."""


Test patch:
diff --git a/t/integration/test_tasks.py b/t/integration/test_tasks.py
--- a/t/integration/test_tasks.py
+++ b/t/integration/test_tasks.py
@@ -8,7 +8,8 @@
 
 from .conftest import get_active_redis_channels
 from .tasks import (ClassBasedAutoRetryTask, ExpectedException, add, add_ignore_result, add_not_typed, fail,
-                    print_unicode, retry, retry_once, retry_once_priority, return_properties, sleeping)
+                    print_unicode, retry, retry_once, retry_once_headers, retry_once_priority, return_properties,
+                    sleeping)
 
 TIMEOUT = 10
 
@@ -267,6 +268,13 @@ def test_task_retried_priority(self, manager):
         res = retry_once_priority.apply_async(priority=7)
         assert res.get(timeout=TIMEOUT) == 7  # retried once with priority 7
 
+    @flaky
+    def test_task_retried_headers(self, manager):
+        res = retry_once_headers.apply_async(headers={'x-test-header': 'test-value'})
+        headers = res.get(timeout=TIMEOUT)
+        assert headers is not None  # retried once with headers
+        assert 'x-test-header' in headers  # retry keeps custom headers
+
     @flaky
     def test_unicode_task(self, manager):
         manager.join(
diff --git a/t/unit/tasks/test_context.py b/t/unit/tasks/test_context.py
--- a/t/unit/tasks/test_context.py
+++ b/t/unit/tasks/test_context.py
@@ -63,3 +63,24 @@ def test_context_get(self):
         ctx_dict = get_context_as_dict(ctx, getter=Context.get)
         assert ctx_dict == expected
         assert get_context_as_dict(Context()) == default_context
+
+    def test_extract_headers(self):
+        # Should extract custom headers from the request dict
+        request = {
+            'task': 'test.test_task',
+            'id': 'e16eeaee-1172-49bb-9098-5437a509ffd9',
+            'custom-header': 'custom-value',
+        }
+        ctx = Context(request)
+        assert ctx.headers == {'custom-header': 'custom-value'}
+
+    def test_dont_override_headers(self):
+        # Should not override headers if defined in the request
+        request = {
+            'task': 'test.test_task',
+            'id': 'e16eeaee-1172-49bb-9098-5437a509ffd9',
+            'headers': {'custom-header': 'custom-value'},
+            'custom-header-2': 'custom-value-2',
+        }
+        ctx = Context(request)
+        assert ctx.headers == {'custom-header': 'custom-value'}