prefect 3.4.5.dev5__tar.gz → 3.4.6.dev1__tar.gz

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: prefect
-Version: 3.4.5.dev5
+Version: 3.4.6.dev1
 Summary: Workflow orchestration and management.
 Project-URL: Changelog, https://github.com/PrefectHQ/prefect/releases
 Project-URL: Documentation, https://docs.prefect.io
@@ -53,7 +53,7 @@ Requires-Dist: packaging<25.1,>=21.3
 Requires-Dist: pathspec>=0.8.0
 Requires-Dist: pendulum<4,>=3.0.0; python_version < '3.13'
 Requires-Dist: prometheus-client>=0.20.0
-Requires-Dist: pydantic!=2.10.0,<3.0.0,>=2.9
+Requires-Dist: pydantic!=2.10.0,!=2.11.0,!=2.11.1,!=2.11.2,!=2.11.3,!=2.11.4,<3.0.0,>=2.9
 Requires-Dist: pydantic-core<3.0.0,>=2.12.0
 Requires-Dist: pydantic-extra-types<3.0.0,>=2.8.2
 Requires-Dist: pydantic-settings!=2.9.0,<3.0.0,>2.2.1

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/pyproject.toml

@@ -64,7 +64,7 @@ dependencies = [
     "pathspec>=0.8.0",
     "pendulum>=3.0.0,<4; python_version<'3.13'",
     "prometheus-client>=0.20.0",
-    "pydantic>=2.9,<3.0.0,!=2.10.0",
+    "pydantic>=2.9,<3.0.0,!=2.10.0,!=2.11.0,!=2.11.1,!=2.11.2,!=2.11.3,!=2.11.4",
     "pydantic_core>=2.12.0,<3.0.0",
     "pydantic_extra_types>=2.8.2,<3.0.0",
     "pydantic_settings>2.2.1,<3.0.0,!=2.9.0",

prefect-3.4.6.dev1/src/prefect/_build_info.py

@@ -0,0 +1,5 @@
+# Generated by versioningit
+__version__ = "3.4.6.dev1"
+__build_date__ = "2025-06-10 08:09:59.076742+00:00"
+__git_commit__ = "cd3c98b5dbb76efc127c1ee052df55873c1723dc"
+__dirty__ = False

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/assets/core.py

@@ -47,6 +47,9 @@ class Asset(PrefectBaseModel):
     def __repr__(self) -> str:
         return f"Asset(key={self.key!r})"
 
+    def __hash__(self) -> int:
+        return hash(self.key)
+
     def add_metadata(self, metadata: dict[str, Any]) -> None:
         from prefect.context import AssetContext
 
@@ -59,7 +62,7 @@ class Asset(PrefectBaseModel):
         asset_ctx.add_asset_metadata(self.key, metadata)
 
 
-def add_asset_metadata(asset_key: str, metadata: dict[str, Any]) -> None:
+def add_asset_metadata(asset: str | Asset, metadata: dict[str, Any]) -> None:
     from prefect.context import AssetContext
 
     asset_ctx = AssetContext.get()
@@ -68,4 +71,5 @@ def add_asset_metadata(asset_key: str, metadata: dict[str, Any]) -> None:
             "Unable to call `add_asset_metadata` when not inside of an AssetContext"
         )
 
+    asset_key = asset if isinstance(asset, str) else asset.key
     asset_ctx.add_asset_metadata(asset_key, metadata)

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/context.py

@@ -400,7 +400,7 @@ class EngineContext(RunContext):
 
     # Tracking information needed to track asset linage between
     # tasks and materialization
-    task_run_assets: dict[UUID, list[Asset]] = Field(default_factory=dict)
+    task_run_assets: dict[UUID, set[Asset]] = Field(default_factory=dict)
 
     # Events worker to emit events
     events: Optional[EventsWorker] = None
