diff --git a/airflow-core/newsfragments/60803.significant.rst b/airflow-core/newsfragments/60803.significant.rst
new file mode 100644
index 0000000000000..1054258ad3852
--- /dev/null
+++ b/airflow-core/newsfragments/60803.significant.rst
@@ -0,0 +1 @@
+Move ``upstream_map_indexes`` computation from API server to Task SDK, reducing memory usage on task start by eliminating ``SerializedDAG`` loads.
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
index e7ebee9ebe790..513a99f6dc996 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py
@@ -324,8 +324,6 @@ class TIRunContext(BaseModel):
     connections: Annotated[list[ConnectionResponse], Field(default_factory=list)]
     """Connections that can be accessed by the task instance."""
 
-    upstream_map_indexes: dict[str, int | list[int] | None] | None = None
-
     next_method: str | None = None
     """Method to call. Set when task resumes from a trigger."""
     next_kwargs: dict[str, Any] | str | None = None
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 a729a1ee83b30..a73145b30ab5a 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
@@ -31,7 +31,7 @@
 from fastapi import Body, HTTPException, Query, status
 from pydantic import JsonValue
 from sqlalchemy import func, or_, tuple_, update
-from sqlalchemy.engine import CursorResult, Row
+from sqlalchemy.engine import CursorResult
 from sqlalchemy.exc import NoResultFound, SQLAlchemyError
 from sqlalchemy.orm import joinedload
 from sqlalchemy.sql import select
@@ -64,14 +64,11 @@
 from airflow.models.asset import AssetActive
 from airflow.models.dag import DagModel
 from airflow.models.dagrun import DagRun as DR
-from airflow.models.expandinput import NotFullyPopulated
 from airflow.models.taskinstance import TaskInstance as TI, _stop_remaining_tasks
 from airflow.models.taskreschedule import TaskReschedule
 from airflow.models.trigger import Trigger
 from airflow.models.xcom import XComModel
 from airflow.serialization.definitions.assets import SerializedAsset, SerializedAssetUniqueKey
-from airflow.serialization.definitions.dag import SerializedDAG
-from airflow.task.trigger_rule import TriggerRule
 from airflow.utils.sqlalchemy import get_dialect_name
 from airflow.utils.state import DagRunState, TaskInstanceState, TerminalTIState
 
@@ -251,17 +248,6 @@ def ti_run(
             or 0
         )
 
-        if dag := dag_bag.get_dag_for_run(dag_run=dr, session=session):
-            upstream_map_indexes = dict(
-                _get_upstream_map_indexes(
-                    serialized_dag=dag,
-                    ti=ti,
-                    session=session,
-                )
-            )
-        else:
-            upstream_map_indexes = None
-
         context = TIRunContext(
             dag_run=dr,
             task_reschedule_count=task_reschedule_count,
@@ -271,7 +257,6 @@ def ti_run(
             connections=[],
             xcom_keys_to_clear=xcom_keys,
             should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries),
-            upstream_map_indexes=upstream_map_indexes,
         )
 
         # Only set if they are non-null
@@ -287,47 +272,6 @@ def ti_run(
         )
 
 
-def _get_upstream_map_indexes(
-    *,
-    serialized_dag: SerializedDAG,
-    ti: TI | Row,
-    session: SessionDep,
-) -> Iterator[tuple[str, int | list[int] | None]]:
-    task = serialized_dag.get_task(ti.task_id)
-    for upstream_task in task.upstream_list:
-        map_indexes: int | list[int] | None
-        if (upstream_mapped_group := upstream_task.get_closest_mapped_task_group()) is None:
-            # regular tasks or non-mapped task groups
-            map_indexes = None
-        elif task.get_closest_mapped_task_group() is upstream_mapped_group:
-            # tasks in the same mapped task group hierarchy
-            map_indexes = ti.map_index
-        else:
-            # tasks not in the same mapped task group
-            # the upstream mapped task group should combine the return xcom as a list and return it
-            mapped_ti_count: int | None = None
-
-            try:
-                # First try: without resolving XCom
-                mapped_ti_count = upstream_mapped_group.get_parse_time_mapped_ti_count()
-            except NotFullyPopulated:
-                # Second try: resolve XCom for correct count
-                try:
-                    expand_input = upstream_mapped_group._expand_input
-                    mapped_ti_count = expand_input.get_total_map_length(ti.run_id, session=session)
-                except NotFullyPopulated:
-                    # For these trigger rules, unresolved map indexes are acceptable.
-                    # The success of the upstream task is not the main reason for triggering the current task.
-                    # Therefore, whether the upstream task is fully populated can be ignored.
-                    if task.trigger_rule != TriggerRule.ALL_SUCCESS:
-                        mapped_ti_count = None
-
-            # Compute map indexes if we have a valid count
-            map_indexes = list(range(mapped_ti_count)) if mapped_ti_count is not None else None
-
-        yield upstream_task.task_id, map_indexes
-
-
 @ti_id_router.patch(
     "/{task_instance_id}/state",
     status_code=status.HTTP_204_NO_CONTENT,
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index 36c1b31b959bf..30d4159f7453c 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -33,11 +33,14 @@
     AddDagRunDetailEndpoint,
     MovePreviousRunEndpoint,
 )
