diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
index e22d0a5f34d0d..eef96323d2b7d 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py
@@ -23,7 +23,7 @@
 from collections import defaultdict
 from collections.abc import Iterator
 from datetime import datetime
-from typing import TYPE_CHECKING, Annotated, Any
+from typing import TYPE_CHECKING, Annotated, Any, cast
 from uuid import UUID
 
 import attrs
@@ -73,6 +73,7 @@
 if TYPE_CHECKING:
     from sqlalchemy.sql.dml import Update
 
+    from airflow.models.expandinput import SchedulerExpandInput
     from airflow.sdk.types import Operator
 
 
@@ -308,9 +309,9 @@ def _get_upstream_map_indexes(
                 mapped_ti_count = upstream_mapped_group.get_parse_time_mapped_ti_count()
             except NotFullyPopulated:
                 # for cases that needs to resolve xcom to get the correct count
-                mapped_ti_count = upstream_mapped_group._expand_input.get_total_map_length(
-                    run_id, session=session
-                )
+                mapped_ti_count = cast(
+                    "SchedulerExpandInput", upstream_mapped_group._expand_input
+                ).get_total_map_length(run_id, session=session)
             map_indexes = list(range(mapped_ti_count)) if mapped_ti_count is not None else None
 
         yield upstream_task.task_id, map_indexes