@@ -486,9 +486,9 @@ class AssetContext(ContextModel):
         materialization_metadata: Metadata for materialized assets
     """
 
-    direct_asset_dependencies: list[Asset] = Field(default_factory=list)
-    downstream_assets: list[Asset] = Field(default_factory=list)
-    upstream_assets: list[Asset] = Field(default_factory=list)
+    direct_asset_dependencies: set[Asset] = Field(default_factory=set)
+    downstream_assets: set[Asset] = Field(default_factory=set)
+    upstream_assets: set[Asset] = Field(default_factory=set)
     materialized_by: Optional[str] = None
     task_run_id: Optional[UUID] = None
     materialization_metadata: dict[str, dict[str, Any]] = Field(default_factory=dict)
@@ -516,7 +516,7 @@ class AssetContext(ContextModel):
         from prefect.client.schemas import TaskRunResult
         from prefect.tasks import MaterializingTask
 
-        upstream_assets: list[Asset] = []
+        upstream_assets: set[Asset] = set()
 
         # Get upstream assets from engine context instead of TaskRunResult.assets
         flow_ctx = FlowRunContext.get()
@@ -527,13 +527,15 @@ class AssetContext(ContextModel):
                         # Look up assets in the engine context
                         task_assets = flow_ctx.task_run_assets.get(task_input.id)
                         if task_assets:
-                            upstream_assets.extend(task_assets)
+                            upstream_assets.update(task_assets)
 
         ctx = cls(
-            direct_asset_dependencies=task.asset_deps[:] if task.asset_deps else [],
-            downstream_assets=task.assets[:]
+            direct_asset_dependencies=set(task.asset_deps)
+            if task.asset_deps
+            else set(),
+            downstream_assets=set(task.assets)
             if isinstance(task, MaterializingTask) and task.assets
-            else [],
+            else set(),
             upstream_assets=upstream_assets,
             materialized_by=task.materialized_by
             if isinstance(task, MaterializingTask)
@@ -551,7 +553,15 @@ class AssetContext(ContextModel):
         Args:
             asset_key: The asset key
             metadata: Metadata dictionary to add
+
+        Raises:
+            ValueError: If asset_key is not in downstream_assets
         """
+        downstream_keys = {asset.key for asset in self.downstream_assets}
+        if asset_key not in downstream_keys:
+            raise ValueError(
+                "Can only add metadata to assets that are arguments to @materialize"
+            )
 
         existing = self.materialization_metadata.get(asset_key, {})
         self.materialization_metadata[asset_key] = existing | metadata
@@ -596,46 +606,46 @@ class AssetContext(ContextModel):
 
     def emit_events(self, state: State) -> None:
         """
-        Emit asset reference and materialization events based on task completion.
+        Emit asset events
         """
 
         from prefect.events import emit_event
 
         if state.name == "Cached":
             return
-        if state.is_failed():
+        elif state.is_failed():
             event_status = "failed"
         elif state.is_completed():
             event_status = "succeeded"
         else:
             return
 
-        asset_deps_related: list[Asset] = []
+        # If we have no downstream assets, this not a materialization
+        if not self.downstream_assets:
+            return
 
-        # Emit reference events for direct asset dependencies
-        for asset in self.direct_asset_dependencies:
+        # Emit reference events for all upstream assets (direct + inherited)
+        all_upstream_assets = self.upstream_assets | self.direct_asset_dependencies
+        for asset in all_upstream_assets:
             emit_event(
-                event=f"prefect.asset.reference.{event_status}",
+                event="prefect.asset.referenced",
                 resource=self.asset_as_resource(asset),
                 related=[],
             )
-            asset_deps_related.append(self.asset_as_related(asset))
 
         # Emit materialization events for downstream assets
