From 29ffc0b541969c5b9174765e5bb37aeea5cca114 Mon Sep 17 00:00:00 2001
From: James Stuart Milne <jstuartmilne@gmail.com>
Date: Tue, 10 Dec 2024 19:22:14 -0300
Subject: [PATCH] fix: find task by id or taskReferenceName

---
 ui/src/pages/execution/RightPanel.jsx |  2 +-
 ui/src/utils/helpers.js               | 28 ++++++++++++++++++---------
 2 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/ui/src/pages/execution/RightPanel.jsx b/ui/src/pages/execution/RightPanel.jsx
index b3dc7b79e..839a2692c 100644
--- a/ui/src/pages/execution/RightPanel.jsx
+++ b/ui/src/pages/execution/RightPanel.jsx
@@ -46,7 +46,7 @@ export default function RightPanel({
   const taskResult =
     selectedNode?.data?.task?.executionData?.status === "PENDING"
       ? pendingTaskSelection(selectedNode?.data?.task)
-      : taskWithLatestIteration(execution?.tasks, selectedTask?.ref);
+      : taskWithLatestIteration(execution?.tasks, selectedTask);
 
   const dfOptions = useMemo(
     () => dag && dag.getSiblings(selectedTask),
diff --git a/ui/src/utils/helpers.js b/ui/src/utils/helpers.js
index 407732f63..405ba06da 100644
--- a/ui/src/utils/helpers.js
+++ b/ui/src/utils/helpers.js
@@ -48,7 +48,7 @@ export function astToQuery(node) {
 
     return `${wrapper ? "(" : ""}${clauses.join(` ${combinator} `)}${
       wrapper ? ")" : ""
-    }`;
+      }`;
   } else {
     return "";
   }
@@ -99,16 +99,24 @@ export function getBasename() {
   return _.isEmpty(basename) ? "/" : basename;
 }
 
-export const taskWithLatestIteration = (tasksList, taskReferenceName) => {
-  const filteredTasks = tasksList?.filter(
-    (task) =>
-      task?.workflowTask?.taskReferenceName === taskReferenceName ||
-      task?.referenceTaskName === taskReferenceName
-  );
+export const taskWithLatestIteration = (tasksList = [], selectedTask) => {
+  const taskReferenceName = selectedTask?.ref;
+
+  const findTaskByReferenceName = (task) =>
+    task?.workflowTask?.taskReferenceName === taskReferenceName ||
+    task?.referenceTaskName === taskReferenceName;
+
+  const findTaskById = (task) => task?.taskId === selectedTask?.id;
+
+  // If reference name is not provided, use taskId to find the task
+  const findTask = selectedTask?.ref == null ? findTaskById : findTaskByReferenceName;
+
+  const filteredTasks = tasksList?.filter(findTask);
 
   if (filteredTasks && filteredTasks.length === 1) {
     // task without any retry/iteration
-    return _nth(filteredTasks, 0);
+    const targetTask = _nth(filteredTasks, 0);
+    return targetTask;
   } else if (filteredTasks && filteredTasks.length > 1) {
     const result = filteredTasks.reduce(
       (acc, task, idx) => {
@@ -121,9 +129,11 @@ export const taskWithLatestIteration = (tasksList, taskReferenceName) => {
     );
 
     if (result?.idx > -1) {
-      return _nth(filteredTasks, result.idx);
+      const targetTask = _nth(filteredTasks, result.idx);
+      return targetTask;
     }
   }
+
   return undefined;
 };