diff --git a/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py b/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
index 4435380644a8c..0536e6a0bf538 100644
--- a/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
+++ b/providers/common/ai/src/airflow/providers/common/ai/plugins/hitl_review.py
@@ -79,8 +79,8 @@ def _get_bundle_url() -> str:
         HumanFeedbackRequest,
         SessionStatus,
     )
+    from airflow.sdk import TaskInstanceState
     from airflow.utils.session import create_session
-    from airflow.utils.state import TaskInstanceState
 
     def _get_session():
         with create_session(scoped=False) as session: