feat: refine tab preservation logic and fix url parsing contexts

Author: RoyLee1224

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

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { useEffect } from "react";
+import { useEffect, useRef } from "react";
 import { useNavigate } from "react-router-dom";
 import { useLocalStorage } from "usehooks-ts";
 
@@ -43,50 +43,55 @@ type UseTabMemoryOptions = {
 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);
 
-  const normalizedPath = currentPath.replace(/\/$/u, "");
-  const segments = normalizedPath.split("/");
-  const lastSegment = segments[segments.length - 1] ?? "";
-  const isTabSegment = tabs.some((tab) => tab.value === lastSegment);
-  const baseUrl = isTabSegment && lastSegment !== "" ? segments.slice(0, -1).join("/") : normalizedPath;
-
-  // Use entity-based storage key instead of URL-based key
-  const [lastTab, setLastTab] = useLocalStorage<string>(storageKey, tabs[0]?.value ?? "");
+  // Track previous baseUrl to detect entity changes
+  const previousBaseUrlRef = useRef<string | null>(null);
 
   useEffect(() => {
     if (!enabled || tabs.length === 0) {
       return;
     }
 
-    const normalizedCurrentPath = currentPath.replace(/\/$/u, "");
-    const normalizedBase = baseUrl.replace(/\/$/u, "");
+    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);
 
-    const pathSegments = normalizedCurrentPath.split("/");
-    const baseSegments = normalizedBase.split("/");
-    const currentTab = pathSegments[baseSegments.length] ?? "";
+    // Check if we've switched to a different entity (different baseUrl)
+    const hasEntityChanged = previousBaseUrlRef.current !== null && previousBaseUrlRef.current !== baseUrl;
+    const isFirstVisit = previousBaseUrlRef.current === null;
 
-    const isValidTab = tabs.some((tab) => tab.value === currentTab);
-    const isLastTabInCurrentTabs = !lastTab || tabs.some((tab) => tab.value === lastTab);
+    previousBaseUrlRef.current = baseUrl;
 
-    // Only update localStorage if both currentTab is valid AND lastTab is also in current tabs
-    // This prevents different page types (e.g., Task Instance vs Task Group) from overwriting each other's tab memory
-    if (isValidTab && isLastTabInCurrentTabs) {
-      setLastTab(currentTab);
+    // 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 isAtBaseUrl =
-      normalizedCurrentPath === normalizedBase || normalizedCurrentPath === `${normalizedBase}/`;
-    const hasValidSavedTab = Boolean(lastTab) && tabs.some((tab) => tab.value === lastTab);
-    const shouldRedirect = isAtBaseUrl && hasValidSavedTab && lastTab !== tabs[0]?.value;
+    const shouldRedirect =
+      isAtBase && isSavedTabValid && savedTab !== defaultTab && (isFirstVisit || hasEntityChanged);
 
     if (shouldRedirect) {
-      const redirectPath = lastTab ? `${normalizedBase}/${lastTab}` : normalizedBase;
-
       void Promise.resolve(
-        navigate(redirectPath, {
+        navigate(`${baseUrl}/${savedTab}`, {
           replace: true,
         }),
       );
     }
-  }, [baseUrl, currentPath, enabled, lastTab, navigate, setLastTab, tabs]);
+  }, [currentPath, defaultTab, enabled, navigate, savedTab, setSavedTab, tabs]);
 };

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,59 @@ 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 also work for TaskInstance paths (backward compatibility)", () => {
+    // getTaskAdditionalPath is more permissive and will extract from any /tasks/ path
+    // This provides backward compatibility if needed
+    expect(getTaskAdditionalPath("/dags/my_dag/runs/run_1/tasks/task_1/details")).toBe("/details");
+    expect(getTaskAdditionalPath("/dags/my_dag/tasks/task_1/task_instances")).toBe("/task_instances");
+  });
 });
 
 describe("buildTaskInstanceUrl", () => {
@@ -255,5 +315,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");
   });
 });

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

@@ -49,16 +49,16 @@ export const getRedirectPath = (targetPath: string): string => {
 
 export const getTaskInstanceAdditionalPath = (pathname: string): string => {
   const subRoutes = taskInstanceRoutes.filter((route) => route.path !== undefined).map((route) => route.path);
-  // Look for patterns like /tasks/{taskId}/mapped/{mapIndex}/{sub-route}
-  const mappedRegex = /\/tasks\/[^/]+\/mapped\/[^/]+\/(?<subRoute>.+)$/u;
+
+  const mappedRegex = /\/runs\/[^/]+\/tasks\/[^/]+\/mapped\/[^/]+\/(?<subRoute>.+)$/u;
   const mappedMatch = mappedRegex.exec(pathname);
 
   if (mappedMatch?.groups?.subRoute !== undefined) {
     return `/${mappedMatch.groups.subRoute}`;
   }
 
-  // Look for patterns like /tasks/{taskId}/{sub-route} or /tasks/group/{groupId}/{sub-route}
-  const taskRegex = /\/tasks\/(?:group\/)?[^/]+\/(?<subRoute>.+)$/u;
+  // Must include /runs/ to avoid matching Task pages (/dags/xxx/tasks/yyy/task_instances)
+  const taskRegex = /\/runs\/[^/]+\/tasks\/(?:group\/)?[^/]+\/(?<subRoute>.+)$/u;
   const taskMatch = taskRegex.exec(pathname);
 
   if (taskMatch?.groups?.subRoute !== undefined) {
@@ -73,6 +73,17 @@ export const getTaskInstanceAdditionalPath = (pathname: string): string => {
   return "";
 };
 
+export const getTaskAdditionalPath = (pathname: string): string => {
+  const taskPageRegex = /\/dags\/[^/]+\/tasks\/(?:group\/)?[^/]+\/(?<subRoute>.+)$/u;
+  const match = taskPageRegex.exec(pathname);
+
+  if (match?.groups?.subRoute !== undefined) {
+    return `/${match.groups.subRoute}`;
+  }
+
+  return "";
+};
+
 export const buildTaskInstanceUrl = (params: {
   currentPathname: string;
   dagId: string;
@@ -84,8 +95,11 @@ export const buildTaskInstanceUrl = (params: {
 }): string => {
   const { currentPathname, dagId, isGroup = false, isMapped = false, mapIndex, runId, taskId } = params;
   const groupPath = isGroup ? "group/" : "";
-  // Task groups only have "Task Instances" tab, so never preserve tabs for groups
-  const additionalPath = isGroup ? "" : getTaskInstanceAdditionalPath(currentPathname);
+
+  // Only preserve tabs for specific task instances, excluding groups and mapped lists
+  const isMappedOverview = isMapped && (mapIndex === undefined || mapIndex === "-1");
+  const shouldPreserveTabs = !isGroup && !isMappedOverview;
+  const additionalPath = shouldPreserveTabs ? getTaskInstanceAdditionalPath(currentPathname) : "";
 
   let basePath = `/dags/${dagId}/runs/${runId}/tasks/${groupPath}${taskId}`;