-from airflow.api_fastapi.execution_api.versions.v2026_03_31 import ModifyDeferredTaskKwargsToJsonValue
+from airflow.api_fastapi.execution_api.versions.v2026_03_31 import (
+    ModifyDeferredTaskKwargsToJsonValue,
+    RemoveUpstreamMapIndexesField,
+)
 
 bundle = VersionBundle(
     HeadVersion(),
-    Version("2026-03-31", ModifyDeferredTaskKwargsToJsonValue),
+    Version("2026-03-31", ModifyDeferredTaskKwargsToJsonValue, RemoveUpstreamMapIndexesField),
     Version("2025-12-08", MovePreviousRunEndpoint, AddDagRunDetailEndpoint),
     Version("2025-11-07", AddPartitionKeyField),
     Version("2025-11-05", AddTriggeringUserNameField),
diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py
index 48630e0b50c09..72e193426dac6 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_03_31.py
@@ -19,9 +19,9 @@
 
 from typing import Any
 
-from cadwyn import VersionChange, schema
+from cadwyn import ResponseInfo, VersionChange, convert_response_to_previous_version_for, schema
 
-from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIDeferredStatePayload
+from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIDeferredStatePayload, TIRunContext
 
 
 class ModifyDeferredTaskKwargsToJsonValue(VersionChange):
@@ -33,3 +33,20 @@ class ModifyDeferredTaskKwargsToJsonValue(VersionChange):
         schema(TIDeferredStatePayload).field("trigger_kwargs").had(type=dict[str, Any] | str),
         schema(TIDeferredStatePayload).field("next_kwargs").had(type=dict[str, Any]),
     )
+
+
+class RemoveUpstreamMapIndexesField(VersionChange):
+    """Remove upstream_map_indexes field from TIRunContext - now computed by Task SDK."""
+
+    description = __doc__
+
+    instructions_to_migrate_to_previous_version = (
+        schema(TIRunContext)
+        .field("upstream_map_indexes")
+        .existed_as(type=dict[str, int | list[int] | None] | None),
+    )
+
+    @convert_response_to_previous_version_for(TIRunContext)  # type: ignore[arg-type]
+    def add_upstream_map_indexes_field(response: ResponseInfo) -> None:  # type: ignore[misc]
+        """Add upstream_map_indexes field with None for older API versions."""
+        response.body["upstream_map_indexes"] = None
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
index e99b7d9380f03..cfee6c9d46cb1 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py
@@ -202,7 +202,6 @@ def test_ti_run_state_to_running(
                 "partition_key": None,
             },
             "task_reschedule_count": 0,
-            "upstream_map_indexes": {},
             "max_tries": max_tries,
             "should_retry": should_retry,
             "variables": [],
@@ -249,10 +248,7 @@ def test_ti_run_state_to_running(
         assert response.status_code == 409
 
     def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker):
-        """
-        Test that the Task Instance upstream_map_indexes is correctly fetched when to running the Task Instances
-        """
-
+        """Test that dynamic task mapping works correctly with parse-time values."""
         with dag_maker("test_dynamic_task_mapping_with_parse_time_value", serialized=True):
 
             @task_group
@@ -278,23 +274,6 @@ def task2():
             ti.set_state(State.QUEUED)
         dag_maker.session.flush()
 
-        # key: (task_id, map_index)
-        # value: result upstream_map_indexes ({task_id: map_indexes})
-        expected_upstream_map_indexes = {
-            # no upstream task for task_group_1.group_task_1
-            ("task_group_1.group1_task_1", 0): {},
-            ("task_group_1.group1_task_1", 1): {},
-            # the upstream task for task_group_1.group_task_2 is task_group_1.group_task_2
-            # since they are in the same task group, the upstream map index should be the same as the task
-            ("task_group_1.group1_task_2", 0): {"task_group_1.group1_task_1": 0},
-            ("task_group_1.group1_task_2", 1): {"task_group_1.group1_task_1": 1},
-            # the upstream task for task2 is the last tasks of task_group_1, which is
-            # task_group_1.group_task_2
-            # since they are not in the same task group, the upstream map index should include all the
-            # expanded tasks
-            ("task2", -1): {"task_group_1.group1_task_2": [0, 1]},
-        }
-
         for ti in dr.get_task_instances():
             response = client.patch(
                 f"/execution/task-instances/{ti.id}/run",
@@ -308,13 +287,9 @@ def task2():
             )
 
             assert response.status_code == 200
-            upstream_map_indexes = response.json()["upstream_map_indexes"]
-            assert upstream_map_indexes == expected_upstream_map_indexes[(ti.task_id, ti.map_index)]
 
-    def test_nested_mapped_task_group_upstream_indexes(self, client, dag_maker):
-        """
-        Test that upstream_map_indexes are correctly computed for tasks in nested mapped task groups.
-        """
+    def test_nested_mapped_task_group(self, client, dag_maker):
+        """Test that nested mapped task groups work correctly."""
         with dag_maker("test_nested_mapped_tg", serialized=True):
 
             @task
@@ -346,25 +321,11 @@ def expandable_task_group(param: str) -> None:
                 ti.set_state(State.QUEUED)
         dag_maker.session.flush()
 
-        # Expected upstream_map_indexes for each print_task instance
-        expected_upstream_map_indexes = {
-            ("expandable_task_group.inner_task_group.print_task", 0): {
-                "expandable_task_group.inner_task_group.alter_input": 0
-            },
-            ("expandable_task_group.inner_task_group.print_task", 1): {
-                "expandable_task_group.inner_task_group.alter_input": 1
-            },
-            ("expandable_task_group.inner_task_group.print_task", 2): {
-                "expandable_task_group.inner_task_group.alter_input": 2
-            },
-        }
-
         # Get only the expanded print_task instances (not the template)
         print_task_tis = [
             ti for ti in dr.get_task_instances() if "print_task" in ti.task_id and ti.map_index >= 0
         ]
 
-        # Test each print_task instance
         for ti in print_task_tis:
             response = client.patch(
                 f"/execution/task-instances/{ti.id}/run",
@@ -378,18 +339,9 @@ def expandable_task_group(param: str) -> None:
             )
 
             assert response.status_code == 200
