From 745d9e414da5285b332be7750be990c0aa472cb8 Mon Sep 17 00:00:00 2001
From: Gayathri Srividya Rajavarapu <gayathrir@Gayathris-MacBook-Air.local>
Date: Wed, 3 Jun 2026 16:01:10 +0530
Subject: [PATCH] UI: Fix task instance state badge staying stale after Mark-as
 action

---
 .../ui/src/queries/usePatchTaskInstance.ts    | 26 +------------------
 1 file changed, 1 insertion(+), 25 deletions(-)

diff --git a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts
index 0f0a5988ed8d3..6dfcc33cdb8d8 100644
--- a/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts
+++ b/airflow-core/src/airflow/ui/src/queries/usePatchTaskInstance.ts
@@ -20,7 +20,6 @@ import { useQueryClient } from "@tanstack/react-query";
 import { useTranslation } from "react-i18next";
 
 import {
-  UseTaskInstanceServiceGetMappedTaskInstanceKeyFn,
   useTaskInstanceServiceGetMappedTaskInstanceKey,
   UseTaskInstanceServiceGetTaskInstanceKeyFn,
   useTaskInstanceServiceGetTaskInstancesKey,
@@ -65,36 +64,13 @@ export const usePatchTaskInstance = ({
       [useTaskInstanceServiceGetTaskInstancesKey],
       [usePatchTaskInstanceDryRunKey, dagId, dagRunId, { mapIndex, taskId }],
       [useClearTaskInstancesDryRunKey, dagId],
+      [useTaskInstanceServiceGetMappedTaskInstanceKey, { dagId, dagRunId, taskId }],
       ...tiPerAttemptQueryKeys,
     ];
 
-    if (mapIndex !== undefined) {
-      queryKeys.push(UseTaskInstanceServiceGetMappedTaskInstanceKeyFn({ dagId, dagRunId, mapIndex, taskId }));
-    }
-
     await Promise.all([
       ...gridQueryKeys(dagId).map((key) => queryClient.invalidateQueries({ queryKey: key })),
       ...queryKeys.map((key) => queryClient.invalidateQueries({ queryKey: key })),
-      // Wildcard match when patching every mapped TI (mapIndex undefined):
-      // invalidate the mapped-TI cache for every map_index of this task.
-      // The mapped-TI key embeds its params as a single object, so prefix
-      // matching on a partial key doesn't catch them — match via predicate.
-      ...(mapIndex === undefined
-        ? [
-            queryClient.invalidateQueries({
-              predicate: ({ queryKey }) => {
-                if (queryKey[0] !== useTaskInstanceServiceGetMappedTaskInstanceKey) {
-                  return false;
-                }
-                const params = queryKey[1] as
-                  | { dagId?: string; dagRunId?: string; taskId?: string }
-                  | undefined;
-
-                return params?.dagId === dagId && params.dagRunId === dagRunId && params.taskId === taskId;
-              },
-            }),
-          ]
-        : []),
     ]);
 
     if (onSuccess) {