From f09be969bb22e7e3acc78a00113ce8fad8243112 Mon Sep 17 00:00:00 2001
From: YBoy <ashleewunsch1226@gmail.com>
Date: Tue, 7 Apr 2026 15:18:00 +0200
Subject: [PATCH] refactor(api): type single-node graph structure with
 TypedDicts in workflow_entry (#34671)

---
 api/core/workflow/workflow_entry.py | 36 ++++++++++++++++++++++-------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/api/core/workflow/workflow_entry.py b/api/core/workflow/workflow_entry.py
index 2346a95d6a..cecc20145a 100644
--- a/api/core/workflow/workflow_entry.py
+++ b/api/core/workflow/workflow_entry.py
@@ -1,7 +1,7 @@
 import logging
 import time
 from collections.abc import Generator, Mapping, Sequence
-from typing import Any
+from typing import Any, TypedDict
 
 from graphon.entities import GraphInitParams
 from graphon.entities.graph_config import NodeConfigDictAdapter
@@ -107,6 +107,26 @@ class _WorkflowChildEngineBuilder:
         return child_engine
 
 
+class _NodeConfigDict(TypedDict):
+    id: str
+    width: int
+    height: int
+    type: str
+    data: dict[str, Any]
+
+
+class _EdgeConfigDict(TypedDict):
+    source: str
+    target: str
+    sourceHandle: str
+    targetHandle: str
+
+
+class SingleNodeGraphDict(TypedDict):
+    nodes: list[_NodeConfigDict]
+    edges: list[_EdgeConfigDict]
+
+
 class WorkflowEntry:
     def __init__(
         self,
@@ -318,7 +338,7 @@ class WorkflowEntry:
         node_data: dict[str, Any],
         node_width: int = 114,
         node_height: int = 514,
-    ) -> dict[str, Any]:
+    ) -> SingleNodeGraphDict:
         """
         Create a minimal graph structure for testing a single node in isolation.
 
@@ -328,14 +348,14 @@ class WorkflowEntry:
         :param node_height: height for UI layout (default: 100)
         :return: graph dictionary with start node and target node
         """
-        node_config = {
+        node_config: _NodeConfigDict = {
             "id": node_id,
             "width": node_width,
             "height": node_height,
             "type": "custom",
             "data": node_data,
         }
-        start_node_config = {
+        start_node_config: _NodeConfigDict = {
             "id": "start",
             "width": node_width,
             "height": node_height,
@@ -346,9 +366,9 @@ class WorkflowEntry:
                 "desc": "Start",
             },
         }
-        return {
-            "nodes": [start_node_config, node_config],
-            "edges": [
+        return SingleNodeGraphDict(
+            nodes=[start_node_config, node_config],
+            edges=[
                 {
                     "source": "start",
                     "target": node_id,
@@ -356,7 +376,7 @@ class WorkflowEntry:
                     "targetHandle": "target",
                 }
             ],
-        }
+        )
 
     @classmethod
     def run_free_node(