-            upstream_map_indexes = response.json()["upstream_map_indexes"]
-            expected = expected_upstream_map_indexes[(ti.task_id, ti.map_index)]
-
-            assert upstream_map_indexes == expected, (
-                f"Task {ti.task_id}[{ti.map_index}] should have upstream_map_indexes {expected}, "
-                f"but got {upstream_map_indexes}"
-            )
 
     def test_dynamic_task_mapping_with_xcom(self, client: Client, dag_maker: DagMaker, session: Session):
-        """
-        Test that the Task Instance upstream_map_indexes is correctly fetched when to running the Task Instances with xcom
-        """
+        """Test that dynamic task mapping works correctly with XCom values."""
         from airflow.models.taskmap import TaskMap
 
         with dag_maker(session=session, serialized=True):
@@ -442,13 +394,10 @@ def task_3():
                 "start_date": "2024-09-30T12:00:00Z",
             },
         )
-        assert response.json()["upstream_map_indexes"] == {"tg.task_2": [0, 1, 2, 3, 4, 5]}
+        assert response.status_code == 200
 
     def test_dynamic_task_mapping_with_all_success_trigger_rule(self, dag_maker: DagMaker, session: Session):
-        """
-        Test that the Task Instance upstream_map_indexes is not populuated but
-        the downstream task should not be run.
-        """
+        """Test that with ALL_SUCCESS trigger rule and skipped upstream, downstream should not run."""
 
         with dag_maker(session=session, serialized=True):
 
@@ -504,10 +453,7 @@ def task_4():
     def test_dynamic_task_mapping_with_non_all_success_trigger_rule(
         self, client: Client, dag_maker: DagMaker, session: Session, trigger_rule: TriggerRule
     ):
-        """
-        Test that the Task Instance upstream_map_indexes is not populuated but
-        the downstream task should still be run due to trigger rule.
-        """
+        """Test that with non-ALL_SUCCESS trigger rule, downstream task should still run."""
 
         with dag_maker(session=session, serialized=True):
 
@@ -564,7 +510,7 @@ def task_4():
                 "start_date": "2024-09-30T12:00:00Z",
             },
         )
-        assert response.json()["upstream_map_indexes"] == {"tg.task_2": None}
+        assert response.status_code == 200
 
     def test_next_kwargs_still_encoded(self, client, session, create_task_instance, time_machine):
         instant_str = "2024-09-30T12:00:00Z"
