feat: allow defining custom topology nodes from tasks

Part-of: kestra-io/kestra-ee#4729

Author: loicmathieu

core/src/main/java/io/kestra/core/models/hierarchies/CustomGraphCluster.java

@@ -0,0 +1,27 @@
+package io.kestra.core.models.hierarchies;
+
+import lombok.Getter;
+
+import java.util.List;
+
+@SuppressWarnings("this-escape")
+@Getter
+public class CustomGraphCluster extends GraphCluster {
+    public CustomGraphCluster(String uid, GraphTask rootTask, List<CustomGraphNode> nodes) {
+        super(rootTask, uid, RelationType.SEQUENTIAL); // TODO should we add a custom relation type?
+
+        this.getGraph().addNode(rootTask);
+        this.addEdge(this.getRoot(), rootTask, new Relation());
+
+        this.getGraph().removeEdge(rootTask, this.getFinally());
+        this.getGraph().removeEdge(rootTask, this.getAfterExecution());
+        this.getGraph().removeNode(this.getFinally());
+        this.getGraph().removeNode(this.getAfterExecution());
+
+        nodes.forEach(node -> {
+            this.getGraph().addNode(node);
+            this.addEdge(rootTask, node, new Relation(RelationType.SEQUENTIAL, null));
+            this.addEdge(node, this.getEnd(), new Relation());
+        });
+    }
+}

core/src/main/java/io/kestra/core/models/hierarchies/CustomGraphNode.java

@@ -0,0 +1,24 @@
+package io.kestra.core.models.hierarchies;
+
+import io.kestra.core.models.Plugin;
+
+public class CustomGraphNode extends AbstractGraph {
+    private final String label;
+    private final Plugin plugin;
+
+    public CustomGraphNode(String uid, String label, Plugin plugin) {
+        super(uid);
+
+        this.label = label;
+        this.plugin = plugin;
+    }
+
+    @Override
+    public String getLabel() {
+        return label;
+    }
+
+    public Plugin getPlugin() {
+        return plugin;
+    }
+}

core/src/main/java/io/kestra/core/models/tasks/RunnableTask.java

@@ -2,8 +2,14 @@
 
 import io.kestra.core.models.Plugin;
 import io.kestra.core.models.WorkerJobLifecycle;
+import io.kestra.core.models.executions.TaskRun;
+import io.kestra.core.models.hierarchies.AbstractGraph;
+import io.kestra.core.models.hierarchies.GraphTask;
+import io.kestra.core.models.hierarchies.RelationType;
 import io.kestra.core.runners.RunContext;
 
+import java.util.List;
+
 /**
  * Interface for tasks that are run in the Worker.
  */
@@ -12,4 +18,13 @@ public interface RunnableTask <T extends Output> extends Plugin, WorkerJobLifecy
      * This method is called inside the Worker to run (execute) the task.
      */
     T run(RunContext runContext) throws Exception;
+
+    /**
+     * Create the topology representation of a runnable task.
+     * <p>
+     * By default, it returns a single GraphTask, tasks may override it to provide a custom topology representation.
+     */
+    default AbstractGraph graph(TaskRun taskRun, List<String> values, RelationType relationType) {
+        return new GraphTask((Task) this, taskRun, values, relationType);
+    }
 }

core/src/main/java/io/kestra/core/utils/GraphUtils.java

@@ -7,6 +7,7 @@
 import io.kestra.core.models.hierarchies.*;
 import io.kestra.core.models.tasks.ExecutableTask;
 import io.kestra.core.models.tasks.FlowableTask;
+import io.kestra.core.models.tasks.RunnableTask;
 import io.kestra.core.models.tasks.Task;
 import io.kestra.core.models.triggers.AbstractTrigger;
 import io.kestra.core.models.triggers.Trigger;
@@ -391,6 +392,8 @@ private static void fillGraph(
                     currentGraph = flowableTask.tasksTree(execution, currentTaskRun, parentValues);
                 } else if (currentTask instanceof ExecutableTask<?> subflowTask) {
                     currentGraph = new SubflowGraphTask(subflowTask, currentTaskRun, parentValues, relationType);
+                } else if (currentTask instanceof RunnableTask<?> runnableTask) {
+                    currentGraph = runnableTask.graph(currentTaskRun, parentValues, relationType);
                 } else {
                     currentGraph = new GraphTask(currentTask, currentTaskRun, parentValues, relationType);
                 }