Preserve active tabs during navigation

airflow-core/src/airflow/ui/src/hooks/navigation/useNavigation.ts

@@ -21,7 +21,7 @@ import { useLocation, useNavigate, useParams } from "react-router-dom";
 
 import type { GridRunsResponse } from "openapi/requests";
 import type { GridTask } from "src/layouts/Details/Grid/utils";
-import { buildTaskInstanceUrl } from "src/utils/links";
+import { buildTaskInstanceUrl, getTaskAdditionalPath } from "src/utils/links";
 
 import {
   NavigationModes,
@@ -77,8 +77,11 @@ const buildPath = (params: {
   switch (mode) {
     case NavigationModes.RUN:
       return `/dags/${dagId}/runs/${run.run_id}`;
-    case NavigationModes.TASK:
-      return `/dags/${dagId}/tasks/${groupPath}${task.id}`;
+    case NavigationModes.TASK: {
+      const additionalPath = getTaskAdditionalPath(pathname);
+
+      return `/dags/${dagId}/tasks/${groupPath}${task.id}${additionalPath}`;
+    }
     case NavigationModes.TI:
       return buildTaskInstanceUrl({
         currentPathname: pathname,

airflow-core/src/airflow/ui/src/hooks/useTabMemory.ts

@@ -0,0 +1,97 @@
+/*!
+ * 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.
+ */
+import { useEffect, useRef } from "react";
+import { useNavigate } from "react-router-dom";
+import { useLocalStorage } from "usehooks-ts";
+
+import type { TabItem } from "./useRequiredActionTabs";
+
+export const TabStorageKeys = {
+  DAG: "tab_view_dag",
+  MAPPED_TASK_INSTANCE: "tab_view_mapped_ti",
+  RUN: "tab_view_run",
+  TASK: "tab_view_task",
+  TASK_GROUP: "tab_view_task_group",
+  TASK_INSTANCE: "tab_view_ti",
+} as const;
+
+export type TabStorageKey = (typeof TabStorageKeys)[keyof typeof TabStorageKeys];
+
+type UseTabMemoryOptions = {
+  currentPath: string;
+  enabled?: boolean;
+  storageKey: TabStorageKey;
+  tabs: Array<TabItem>;
+};
+
+export const useTabMemory = (options: UseTabMemoryOptions) => {
+  const { currentPath, enabled = true, storageKey, tabs } = options;
+  const navigate = useNavigate();
+  const defaultTab = tabs[0]?.value ?? "";
+  const [savedTab, setSavedTab] = useLocalStorage<string>(storageKey, defaultTab);
+
+  // Track previous baseUrl to detect entity changes
+  const previousBaseUrlRef = useRef<string | null>(null);
+
+  useEffect(() => {
+    if (!enabled || tabs.length === 0) {
+      return;
+    }
+
+    const normalizedPath = currentPath.replace(/\/$/u, "");
+    const pathSegments = normalizedPath.split("/");
+    const lastSegment = pathSegments[pathSegments.length - 1] ?? "";
+    const isTabSegment = tabs.some((tab) => tab.value === lastSegment);
+
+    const baseUrl = isTabSegment && lastSegment !== "" ? pathSegments.slice(0, -1).join("/") : normalizedPath;
+    const currentTab = isTabSegment ? lastSegment : "";
+
+    const isAtBase = normalizedPath === baseUrl;
+    const isSavedTabValid = Boolean(savedTab) && tabs.some((tab) => tab.value === savedTab);
+
+    // Check if we've switched to a different entity (different baseUrl)
+    const hasEntityChanged = previousBaseUrlRef.current !== null && previousBaseUrlRef.current !== baseUrl;
+    const isFirstVisit = previousBaseUrlRef.current === null;
+
+    previousBaseUrlRef.current = baseUrl;
+
+    // Check if current tab is valid (including empty string for default tab)
+    const isCurrentTabValid = tabs.some((tab) => tab.value === currentTab);
+
+    // Update saved preference when viewing a valid tab
+    if (isCurrentTabValid) {
+      const isPreviousSavedTabValid = !savedTab || tabs.some((tab) => tab.value === savedTab);
+
+      if (isPreviousSavedTabValid && currentTab !== savedTab) {
+        setSavedTab(currentTab);
+      }
+    }
+
+    const shouldRedirect =
+      isAtBase && isSavedTabValid && savedTab !== defaultTab && (isFirstVisit || hasEntityChanged);
+
+    if (shouldRedirect) {
+      void Promise.resolve(
+        navigate(`${baseUrl}/${savedTab}`, {
+          replace: true,
+        }),
+      );
+    }
+  }, [currentPath, defaultTab, enabled, navigate, savedTab, setSavedTab, tabs]);
+};

airflow-core/src/airflow/ui/src/layouts/Details/DetailsLayout.tsx

@@ -16,6 +16,8 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
+/* eslint-disable max-lines */
 import { Box, HStack, Flex, useDisclosure, IconButton } from "@chakra-ui/react";
 import { useReactFlow } from "@xyflow/react";
 import { useRef, useState } from "react";
@@ -29,7 +31,7 @@ import {
   PanelGroup,
   PanelResizeHandle,
 } from "react-resizable-panels";
-import { Outlet, useParams } from "react-router-dom";
+import { Outlet, useLocation, useParams } from "react-router-dom";
 import { useLocalStorage } from "usehooks-ts";
 
 import { useDagServiceGetDag, useDagWarningServiceListDagWarnings } from "openapi/queries";
@@ -43,6 +45,8 @@ import { Toaster } from "src/components/ui";
 import { Tooltip } from "src/components/ui/Tooltip";
 import { HoverProvider } from "src/context/hover";
 import { OpenGroupsProvider } from "src/context/openGroups";
+import type { TabStorageKey } from "src/hooks/useTabMemory";
+import { TabStorageKeys, useTabMemory } from "src/hooks/useTabMemory";
 
 import { DagBreadcrumb } from "./DagBreadcrumb";
 import { Gantt } from "./Gantt/Gantt";
@@ -54,14 +58,29 @@ import { PanelButtons } from "./PanelButtons";
 type Props = {
   readonly error?: unknown;
   readonly isLoading?: boolean;
+  readonly storageKey?: TabStorageKey;
   readonly tabs: Array<{ icon: ReactNode; label: string; value: string }>;
 } & PropsWithChildren;
 
-export const DetailsLayout = ({ children, error, isLoading, tabs }: Props) => {
+export const DetailsLayout = ({
+  children,
+  error,
+  isLoading,
+  storageKey: providedStorageKey,
+  tabs,
+}: Props) => {
   const { t: translate } = useTranslation();
   const { dagId = "", runId } = useParams();
+  const location = useLocation();
   const { data: dag } = useDagServiceGetDag({ dagId });
   const [defaultDagView] = useLocalStorage<"graph" | "grid">("default_dag_view", "grid");
+  const storageKey = providedStorageKey ?? (Boolean(runId) ? TabStorageKeys.RUN : TabStorageKeys.DAG);
+
+  useTabMemory({
+    currentPath: location.pathname,
+    storageKey,
+    tabs,
+  });
   const panelGroupRef = useRef<ImperativePanelGroupHandle | null>(null);
   const [dagView, setDagView] = useLocalStorage<"graph" | "grid">(`dag_view-${dagId}`, defaultDagView);
   const [limit, setLimit] = useLocalStorage<number>(`dag_runs_limit-${dagId}`, 10);

airflow-core/src/airflow/ui/src/pages/GroupTaskInstance/GroupTaskInstance.tsx

@@ -21,6 +21,7 @@ import { useTranslation } from "react-i18next";
 import { MdOutlineTask } from "react-icons/md";
 import { useParams } from "react-router-dom";
 
+import { TabStorageKeys } from "src/hooks/useTabMemory";
 import { DetailsLayout } from "src/layouts/Details/DetailsLayout";
 import { useGridTiSummaries } from "src/queries/useGridTISummaries.ts";
 
@@ -36,7 +37,7 @@ export const GroupTaskInstance = () => {
 
   return (
     <ReactFlowProvider>
-      <DetailsLayout tabs={tabs}>
+      <DetailsLayout storageKey={TabStorageKeys.TASK_GROUP} tabs={tabs}>
         {taskInstance === undefined ? undefined : <Header taskInstance={taskInstance} />}
       </DetailsLayout>
     </ReactFlowProvider>

airflow-core/src/airflow/ui/src/pages/MappedTaskInstance/MappedTaskInstance.tsx

@@ -21,6 +21,7 @@ import { useTranslation } from "react-i18next";
 import { MdOutlineTask } from "react-icons/md";
 import { useParams } from "react-router-dom";
 
+import { TabStorageKeys } from "src/hooks/useTabMemory";
 import { DetailsLayout } from "src/layouts/Details/DetailsLayout";
 import { useGridTiSummaries } from "src/queries/useGridTISummaries.ts";
 
@@ -44,7 +45,7 @@ export const MappedTaskInstance = () => {
 
   return (
     <ReactFlowProvider>
-      <DetailsLayout tabs={tabs}>
+      <DetailsLayout storageKey={TabStorageKeys.MAPPED_TASK_INSTANCE} tabs={tabs}>
         {taskInstance === undefined ? undefined : <Header taskInstance={taskInstance} />}
       </DetailsLayout>
     </ReactFlowProvider>

airflow-core/src/airflow/ui/src/pages/Task/Task.tsx

@@ -26,6 +26,7 @@ import { useParams } from "react-router-dom";
 import { useTaskServiceGetTask } from "openapi/queries";
 import { usePluginTabs } from "src/hooks/usePluginTabs";
 import { useRequiredActionTabs } from "src/hooks/useRequiredActionTabs";
+import { TabStorageKeys } from "src/hooks/useTabMemory";
 import { DetailsLayout } from "src/layouts/Details/DetailsLayout";
 import { useGridStructure } from "src/queries/useGridStructure.ts";
 import { getGroupTask } from "src/utils/groupTask";
@@ -77,7 +78,7 @@ export const Task = () => {
 
   return (
     <ReactFlowProvider>
-      <DetailsLayout error={error} isLoading={isLoading} tabs={displayTabs}>
+      <DetailsLayout error={error} isLoading={isLoading} storageKey={TabStorageKeys.TASK} tabs={displayTabs}>
         {task === undefined ? undefined : <Header task={task} />}
         {groupTask ? <GroupTaskHeader title={groupTask.label} /> : undefined}
       </DetailsLayout>

airflow-core/src/airflow/ui/src/pages/TaskInstance/TaskInstance.tsx

@@ -27,6 +27,7 @@ import { useParams } from "react-router-dom";
 import { useTaskInstanceServiceGetMappedTaskInstance } from "openapi/queries";
 import { usePluginTabs } from "src/hooks/usePluginTabs";
 import { useRequiredActionTabs } from "src/hooks/useRequiredActionTabs";
+import { TabStorageKeys } from "src/hooks/useTabMemory";
 import { DetailsLayout } from "src/layouts/Details/DetailsLayout";
 import { useGridTiSummaries } from "src/queries/useGridTISummaries.ts";
 import { isStatePending, useAutoRefresh } from "src/utils";
@@ -103,9 +104,14 @@ export const TaskInstance = () => {
     refetchInterval: isStatePending(taskInstance?.state) ? refetchInterval : false,
   });
 
+  const storageKey =
+    taskInstance && taskInstance.map_index > -1
+      ? TabStorageKeys.MAPPED_TASK_INSTANCE
+      : TabStorageKeys.TASK_INSTANCE;
+
   return (
     <ReactFlowProvider>
-      <DetailsLayout error={error} isLoading={isLoading} tabs={displayTabs}>
+      <DetailsLayout error={error} isLoading={isLoading} storageKey={storageKey} tabs={displayTabs}>
         {taskInstance === undefined ? (
           <Heading p={2} size="lg">
             {translate("common:noItemsFound", { modelName: translate("common:taskInstance_one") })}

airflow-core/src/airflow/ui/src/utils/links.test.ts

@@ -16,11 +16,18 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+
+/* eslint-disable max-lines */
 import { describe, it, expect } from "vitest";
 
 import type { TaskInstanceResponse } from "openapi/requests/types.gen";
 
-import { buildTaskInstanceUrl, getTaskInstanceAdditionalPath, getTaskInstanceLink } from "./links";
+import {
+  buildTaskInstanceUrl,
+  getTaskAdditionalPath,
+  getTaskInstanceAdditionalPath,
+  getTaskInstanceLink,
+} from "./links";
 
 describe("getTaskInstanceLink", () => {
   const testCases = [
@@ -148,6 +155,61 @@ describe("getTaskInstanceAdditionalPath", () => {
       getTaskInstanceAdditionalPath("/dags/my-dag_v2/runs/run_1-test/tasks/task.1/rendered_templates"),
     ).toBe("/rendered_templates");
   });
+
+  it("should NOT preserve sub-routes from Task pages (without /runs/)", () => {
+    // Task page with task_instances tab - should NOT be preserved
+    expect(getTaskInstanceAdditionalPath("/dags/my_dag/tasks/task_1/task_instances")).toBe("");
+
+    // Task page with overview/events tab - should NOT be preserved
+    expect(getTaskInstanceAdditionalPath("/dags/my_dag/tasks/task_1/events")).toBe("");
+
+    // Task group page - should NOT be preserved
+    expect(getTaskInstanceAdditionalPath("/dags/my_dag/tasks/group/my_group/task_instances")).toBe("");
+
+    // Only TaskInstance pages (with /runs/) should preserve sub-routes
+    expect(getTaskInstanceAdditionalPath("/dags/my_dag/runs/run_1/tasks/task_1/task_instances")).toBe(
+      "/task_instances",
+    );
+  });
+});
+
+describe("getTaskAdditionalPath", () => {
+  it("should return empty string for basic task path", () => {
+    const result = getTaskAdditionalPath("/dags/my_dag/tasks/task_1");
+
+    expect(result).toBe("");
+  });
+
+  it("should extract sub-route from task path", () => {
+    const result = getTaskAdditionalPath("/dags/my_dag/tasks/task_1/task_instances");
+
+    expect(result).toBe("/task_instances");
+  });
+
+  it("should extract sub-route from group task path", () => {
+    const result = getTaskAdditionalPath("/dags/my_dag/tasks/group/my_group/events");
+
+    expect(result).toBe("/events");
+  });
+
+  it("should handle all known task routes", () => {
+    const knownRoutes = ["task_instances", "required_actions", "events"];
+
+    for (const route of knownRoutes) {
+      const result = getTaskAdditionalPath(`/dags/test/tasks/task_1/${route}`);
+
+      expect(result).toBe(`/${route}`);
+    }
+  });
+
+  it("should NOT extract from TaskInstance paths (with /runs/)", () => {
+    // getTaskAdditionalPath is specifically for Task pages (without /runs/)
+    // TaskInstance paths should use getTaskInstanceAdditionalPath instead
+    expect(getTaskAdditionalPath("/dags/my_dag/runs/run_1/tasks/task_1/details")).toBe("");
+
+    // But it should work for Task pages
+    expect(getTaskAdditionalPath("/dags/my_dag/tasks/task_1/task_instances")).toBe("/task_instances");
+  });
 });
 
 describe("buildTaskInstanceUrl", () => {
@@ -255,5 +317,41 @@ describe("buildTaskInstanceUrl", () => {
         taskId: "new_group",
       }),
     ).toBe("/dags/new_dag/runs/new_run/tasks/group/new_group/mapped/3");
+
+    // Mapped task overview (no mapIndex) should never preserve tabs
+    expect(
+      buildTaskInstanceUrl({
+        currentPathname: "/dags/old/runs/old/tasks/old_task/xcom",
+        dagId: "new_dag",
+        isMapped: true,
+        mapIndex: undefined,
+        runId: "new_run",
+        taskId: "mapped_task",
+      }),
+    ).toBe("/dags/new_dag/runs/new_run/tasks/mapped_task/mapped");
+
+    // Mapped task overview with mapIndex="-1" should also not preserve
+    expect(
+      buildTaskInstanceUrl({
+        currentPathname: "/dags/old/runs/old/tasks/old_task/details",
+        dagId: "new_dag",
+        isMapped: true,
+        mapIndex: "-1",
+        runId: "new_run",
+        taskId: "mapped_task",
+      }),
+    ).toBe("/dags/new_dag/runs/new_run/tasks/mapped_task/mapped");
+
+    // But specific mapped instance (with mapIndex) SHOULD preserve tabs
+    expect(
+      buildTaskInstanceUrl({
+        currentPathname: "/dags/old/runs/old/tasks/old_task/mapped/0/xcom",
+        dagId: "new_dag",
+        isMapped: true,
+        mapIndex: "5",
+        runId: "new_run",
+        taskId: "mapped_task",
+      }),
+    ).toBe("/dags/new_dag/runs/new_run/tasks/mapped_task/mapped/5/xcom");
   });
 });