Add support for operators that achieve durable execution manually in registry

dev/registry/extract_parameters.py

@@ -386,11 +386,29 @@ def load_resumable_job_mixin() -> type | None:
 
 
 def is_durable_capable(cls: type, resumable_mixin: type | None) -> bool:
-    """Return True if a class fully implements ResumableJobMixin's crash-recovery contract.
-
-    Inheriting the mixin is not sufficient: a complete override is inert unless
-    execute() actually calls execute_resumable().
+    """Return True if a class implements durable/crash-safe execution.
+
+    Two ways to qualify:
+    1. A class-level `__supports_durable_execution = True`
+    declaration (for operators that implement this directly against
+    task_state_store, without ResumableJobMixin -- e.g. KubernetesPodOperator,
+    AgentOperator).
+    2. Genuinely implementing ResumableJobMixin's contract.
+
+    The first path deliberately looks up the class prefixed attribute
+    (`_{ClassName}__supports_durable_execution`) rather than a fixed string.
+    A subclass that overrides execute() itself (e.g. SparkKubernetesOperator)
+    may not preserve the parent's task_state_store reconnect behavior, so the
+    declaration must not be inherited -- only the exact class that wrote
+    `__supports_durable_execution` in its own body qualifies this way.
+
+    Inheriting the mixin alone is not sufficient for the second path: a
+    complete override is inert unless execute() actually calls
+    execute_resumable().
     """
+    if getattr(cls, f"_{cls.__name__}__supports_durable_execution", None) is True:
+        return True
+
     if resumable_mixin is None or resumable_mixin not in cls.__mro__:
         return False
 

dev/registry/tests/test_extract_parameters.py

@@ -225,6 +225,24 @@ def execute(self, context):
         return None
 
 
+class ManuallyDurableOperator:
+    """Implements durable execution directly (e.g. via task_state_store), without
+    ResumableJobMixin -- mirrors KubernetesPodOperator/AgentOperator."""
+
+    __supports_durable_execution = True
+
+    def execute(self, context):
+        return None
+
+
+class ManuallyDurableSubclass(ManuallyDurableOperator):
+    """Overrides execute() itself -- must NOT inherit the parent's declaration,
+    since nothing here verifies it preserves the reconnect behavior."""
+
+    def execute(self, context):
+        return "something else entirely"
+
+
 class TestIsDurableCapable:
     def test_fully_implemented_and_wired_qualifies(self):
         assert is_durable_capable(FullyImplementedResumableOperator, FakeResumableJobMixin) is True
@@ -241,6 +259,12 @@ def test_no_mixin_in_mro_disqualifies(self):
     def test_mixin_unavailable_disqualifies(self):
         assert is_durable_capable(FullyImplementedResumableOperator, None) is False
 
+    def test_manual_durable_marker_qualifies_without_mixin(self):
+        assert is_durable_capable(ManuallyDurableOperator, None) is True
+
+    def test_subclass_not_redeclaring_marker_disqualifies(self):
+        assert is_durable_capable(ManuallyDurableSubclass, FakeResumableJobMixin) is False
+
 
 # ---------------------------------------------------------------------------
 # Module dataclass

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py

@@ -34,7 +34,7 @@
 from contextlib import AbstractContextManager, suppress
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any, ClassVar, Literal
 
 import kubernetes
 import pendulum
@@ -283,6 +283,10 @@ class KubernetesPodOperator(BaseOperator):
     # !!! Changes in KubernetesPodOperator's arguments should be also reflected in !!!
     #  - airflow-core/src/airflow/decorators/__init__.pyi  (by a separate PR)
 
+    # This operator supports durable execution directly, without ResumableJobMixin --
+    # it reconnects via task_state_store on retry instead of resubmitting.
+    __supports_durable_execution: ClassVar[bool] = True
+
     # This field can be overloaded at the instance level via base_container_name
     BASE_CONTAINER_NAME = "base"
     ISTIO_CONTAINER_NAME = "istio-proxy"

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py

@@ -2776,6 +2776,9 @@ def test_reattach_on_restart_via_default_args_reaches_durable(self, dag_maker):
         assert k.durable is False
         assert k.reattach_on_restart is False
 
+    def test_supports_durable_execution_marker(self):
+        assert KubernetesPodOperator._KubernetesPodOperator__supports_durable_execution is True
+
 
 class TestSuppress:
     def test__suppress(self, caplog):

providers/common/ai/src/airflow/providers/common/ai/operators/agent.py

@@ -221,6 +221,10 @@ class AgentOperator(BaseOperator, HITLReviewMixin):
 
     deserialization_allowed_class_fields: ClassVar[tuple[str, ...]] = ("output_type",)
 
+    # This operator supports durable execution directly, without ResumableJobMixin --
+    # it caches step results via task_state_store for replay on retry.
+    __supports_durable_execution: ClassVar[bool] = True
+
     template_fields: Sequence[str] = (
         "prompt",
         "llm_conn_id",

providers/common/ai/tests/unit/common/ai/operators/test_agent.py

@@ -759,6 +759,9 @@ def test_cleanup_skipped_when_post_run_step_fails(self, mock_hook_cls, mock_buil
 
         storage.cleanup.assert_not_called()
 
+    def test_supports_durable_execution_marker(self):
+        assert AgentOperator._AgentOperator__supports_durable_execution is True
+
 
 @pytest.mark.skipif(
     not AIRFLOW_V_3_1_PLUS, reason="Human in the loop is only compatible with Airflow >= 3.1.0"