-        if self.downstream_assets:
-            upstream_related = [self.asset_as_related(a) for a in self.upstream_assets]
-            all_related = upstream_related + asset_deps_related
-
-            if self.materialized_by:
-                all_related.append(self.related_materialized_by(self.materialized_by))
-
-            for asset in self.downstream_assets:
-                emit_event(
-                    event=f"prefect.asset.materialization.{event_status}",
-                    resource=self.asset_as_resource(asset),
-                    related=all_related,
-                    payload=self.materialization_metadata.get(asset.key),
-                )
+        upstream_related = [self.asset_as_related(a) for a in all_upstream_assets]
+
+        if self.materialized_by:
+            upstream_related.append(self.related_materialized_by(self.materialized_by))
+
+        for asset in self.downstream_assets:
+            emit_event(
+                event=f"prefect.asset.materialization.{event_status}",
+                resource=self.asset_as_resource(asset),
+                related=upstream_related,
+                payload=self.materialization_metadata.get(asset.key),
+            )
 
     def update_tracked_assets(self) -> None:
         """
@@ -649,11 +659,11 @@ class AssetContext(ContextModel):
 
         if self.downstream_assets:
             # MaterializingTask: propagate the downstream assets (what we create)
-            assets_for_downstream = self.downstream_assets[:]
+            assets_for_downstream = set(self.downstream_assets)
         else:
             # Regular task: propagate upstream assets + direct dependencies
-            assets_for_downstream = (
-                list(self.upstream_assets) + self.direct_asset_dependencies
+            assets_for_downstream = set(
+                self.upstream_assets | self.direct_asset_dependencies
             )
 
         flow_run_context.task_run_assets[self.task_run_id] = assets_for_downstream
@@ -661,6 +671,9 @@ class AssetContext(ContextModel):
     def serialize(self: Self, include_secrets: bool = True) -> dict[str, Any]:
         """Serialize the AssetContext for distributed execution."""
         return self.model_dump(
+            # use json serialization so fields that are
+            # sets of pydantic models are serialized
+            mode="json",
             exclude_unset=True,
             serialize_as_any=True,
             context={"include_secrets": include_secrets},

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/events/clients.py

@@ -628,7 +628,7 @@ class PrefectEventSubscriber:
         try:
             await self._reconnect()
         finally:
-            EVENT_WEBSOCKET_CONNECTIONS.labels(self.client_name, "out", "initial")
+            EVENT_WEBSOCKET_CONNECTIONS.labels(self.client_name, "out", "initial").inc()
         return self
 
     async def _reconnect(self) -> None:
@@ -709,7 +709,7 @@ class PrefectEventSubscriber:
                     finally:
                         EVENT_WEBSOCKET_CONNECTIONS.labels(
                             self.client_name, "out", "reconnect"
-                        )
+                        ).inc()
                     assert self._websocket
 
                 while True:

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/server/database/alembic.ini

@@ -5,3 +5,5 @@ prepend_sys_path = .
 revision_environment = true
 
 file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d%%(second).2d_%%(rev)s_%%(slug)s
+
+path_separator = os

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/server/database/configurations.py

@@ -427,8 +427,8 @@ class AioSqliteConfiguration(BaseDatabaseConfiguration):
                 )
 
             engine = create_async_engine(self.connection_url, echo=self.echo, **kwargs)
-            sa.event.listen(engine.sync_engine, "connect", self.setup_sqlite)
-            sa.event.listen(engine.sync_engine, "begin", self.begin_sqlite_stmt)
+            event.listen(engine.sync_engine, "connect", self.setup_sqlite)
+            event.listen(engine.sync_engine, "begin", self.begin_sqlite_stmt)
 
             if logfire:
                 logfire.instrument_sqlalchemy(engine)  # pyright: ignore

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/server/database/orm_models.py

@@ -19,6 +19,7 @@ from sqlalchemy.orm import (
     relationship,
     synonym,
 )
+from sqlalchemy.orm.decl_api import registry as RegistryType
 from sqlalchemy.sql import roles
 from sqlalchemy.sql.functions import coalesce
 
@@ -59,7 +60,7 @@ class Base(DeclarativeBase):
     and provides ID, created, and updated columns
     """
 
