Repository: PrefectHQ/prefect

Issue:
`AttributeError: 'NoneType' object has no attribute 'state_details'` after database error during state transition
### 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

This bug was observed in Prefect Cloud

### Reproduction

```python3
n/a
```


### Error

_No response_

### Versions

```Text
n/a
```


### Additional context

_No response_


Discussion:

Resolving patch:
diff --git a/src/prefect/orion/orchestration/core_policy.py b/src/prefect/orion/orchestration/core_policy.py
--- a/src/prefect/orion/orchestration/core_policy.py
+++ b/src/prefect/orion/orchestration/core_policy.py
@@ -192,6 +192,9 @@ async def after_transition(
         context: TaskOrchestrationContext,
         db: OrionDBInterface,
     ) -> None:
+        if not validated_state or not context.session:
+            return
+
         cache_key = validated_state.state_details.cache_key
         if cache_key:
             new_cache_item = db.TaskRunStateCache(


Test patch:
diff --git a/tests/orion/orchestration/test_core_policy.py b/tests/orion/orchestration/test_core_policy.py
--- a/tests/orion/orchestration/test_core_policy.py
+++ b/tests/orion/orchestration/test_core_policy.py
@@ -168,6 +168,41 @@ async def test_set_and_retrieve_unexpired_cached_states(
         assert ctx2.response_status == expected_status
         assert ctx2.validated_state.name == expected_name
 
+    async def test_cache_insertion_requires_validated_state(
+        self,
+        session,
+        initialize_orchestration,
+    ):
+        """Regression test for the bug observed when exiting the CacheInsertion
+        rule after a database error that prevented committing the validated state"""
+        initial_state_type = states.StateType.RUNNING
+        proposed_state_type = states.StateType.COMPLETED
+        intended_transition = (initial_state_type, proposed_state_type)
+        expiration = pendulum.now().subtract(days=1)
+
+        ctx1 = await initialize_orchestration(
+            session,
+            "task",
+            *intended_transition,
+            initial_details={"cache_key": "cache-hit", "cache_expiration": expiration},
+            proposed_details={"cache_key": "cache-hit", "cache_expiration": expiration},
+        )
+
+        with pytest.raises(ValueError, match="this better be mine"):
+            async with contextlib.AsyncExitStack() as stack:
+                ctx1 = await stack.enter_async_context(
+                    CacheInsertion(ctx1, *intended_transition)
+                )
+                # Simulate an exception (for example, a database error) that happens
+                # within the context manager; when this happens, we've observed the
+                # CacheInsertion to raise:
+                #
+                #   AttributeError: 'NoneType' object has no attribute 'state_details'
+                #
+                # because the `validated_state` passed to its after_transition
+                # handler is None
+                raise ValueError("this better be mine")
+
     @pytest.mark.parametrize(
         "proposed_state_type",
         # Include all state types but COMPLETED; cast to sorted list for determinism