@@ -615,7 +561,6 @@ def test_next_kwargs_still_encoded(self, client, session, create_task_instance,
         assert response.json() == {
             "dag_run": mock.ANY,
             "task_reschedule_count": 0,
-            "upstream_map_indexes": {},
             "max_tries": 0,
             "should_retry": False,
             "variables": [],
@@ -687,7 +632,6 @@ def test_next_kwargs_determines_start_date_update(self, client, session, create_
         assert response.json() == {
             "dag_run": mock.ANY,
             "task_reschedule_count": 0,
-            "upstream_map_indexes": {},
             "max_tries": 0,
             "should_retry": False,
             "variables": [],
diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
index 16371b67d0453..efa29338d9251 100644
--- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
+++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2025_04_28/test_task_instances.py
@@ -17,8 +17,6 @@
 
 from __future__ import annotations
 
-from unittest.mock import patch
-
 import pytest
 
 from airflow._shared.timezones import timezone
@@ -50,52 +48,19 @@ def teardown_method(self):
         clear_db_assets()
         clear_db_runs()
 
-    @pytest.mark.parametrize(
-        ("mock_indexes", "expected_response_indexes"),
-        [
-            pytest.param(
-                [("task_a", 5), ("task_b", 10)],
-                {"task_a": 5, "task_b": 10},
-                id="plain ints",
-            ),
-            pytest.param(
-                [("task_a", [3, 4]), ("task_b", [9])],
-                {"task_a": 3, "task_b": 9},
-                id="list of ints",
-            ),
-            pytest.param(
-                [
-                    ("task_a", None),
-                ],
-                {"task_a": None},
-                id="task has no upstreams",
-            ),
-            pytest.param(
-                [("task_a", None), ("task_b", [6, 7]), ("task_c", 2)],
-                {"task_a": None, "task_b": 6, "task_c": 2},
-                id="mixed types",
-            ),
-        ],
-    )
-    @patch("airflow.api_fastapi.execution_api.routes.task_instances._get_upstream_map_indexes")
     def test_ti_run(
         self,
-        mock_get_upstream_map_indexes,
         ver_client,
         session,
         create_task_instance,
         time_machine,
-        mock_indexes,
-        expected_response_indexes,
         get_execution_app,
     ):
         """
         Test that this version of the endpoint works.
 
-        Later versions modified the type of upstream_map_indexes.
+        upstream_map_indexes is now always None as it's computed by the Task SDK.
         """
-        mock_get_upstream_map_indexes.return_value = mock_indexes
-
         instant_str = "2024-09-30T12:00:00Z"
         instant = timezone.parse(instant_str)
         time_machine.move_to(instant, tick=False)
@@ -124,26 +89,10 @@ def test_ti_run(
         )
 
         assert response.status_code == 200
-        assert response.json() == {
-            "dag_run": {
-                "dag_id": "dag",
-                "run_id": "test",
-                "clear_number": 0,
-                "logical_date": instant_str,
-                "data_interval_start": instant.subtract(days=1).to_iso8601_string(),
-                "data_interval_end": instant_str,
-                "run_after": instant_str,
-                "start_date": instant_str,
-                "end_date": None,
-                "run_type": "manual",
-                "conf": {},
-                "consumed_asset_events": [],
-            },
-            "task_reschedule_count": 0,
-            "upstream_map_indexes": expected_response_indexes,
-            "max_tries": 0,
-            "should_retry": False,
-            "variables": [],
-            "connections": [],
-            "xcom_keys_to_clear": [],
-        }
+        result = response.json()
+        # upstream_map_indexes is now computed by SDK, server returns None
+        assert result["upstream_map_indexes"] is None
+        assert result["dag_run"]["dag_id"] == "dag"
+        assert result["task_reschedule_count"] == 0
+        assert result["max_tries"] == 0
+        assert result["should_retry"] is False
diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py
index bc18c94f17bff..258d3345706cc 100644
--- a/devel-common/src/tests_common/pytest_plugin.py
+++ b/devel-common/src/tests_common/pytest_plugin.py
@@ -2457,7 +2457,6 @@ def _create_task_instance(
         run_type: str = "manual",
         try_number: int = 1,
         map_index: int | None = -1,
-        upstream_map_indexes: dict[str, int | list[int] | None] | None = None,
         task_reschedule_count: int = 0,
         ti_id: UUID | None = None,
         conf: dict[str, Any] | None = None,
@@ -2532,12 +2531,8 @@ def _create_task_instance(
             task_reschedule_count=task_reschedule_count,
             max_tries=task_retries if max_tries is None else max_tries,
             should_retry=should_retry if should_retry is not None else try_number <= task_retries,
-            upstream_map_indexes=upstream_map_indexes,
         )
 
-        if upstream_map_indexes is not None:
-            ti_context.upstream_map_indexes = upstream_map_indexes
-
         compat_fields = {
             "requests_fd": 0,
             "sentry_integration": "",
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index 6a3a07f5e8a92..149aa5d7dcb0d 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -631,9 +631,6 @@ class TIRunContext(BaseModel):
     max_tries: Annotated[int, Field(title="Max Tries")]
     variables: Annotated[list[VariableResponse] | None, Field(title="Variables")] = None
     connections: Annotated[list[ConnectionResponse] | None, Field(title="Connections")] = None
-    upstream_map_indexes: Annotated[
-        dict[str, int | list[int] | None] | None, Field(title="Upstream Map Indexes")
-    ] = None
     next_method: Annotated[str | None, Field(title="Next Method")] = None
     next_kwargs: Annotated[dict[str, Any] | str | None, Field(title="Next Kwargs")] = None
     xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
index fa56ae8465040..b6ffbd2214253 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
@@ -189,7 +189,9 @@ def resolve(self, context: Mapping[str, Any]) -> tuple[Mapping[str, Any], set[in
         if map_index is None or map_index < 0:
             raise RuntimeError("can't resolve task-mapping argument without expanding")
 
-        upstream_map_indexes = getattr(context["ti"], "_upstream_map_indexes", {})
+        # Get pre-computed upstream_map_indexes if available, otherwise default to empty dict.
+        # When empty, individual XComArgs will compute their map_indexes lazily in xcom_arg.py.
+        upstream_map_indexes = getattr(context["ti"], "_upstream_map_indexes", None) or {}
 
         # TODO: This initiates one API call for each XComArg. Would it be
         # more efficient to do one single call and unpack the value here?
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index ef316574df550..88db900924289 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -30,7 +30,7 @@
 from airflow.sdk.definitions._internal.abstractoperator import AbstractOperator
 from airflow.sdk.definitions._internal.mixins import DependencyMixin, ResolveMixin
 from airflow.sdk.definitions._internal.setup_teardown import SetupTeardownContext
-from airflow.sdk.definitions._internal.types import NOTSET, is_arg_set
+from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet, is_arg_set
 from airflow.sdk.exceptions import AirflowException, XComNotFound
 from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
 from airflow.sdk.execution_time.xcom import BaseXCom
@@ -337,10 +337,25 @@ def resolve(self, context: Mapping[str, Any]) -> Any:
             return LazyXComSequence(xcom_arg=self, ti=ti)
         tg = self.operator.get_closest_mapped_task_group()
         if tg is None:
-            map_indexes = None
+            # No mapped task group - pull from unmapped instance
+            map_indexes: int | range | None | ArgNotSet = None
         else:
-            upstream_map_indexes = getattr(ti, "_upstream_map_indexes", {})
-            map_indexes = upstream_map_indexes.get(task_id, None)
+            # Check for pre-computed value from server (backward compatibility)
+            upstream_map_indexes = getattr(ti, "_upstream_map_indexes", None)
+            if upstream_map_indexes is not None:
+                # Use None as default to match original behavior (filter for unmapped XCom)
+                map_indexes = upstream_map_indexes.get(task_id, None)
+            else:
+                # Compute lazily - ti_count will be queried if needed
+                cached_context = getattr(ti, "_cached_template_context", None)
+                ti_count = cached_context.get("expanded_ti_count") if cached_context else None
+                computed = ti.get_relevant_upstream_map_indexes(
+                    upstream=self.operator,
+                    ti_count=ti_count,
+                    session=None,  # Not used in SDK implementation
+                )
+                # None means "no filtering needed" -> use NOTSET to pull all values
+                map_indexes = NOTSET if computed is None else computed
         result = ti.xcom_pull(
             task_ids=task_id,
             key=self.key,
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_mapping.py b/task-sdk/src/airflow/sdk/execution_time/task_mapping.py
new file mode 100644
index 0000000000000..fc54bb794d9ac
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/execution_time/task_mapping.py
@@ -0,0 +1,133 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Utility functions for computing upstream map indexes in the Task SDK."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.sdk.execution_time.comms import GetTICount, TICount
+
+if TYPE_CHECKING:
+    from airflow.sdk import BaseOperator
+    from airflow.sdk.definitions.taskgroup import MappedTaskGroup, TaskGroup
+
+
+def _find_common_ancestor_mapped_group(node1: BaseOperator, node2: BaseOperator) -> MappedTaskGroup | None:
+    """
+    Given two operators, find their innermost common mapped task group.
+
+    :param node1: First operator
+    :param node2: Second operator
+    :return: The common mapped task group, or None if they don't share one
+    """
+    try:
+        dag1 = node1.dag
+        dag2 = node2.dag
+    except RuntimeError:
+        # Operator not assigned to a DAG
+        return None
+
+    if dag1 is None or dag2 is None or node1.dag_id != node2.dag_id:
+        return None
+    parent_group_ids = {g.group_id for g in node1.iter_mapped_task_groups()}
+    common_groups = (g for g in node2.iter_mapped_task_groups() if g.group_id in parent_group_ids)
+    return next(common_groups, None)
+
+
+def _is_further_mapped_inside(operator: BaseOperator, container: TaskGroup) -> bool:
+    """
+    Whether given operator is *further* mapped inside a task group.
+
+    :param operator: The operator to check
+    :param container: The container task group
+    :return: True if the operator is further mapped inside the container
+    """
+    # Use getattr for compatibility with both SDK and serialized operators
+    if getattr(operator, "is_mapped", False):
+        return True
+    task_group = operator.task_group
+    while task_group is not None and task_group.group_id != container.group_id:
+        if getattr(task_group, "is_mapped", False):
+            return True
+        task_group = task_group.parent_group
+    return False
+
+
+def get_ti_count_for_task(task_id: str, dag_id: str, run_id: str) -> int:
+    """
+    Query TI count for a specific task.
+
+    :param task_id: The task ID
+    :param dag_id: The DAG ID
+    :param run_id: The run ID
+    :return: The count of task instances for the task
+    """
+    # Import here because SUPERVISOR_COMMS is set at runtime, not import time
+    from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS
+
+    response = SUPERVISOR_COMMS.send(GetTICount(dag_id=dag_id, task_ids=[task_id], run_ids=[run_id]))
+    if not isinstance(response, TICount):
+        raise RuntimeError(f"Unexpected response type: {type(response)}")
+    return response.count
+
+
+def get_relevant_map_indexes(
+    task: BaseOperator,
+    run_id: str,
+    map_index: int,
+    ti_count: int,
+    relative: BaseOperator,
+    dag_id: str,
+) -> int | range | None:
+    """
+    Determine map indexes for XCom aggregation.
+
+    This is used to figure out which specific map indexes of an upstream task
+    are relevant when resolving XCom values for a task in a mapped task group.
+
+    :param task: The current task
+    :param run_id: The current run ID
+    :param map_index: The map index of the current task instance
+    :param ti_count: The total count of task instances for the current task
+    :param relative: The upstream/downstream task to find relevant map indexes for
+    :param dag_id: The DAG ID
+    :return: None (use entire value), int (single index), or range (subset of indexes)
+    """
+    if not ti_count:
+        return None
+
+    common_ancestor = _find_common_ancestor_mapped_group(task, relative)
+    if common_ancestor is None or common_ancestor.group_id is None:
+        return None  # Different mapping contexts → use whole value
+
+    # Query TI count using the current task, which is in the mapped task group.
+    # This gives us the number of expansion iterations, not total TIs in the group.
+    ancestor_ti_count = get_ti_count_for_task(task.task_id, dag_id, run_id)
+    if not ancestor_ti_count:
+        return None
+
+    ancestor_map_index = map_index * ancestor_ti_count // ti_count
+
+    if not _is_further_mapped_inside(relative, common_ancestor):
+        return ancestor_map_index  # Single index
+
+    # Partial aggregation for selected TIs
+    further_count = ti_count // ancestor_ti_count
+    map_index_start = ancestor_map_index * further_count
+    return range(map_index_start, map_index_start + further_count)
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index 6057706699889..1cfc2c82d4569 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -268,10 +268,10 @@ def get_template_context(self) -> Context:
                     }
                 )
 
-            if from_server.upstream_map_indexes is not None:
-                # We stash this in here for later use, but we purposefully don't want to document it's
-                # existence. Should this be a private attribute on RuntimeTI instead perhaps?
-                setattr(self, "_upstream_map_indexes", from_server.upstream_map_indexes)
+            # Backward compatibility: old servers may still send upstream_map_indexes
+            upstream_map_indexes = getattr(from_server, "upstream_map_indexes", None)
+            if upstream_map_indexes is not None:
+                setattr(self, "_upstream_map_indexes", upstream_map_indexes)
 
         return self._cached_template_context
 
@@ -436,8 +436,35 @@ def xcom_push(self, key: str, value: Any):
     def get_relevant_upstream_map_indexes(
         self, upstream: BaseOperator, ti_count: int | None, session: Any
     ) -> int | range | None:
-        # TODO: Implement this method
-        return None
+        """
+        Compute the relevant upstream map indexes for XCom resolution.
+
+        :param upstream: The upstream operator
+        :param ti_count: The total count of task instances for this task's expansion
+        :param session: Not used (kept for API compatibility)
+        :return: None (use entire value), int (single index), or range (subset of indexes)
+        """
+        from airflow.sdk.execution_time.task_mapping import get_relevant_map_indexes, get_ti_count_for_task
+
+        map_index = self.map_index
+        if map_index is None or map_index < 0:
+            return None
+
+        # If ti_count not provided, we need to query it
+        if ti_count is None:
+            ti_count = get_ti_count_for_task(self.task_id, self.dag_id, self.run_id)
+
+        if not ti_count:
+            return None
+
+        return get_relevant_map_indexes(
+            task=self.task,
+            run_id=self.run_id,
+            map_index=map_index,
+            ti_count=ti_count,
+            relative=upstream,
+            dag_id=self.dag_id,
+        )
 
     def get_first_reschedule_date(self, context: Context) -> AwareDatetime | None:
         """Get the first reschedule date for the task instance if found, none otherwise."""
diff --git a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
index 4d0e52e8f2f75..85b6e9c0b271f 100644
--- a/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
+++ b/task-sdk/tests/task_sdk/definitions/test_mappedoperator.py
@@ -31,7 +31,15 @@
 from airflow.sdk.definitions.dag import DAG
 from airflow.sdk.definitions.mappedoperator import MappedOperator
 from airflow.sdk.definitions.xcom_arg import XComArg
-from airflow.sdk.execution_time.comms import GetXCom, SetXCom, XComResult
+from airflow.sdk.execution_time.comms import (
+    GetTICount,
+    GetXCom,
+    GetXComSequenceSlice,
+    SetXCom,
+    TICount,
+    XComResult,
+    XComSequenceSliceResult,
+)
 
 from tests_common.test_utils.mapping import expand_mapped_task  # noqa: F401
 from tests_common.test_utils.mock_operators import (
@@ -252,9 +260,18 @@ def execute(self, context):
         )
         mapped = callable(mapped, task1.output)
 
-    mock_supervisor_comms.send.return_value = XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=["{{ ds }}"])
+    def mock_comms(msg):
+        if isinstance(msg, GetXCom):
+            return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=["{{ ds }}"])
+        if isinstance(msg, GetXComSequenceSlice):
+            return XComSequenceSliceResult(root=["{{ ds }}"])
+        if isinstance(msg, GetTICount):
+            return TICount(count=1)
+        return mock.DEFAULT
+
+    mock_supervisor_comms.send.side_effect = mock_comms
 
-    mapped_ti = create_runtime_ti(task=mapped, map_index=0, upstream_map_indexes={task1.task_id: 1})
+    mapped_ti = create_runtime_ti(task=mapped, map_index=0)
 
     assert isinstance(mapped_ti.task, MappedOperator)
     mapped_ti.task.render_template_fields(context=mapped_ti.get_template_context())
@@ -273,13 +290,13 @@ def test_mapped_render_nested_template_fields(create_runtime_ti, mock_supervisor
             task_id="t", arg2=NestedFields(field_1="{{ ti.task_id }}", field_2="value_2")
         ).expand(arg1=["{{ ti.task_id }}", ["s", "{{ ti.task_id }}"]])
 
-    ti = create_runtime_ti(task=mapped, map_index=0, upstream_map_indexes={})
+    ti = create_runtime_ti(task=mapped, map_index=0)
     ti.task.render_template_fields(context=ti.get_template_context())
     assert ti.task.arg1 == "t"
     assert ti.task.arg2.field_1 == "t"
     assert ti.task.arg2.field_2 == "value_2"
 
-    ti = create_runtime_ti(task=mapped, map_index=1, upstream_map_indexes={})
+    ti = create_runtime_ti(task=mapped, map_index=1)
     ti.task.render_template_fields(context=ti.get_template_context())
     assert ti.task.arg1 == ["s", "t"]
     assert ti.task.arg2.field_1 == "t"
@@ -300,11 +317,20 @@ def test_expand_kwargs_render_template_fields_validating_operator(
         task1 = BaseOperator(task_id="op1")
         mapped = MockOperator.partial(task_id="a", arg2="{{ ti.task_id }}").expand_kwargs(task1.output)
 
-    mock_supervisor_comms.send.return_value = XComResult(
-        key=BaseXCom.XCOM_RETURN_KEY, value=[{"arg1": "{{ ds }}"}, {"arg1": 2}]
-    )
+    xcom_values = [{"arg1": "{{ ds }}"}, {"arg1": 2}]
+
+    def mock_comms(msg):
+        if isinstance(msg, GetXCom):
+            return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=xcom_values)
+        if isinstance(msg, GetXComSequenceSlice):
+            return XComSequenceSliceResult(root=xcom_values)
+        if isinstance(msg, GetTICount):
+            return TICount(count=2)
+        return mock.DEFAULT
+
+    mock_supervisor_comms.send.side_effect = mock_comms
 
-    ti = create_runtime_ti(task=mapped, map_index=map_index, upstream_map_indexes={})
+    ti = create_runtime_ti(task=mapped, map_index=map_index)
     assert isinstance(ti.task, MappedOperator)
     ti.task.render_template_fields(context=ti.get_template_context())
     assert isinstance(ti.task, MockOperator)
@@ -428,14 +454,29 @@ def show(number, letter):
 
         show.expand(number=emit_numbers(), letter=emit_letters())
 
-    def xcom_get(msg):
-        if not isinstance(msg, GetXCom):
-            return mock.DEFAULT
-        task = dag.get_task(msg.task_id)
-        value = task.python_callable()
-        return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
+    numbers = [1, 2]
+    letters = {"a": "x", "b": "y", "c": "z"}
+
+    def mock_comms(msg):
+        if isinstance(msg, GetXCom):
+            if msg.task_id == "emit_numbers":
+                return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers)
+            if msg.task_id == "emit_letters":
+                return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=letters)
+        elif isinstance(msg, GetXComSequenceSlice):
+            if msg.task_id == "emit_numbers":
+                return XComSequenceSliceResult(root=numbers)
+            if msg.task_id == "emit_letters":
+                # Convert dict items to list for XComSequenceSliceResult
+                return XComSequenceSliceResult(root=list(letters.items()))
+        elif isinstance(msg, GetTICount):
+            # show is mapped by 6 (2 numbers * 3 letters)
+            if msg.task_ids and msg.task_ids[0] == "show":
+                return TICount(count=6)
+            return TICount(count=1)
+        return mock.DEFAULT
 
-    mock_supervisor_comms.send.side_effect = xcom_get
+    mock_supervisor_comms.send.side_effect = mock_comms
 
     states = [run_ti(dag, "show", map_index) for map_index in range(6)]
     assert states == [TaskInstanceState.SUCCESS] * 6
@@ -466,14 +507,23 @@ def show(a, b):
         emit_task = emit_numbers()
         show.expand(a=emit_task, b=emit_task)
 
-    def xcom_get(msg):
-        if not isinstance(msg, GetXCom):
-            return mock.DEFAULT
-        task = dag.get_task(msg.task_id)
-        value = task.python_callable()
-        return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
+    numbers = [1, 2]
+
+    def mock_comms(msg):
+        if isinstance(msg, GetXCom):
+            if msg.task_id == "emit_numbers":
+                return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=numbers)
+        elif isinstance(msg, GetXComSequenceSlice):
+            if msg.task_id == "emit_numbers":
+                return XComSequenceSliceResult(root=numbers)
+        elif isinstance(msg, GetTICount):
+            # show is mapped by 4 (2 * 2 cross product)
+            if msg.task_ids and msg.task_ids[0] == "show":
+                return TICount(count=4)
+            return TICount(count=1)
+        return mock.DEFAULT
 
-    mock_supervisor_comms.send.side_effect = xcom_get
+    mock_supervisor_comms.send.side_effect = mock_comms
 
     states = [run_ti(dag, "show", map_index) for map_index in range(4)]
     assert states == [TaskInstanceState.SUCCESS] * 4
@@ -591,20 +641,37 @@ def tg(va):
         # Aggregates results from task group.
         t.override(task_id="t3")(tg1)
 
-    def xcom_get(msg):
-        if not isinstance(msg, GetXCom):
-            return mock.DEFAULT
-        key = (msg.task_id, msg.map_index)
-        if key in expected_values:
-            value = expected_values[key]
-            return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
-        if msg.map_index is None:
-            # Get all mapped XComValues for this ti
-            value = [v for k, v in expected_values.items() if k[0] == msg.task_id]
-            return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
+    # Map task group IDs to their expansion counts
+    task_group_expansion = {"tg": 3}
+
+    def mock_comms_response(msg):
+        if isinstance(msg, GetXCom):
+            key = (msg.task_id, msg.map_index)
+            if key in expected_values:
+                value = expected_values[key]
+                return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
+            if msg.map_index is None:
+                # Get all mapped XComValues for this ti
+                value = [v for k, v in expected_values.items() if k[0] == msg.task_id]
+                return XComResult(key=BaseXCom.XCOM_RETURN_KEY, value=value)
+        elif isinstance(msg, GetXComSequenceSlice):
+            # Handle sequence slicing for pulling all XCom values from mapped tasks
+            task_id = msg.task_id
+            values = [v for k, v in expected_values.items() if k[0] == task_id and k[1] is not None]
+            return XComSequenceSliceResult(root=values)
+        elif isinstance(msg, GetTICount):
+            # Handle TI count queries for upstream_map_indexes computation
+            if msg.task_ids:
+                task_id = msg.task_ids[0]
+                if task_id in expansion_per_task_id:
+                    return TICount(count=len(list(expansion_per_task_id[task_id])))
+                return TICount(count=1)
+            if msg.task_group_id:
+                return TICount(count=task_group_expansion.get(msg.task_group_id, 0))
+            return TICount(count=0)
         return mock.DEFAULT
 
-    mock_supervisor_comms.send.side_effect = xcom_get
+    mock_supervisor_comms.send.side_effect = mock_comms_response
 
     expected_values = {
         ("tg.t1", 0): ["a", "b"],
@@ -622,21 +689,11 @@ def xcom_get(msg):
         "tg.t2": range(3),
         "t3": [None],
     }
-    upstream_map_indexes_per_task_id = {
-        ("tg.t1", 0): {},
-        ("tg.t1", 1): {},
-        ("tg.t1", 2): {},
-        ("tg.t2", 0): {"tg.t1": 0},
-        ("tg.t2", 1): {"tg.t1": 1},
-        ("tg.t2", 2): {"tg.t1": 2},
-        ("t3", None): {"tg.t2": [0, 1, 2]},
-    }
     for task in dag.tasks:
         for map_index in expansion_per_task_id[task.task_id]:
             mapped_ti = create_runtime_ti(
                 task=task.prepare_for_execution(),
                 map_index=map_index,
-                upstream_map_indexes=upstream_map_indexes_per_task_id[(task.task_id, map_index)],
             )
             context = mapped_ti.get_template_context()
             mapped_ti.task.render_template_fields(context)
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_mapping.py b/task-sdk/tests/task_sdk/execution_time/test_task_mapping.py
new file mode 100644
index 0000000000000..5f6bf044785ab
--- /dev/null
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_mapping.py
@@ -0,0 +1,189 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from airflow.sdk import DAG, BaseOperator
+from airflow.sdk.definitions.taskgroup import TaskGroup
+from airflow.sdk.execution_time.comms import TICount
+from airflow.sdk.execution_time.task_mapping import (
+    _find_common_ancestor_mapped_group,
+    _is_further_mapped_inside,
+    get_relevant_map_indexes,
+    get_ti_count_for_task,
+)
+
+
+class TestFindCommonAncestorMappedGroup:
+    """Tests for _find_common_ancestor_mapped_group function."""
+
+    def test_no_common_group_different_dags(self):
+        """Tasks in different DAGs should return None."""
+        with DAG("dag1"):
+            op1 = BaseOperator(task_id="op1")
+
+        with DAG("dag2"):
+            op2 = BaseOperator(task_id="op2")
+
+        result = _find_common_ancestor_mapped_group(op1, op2)
+        assert result is None
+
+    def test_no_common_group_no_mapped_groups(self):
+        """Tasks not in any mapped group should return None."""
+        with DAG("dag1"):
+            op1 = BaseOperator(task_id="op1")
+            op2 = BaseOperator(task_id="op2")
+
+        result = _find_common_ancestor_mapped_group(op1, op2)
+        assert result is None
+
+    def test_no_dag_returns_none(self):
+        """Tasks without DAG should return None."""
+        op1 = BaseOperator(task_id="op1")
+        op2 = BaseOperator(task_id="op2")
+
+        # Function should handle operators not assigned to a DAG gracefully
+        result = _find_common_ancestor_mapped_group(op1, op2)
+        assert result is None
+
+
+class TestIsFurtherMappedInside:
+    """Tests for _is_further_mapped_inside function."""
+
+    def test_mapped_operator_returns_true(self):
+        """A mapped operator should return True."""
+        with DAG("dag1"):
+            with TaskGroup("tg") as tg:
+                op = BaseOperator(task_id="op")
+
+        # Simulate a mapped operator
+        op._is_mapped = True
+
+        result = _is_further_mapped_inside(op, tg)
+        assert result is True
+
+    def test_non_mapped_operator_returns_false(self):
+        """A non-mapped operator with no mapped parent groups should return False."""
+        with DAG("dag1"):
+            with TaskGroup("tg") as tg:
+                op = BaseOperator(task_id="op")
+
+        result = _is_further_mapped_inside(op, tg)
+        assert result is False
+
+
+class TestGetTiCountForTask:
+    """Tests for get_ti_count_for_task function."""
+
+    def test_queries_supervisor(self, mock_supervisor_comms):
+        """Should send GetTICount message to supervisor with task_ids."""
+        from airflow.sdk.execution_time.comms import TICount
+
+        mock_supervisor_comms.send.return_value = TICount(count=3)
+
+        result = get_ti_count_for_task("task_id", "dag_id", "run_id")
+
+        assert result == 3
+        mock_supervisor_comms.send.assert_called_once()
+        call_args = mock_supervisor_comms.send.call_args[0][0]
+        assert call_args.dag_id == "dag_id"
+        assert call_args.task_ids == ["task_id"]
+        assert call_args.run_ids == ["run_id"]
+
+
+class TestGetRelevantMapIndexes:
+    """Tests for get_relevant_map_indexes function."""
+
+    def test_returns_none_when_no_ti_count(self):
+        """Should return None when ti_count is 0 or None."""
+        with DAG("dag1"):
+            op1 = BaseOperator(task_id="op1")
+            op2 = BaseOperator(task_id="op2")
+
+        result = get_relevant_map_indexes(
+            task=op1,
+            run_id="run_id",
+            map_index=0,
+            ti_count=0,
+            relative=op2,
+            dag_id="dag1",
+        )
+        assert result is None
+
+    def test_returns_none_when_no_common_ancestor(self):
+        """Should return None when tasks have no common mapped ancestor."""
+        with DAG("dag1"):
+            op1 = BaseOperator(task_id="op1")
+            op2 = BaseOperator(task_id="op2")
+
+        result = get_relevant_map_indexes(
+            task=op1,
+            run_id="run_id",
+            map_index=0,
+            ti_count=3,
+            relative=op2,
+            dag_id="dag1",
+        )
+        assert result is None
+
+    def test_same_mapped_group_returns_single_index(self, mock_supervisor_comms):
+        """Tasks in same mapped group should get single index matching their map_index."""
+        with DAG("dag1"):
+            with TaskGroup("tg"):
+                op1 = BaseOperator(task_id="op1")
+                op2 = BaseOperator(task_id="op2")
+                op1 >> op2
+
+        # Mock iter_mapped_task_groups to simulate a mapped task group
+        mock_mapped_tg = MagicMock(spec=TaskGroup)
+        mock_mapped_tg.group_id = "tg"
+        op1.iter_mapped_task_groups = MagicMock(spec=TaskGroup, return_value=iter([mock_mapped_tg]))
+        op2.iter_mapped_task_groups = MagicMock(spec=TaskGroup, return_value=iter([mock_mapped_tg]))
+
+        # Mock: op2 has 3 TIs (mapped by 3)
+        mock_supervisor_comms.send.return_value = TICount(count=3)
+
+        # For map_index=1 with ti_count=3, should return 1 (same index)
+        result = get_relevant_map_indexes(
+            task=op2,
+            run_id="run_id",
+            map_index=1,
+            ti_count=3,
+            relative=op1,
+            dag_id="dag1",
+        )
+        assert result == 1
+
+    def test_unmapped_task_pulling_from_mapped_returns_none(self):
+        """Unmapped task pulling from mapped upstream should return None (pull all)."""
+        with DAG("dag1"):
+            op1 = BaseOperator(task_id="op1")
+            op2 = BaseOperator(task_id="op2")
+            op1 >> op2
+
+        # op2 is not in a mapped group, so there's no common ancestor
+        result = get_relevant_map_indexes(
+            task=op2,
+            run_id="run_id",
+            map_index=0,
+            ti_count=1,
+            relative=op1,
+            dag_id="dag1",
+        )
+        assert result is None