-    registry: ClassVar[sa.orm.registry] = registry(
+    registry: ClassVar[RegistryType] = registry(
         metadata=sa.schema.MetaData(
             # define naming conventions for our Base class to use
             # sqlalchemy will use the following templated strings
@@ -846,10 +847,8 @@ class Deployment(Base):
     concurrency_limit_id: Mapped[Optional[uuid.UUID]] = mapped_column(
         sa.ForeignKey("concurrency_limit_v2.id", ondelete="SET NULL"),
     )
-    global_concurrency_limit: Mapped[Optional["ConcurrencyLimitV2"]] = (
-        sa.orm.relationship(
-            lazy="selectin",
-        )
+    global_concurrency_limit: Mapped[Optional["ConcurrencyLimitV2"]] = relationship(
+        lazy="selectin",
     )
     concurrency_options: Mapped[Optional[schemas.core.ConcurrencyOptions]] = (
         mapped_column(

{prefect-3.4.5.dev5 → prefect-3.4.6.dev1}/src/prefect/server/orchestration/core_policy.py

@@ -302,8 +302,8 @@ class SecureTaskConcurrencySlots(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         self._applied_limits: list[str] = []
@@ -352,8 +352,8 @@ class SecureTaskConcurrencySlots(TaskRunOrchestrationRule):
 
     async def cleanup(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         for tag in self._applied_limits:
@@ -417,8 +417,8 @@ class SecureFlowConcurrencySlots(FlowRunOrchestrationRule):
 
     async def before_transition(  # type: ignore
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: FlowOrchestrationContext,
     ) -> None:
         if not context.session or not context.run.deployment_id:
@@ -485,8 +485,8 @@ class SecureFlowConcurrencySlots(FlowRunOrchestrationRule):
 
     async def cleanup(  # type: ignore
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: FlowOrchestrationContext,
     ) -> None:
         logger = get_logger()
@@ -579,8 +579,8 @@ class CacheInsertion(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if proposed_state is None:
@@ -598,8 +598,8 @@ class CacheInsertion(TaskRunOrchestrationRule):
     async def after_transition(
         self,
         db: PrefectDBInterface,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if not validated_state or not context.session:
@@ -631,8 +631,8 @@ class CacheRetrieval(TaskRunOrchestrationRule):
     async def before_transition(
         self,
         db: PrefectDBInterface,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if not proposed_state:
@@ -679,8 +679,8 @@ class RetryFailedFlows(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -756,8 +756,8 @@ class RetryFailedTasks(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -803,8 +803,8 @@ class EnqueueScheduledTasks(TaskRunOrchestrationRule):
 
     async def after_transition(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if not validated_state:
@@ -837,8 +837,8 @@ class RenameReruns(GenericOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[
             orm_models.Run, core.TaskRunPolicy | core.FlowRunPolicy
         ],
@@ -869,8 +869,8 @@ class CopyScheduledTime(
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[
             orm_models.Run, core.TaskRunPolicy | core.FlowRunPolicy
         ],
@@ -902,8 +902,8 @@ class WaitForScheduledTime(
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[
             orm_models.Run, core.TaskRunPolicy | core.FlowRunPolicy
         ],
@@ -940,8 +940,8 @@ class CopyTaskParametersID(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -963,8 +963,8 @@ class HandlePausingFlows(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if proposed_state is None:
@@ -1010,8 +1010,8 @@ class HandlePausingFlows(FlowRunOrchestrationRule):
 
     async def after_transition(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         updated_policy = context.run.empirical_policy.model_dump()
@@ -1029,8 +1029,8 @@ class HandleResumingPausedFlows(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -1084,8 +1084,8 @@ class HandleResumingPausedFlows(FlowRunOrchestrationRule):
 
     async def after_transition(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         updated_policy = context.run.empirical_policy.model_dump()
@@ -1103,8 +1103,8 @@ class UpdateFlowRunTrackerOnTasks(TaskRunOrchestrationRule):
 
     async def after_transition(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if context.run.flow_run_id is not None:
@@ -1135,8 +1135,8 @@ class HandleTaskTerminalStateTransitions(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -1176,8 +1176,8 @@ class HandleTaskTerminalStateTransitions(TaskRunOrchestrationRule):
 
     async def cleanup(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         # reset run count
@@ -1200,8 +1200,8 @@ class HandleFlowTerminalStateTransitions(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -1255,8 +1255,8 @@ class HandleFlowTerminalStateTransitions(FlowRunOrchestrationRule):
 
     async def cleanup(
         self,
-        initial_state: states.State | None,
-        validated_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        validated_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         context.run.empirical_policy = core.FlowRunPolicy(**self.original_flow_policy)
@@ -1294,8 +1294,8 @@ class PreventPendingTransitions(GenericOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[
             orm_models.Run, Union[core.FlowRunPolicy, core.TaskRunPolicy]
         ],
@@ -1317,8 +1317,8 @@ class EnsureOnlyScheduledFlowsMarkedLate(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -1346,8 +1346,8 @@ class PreventRunningTasksFromStoppedFlows(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         flow_run = await context.flow_run()
@@ -1390,8 +1390,8 @@ class EnforceCancellingToCancelledTransition(TaskRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.TaskRun, core.TaskRunPolicy],
     ) -> None:
         await self.reject_transition(
@@ -1422,8 +1422,8 @@ class BypassCancellingFlowRunsWithNoInfra(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None:
@@ -1468,8 +1468,8 @@ class PreventDuplicateTransitions(FlowRunOrchestrationRule):
 
     async def before_transition(
         self,
-        initial_state: states.State | None,
-        proposed_state: states.State | None,
+        initial_state: states.State[Any] | None,
+        proposed_state: states.State[Any] | None,
         context: OrchestrationContext[orm_models.FlowRun, core.FlowRunPolicy],
     ) -> None:
         if initial_state is None or proposed_state is None: