Repository: PrefectHQ/prefect

Issue:
Task run context doesn't seems to be updated on retries
### First check

- [X] I added a descriptive title to this issue.
- [X] I used the GitHub search to find a similar issue and didn't find it.
- [X] I searched the Prefect documentation for this issue.
- [X] I checked that this issue is related to Prefect and not one of its dependencies.

### Bug summary

When retrying a task, the task run's attributes such as `state_name` and `run_count` don't get updated

The code below always prints run count 0 and state name "Pending"

### Reproduction

```python
import prefect
from prefect import task, flow


@task(retries=3, retry_delay_seconds=1)
def always_failing_task():
    # Run count is always 0 and state is always Pending - why?
    print(f"Run count: {prefect.context.get_run_context().task_run.run_count}")
    print(f"State name: {prefect.context.get_run_context().task_run.state_name}")
    return 1 / 0


@flow
def flow_reacting_to_retries():
    always_failing_task.submit()


if __name__ == "__main__":
    flow_reacting_to_retries()
```


### Error

After 3 retries:
```
ZeroDivisionError: division by zero
Run count: 0
State name: Pending
15:24:19.496 | ERROR   | Task run 'always_failing_task-f278e9ce-0' - Finished in state Failed('Task run encountered an exception.')
15:24:20.008 | ERROR   | Flow run 'ambrosial-koel' - Finished in state Failed('1/1 states failed.')
```


### Versions

```
Version:             2.3.2
API version:         0.8.0
Python version:      3.9.12
Git commit:          6e931ee9
Built:               Tue, Sep 6, 2022 12:36 PM
OS/Arch:             darwin/arm64
Profile:             dev
Server type:         cloud
```


### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/src/prefect/context.py b/src/prefect/context.py
--- a/src/prefect/context.py
+++ b/src/prefect/context.py
@@ -90,7 +90,7 @@ def get(cls: Type[T]) -> Optional[T]:
         return cls.__var__.get(None)
 
     def copy(self, **kwargs):
-        """Remove the token on copy to avoid re-entrance errors"""
+        # Remove the token on copy to avoid re-entrance errors
         new = super().copy(**kwargs)
         new._token = None
         return new
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -1191,6 +1191,9 @@ async def orchestrate_task_run(
 
     # Only run the task if we enter a `RUNNING` state
     while state.is_running():
+        # Retrieve the latest metadata for the task run context
+        task_run = await client.read_task_run(task_run.id)
+
         try:
             args, kwargs = parameters_to_args_kwargs(task.fn, resolved_parameters)
 
@@ -1199,7 +1202,9 @@ async def orchestrate_task_run(
             else:
                 logger.debug(f"Beginning execution...", extra={"state_message": True})
 
-            with task_run_context:
+            with task_run_context.copy(
+                update={"task_run": task_run, "start_time": pendulum.now("UTC")}
+            ):
                 if task.isasync:
                     result = await task.fn(*args, **kwargs)
                 else:


Test patch:
diff --git a/tests/test_tasks.py b/tests/test_tasks.py
--- a/tests/test_tasks.py
+++ b/tests/test_tasks.py
@@ -9,7 +9,7 @@
 import pytest
 
 from prefect import flow, get_run_logger, tags
-from prefect.context import PrefectObjectRegistry
+from prefect.context import PrefectObjectRegistry, TaskRunContext, get_run_context
 from prefect.engine import get_state_for_result
 from prefect.exceptions import (
     MappingLengthMismatch,
@@ -623,6 +623,38 @@ def test_flow():
             "Completed",
         ]
 
+    async def test_task_retries_receive_latest_task_run_in_context(self):
+        contexts: List[TaskRunContext] = []
+
+        @task(retries=3)
+        def flaky_function():
+            contexts.append(get_run_context())
+            raise ValueError()
+
+        @flow
+        def test_flow():
+            flaky_function()
+
+        with pytest.raises(ValueError):
+            test_flow()
+
+        expected_state_names = [
+            "Running",
+            "Retrying",
+            "Retrying",
+            "Retrying",
+        ]
+        assert len(contexts) == len(expected_state_names)
+        for i, context in enumerate(contexts):
+            assert context.task_run.run_count == i + 1
+            assert context.task_run.state_name == expected_state_names[i]
+
+            if i > 0:
+                last_context = contexts[i - 1]
+                assert (
+                    last_context.start_time < context.start_time
+                ), "Timestamps should be increasing"
+
 
 class TestTaskCaching:
     def test_repeated_task_call_within_flow_is_not_cached_by_default(self):