Fix failing main (#20871)

When I merged #18724 the jobs ran successfully but it's now failing in main.
This PR fixes it

Co-authored-by: Kaxil Naik <kaxilnaik@gmail.com>
Co-authored-by: Ash Berlin-Taylor <ash_github@firemirror.com>
GitOrigin-RevId: b1fffe7b2673f5aaad72304d999f4a55b2cc2edc

airflow/api_connexion/endpoints/task_instance_endpoint.py

@@ -326,7 +326,9 @@ def post_set_task_instances_state(*, dag_id: str, session: Session = NEW_SESSION
             detail=f"Task instance not found for task {task_id!r} on execution_date {execution_date}"
         )
 
-    if run_id and not session.query(TI).get({'task_id': task_id, 'dag_id': dag_id, 'run_id': run_id}):
+    if run_id and not session.query(TI).get(
+        {'task_id': task_id, 'dag_id': dag_id, 'run_id': run_id, 'map_index': -1}
+    ):
         error_message = f"Task instance not found for task {task_id!r} on DAG run with ID {run_id!r}"
         raise NotFound(detail=error_message)