diff --git a/studio/frontend/src/features/data-recipes/learning-recipes/text-to-sql.json b/studio/frontend/src/features/data-recipes/learning-recipes/text-to-sql.json index d0776d3ead..d2d9abbd41 100644 --- a/studio/frontend/src/features/data-recipes/learning-recipes/text-to-sql.json +++ b/studio/frontend/src/features/data-recipes/learning-recipes/text-to-sql.json @@ -147,8 +147,8 @@ }, { "id": "note_2", - "x": 1672, - "y": 1577, + "x": 1675.8410596026492, + "y": 1644.2185430463576, "width": 400, "node_type": "markdown_note", "name": "note_2", @@ -158,8 +158,8 @@ }, { "id": "note_3", - "x": 2126, - "y": 1485, + "x": 2198.980132450331, + "y": 1723.1456953642385, "width": 400, "node_type": "markdown_note", "name": "note_3", @@ -180,13 +180,13 @@ }, { "id": "vllm", - "x": 1880, + "x": 1939.5364238410598, "y": 781.25, "width": 400 }, { "id": "sql-pro", - "x": 1880, + "x": 1939.5364238410593, "y": 975.25, "width": 400 }, @@ -216,20 +216,20 @@ }, { "id": "sql_prompt", - "x": 1666.887417218543, - "y": 1378.9437086092717, + "x": 1672.6490066225165, + "y": 1457.6854304635763, "width": 400 }, { "id": "sql", - "x": 2120, - "y": 1236.25, + "x": 2194.9006622516554, + "y": 1457.110927152318, "width": 400 }, { "id": "sql-validator", - "x": 2600, - "y": 1304.75, + "x": 2682.5827814569534, + "y": 1491.0413907284767, "width": 400 } ], diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/derive-display-graph.ts b/studio/frontend/src/features/recipe-studio/utils/graph/derive-display-graph.ts index 79ff76c772..45fbea828e 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/derive-display-graph.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/derive-display-graph.ts @@ -46,8 +46,14 @@ function normalizeEdge( configs: Record, layoutDirection: LayoutDirection, activeEdgeIds: Set, + runningNodeId: string | null, + doneNodeIds: Set, ): Edge { - const isActiveEdge = activeEdgeIds.has(edge.id); + const isActiveByRuntimeTarget = + Boolean(runningNodeId) && + edge.target === runningNodeId && + !isAuxEdge(edge); + const isActiveEdge = activeEdgeIds.has(edge.id) || isActiveByRuntimeTarget; const isAux = isAuxEdge(edge); if (isAux) { return { @@ -58,13 +64,28 @@ function normalizeEdge( }; } - const source = configs[edge.source]; - const target = configs[edge.target]; + const isActiveReversedRuntimeEdge = + Boolean(runningNodeId) && + isActiveEdge && + edge.source === runningNodeId && + doneNodeIds.has(edge.target); + const displayEdge = isActiveReversedRuntimeEdge + ? { + ...edge, + source: edge.target, + target: edge.source, + sourceHandle: getDefaultDataSourceHandle(layoutDirection), + targetHandle: getDefaultDataTargetHandle(layoutDirection), + } + : edge; + + const source = configs[displayEdge.source]; + const target = configs[displayEdge.target]; const semantic = - edge.type === "semantic" || + displayEdge.type === "semantic" || (Boolean(source && target) && isSemanticRelation(source, target)); - const sourceHandleNormalized = normalizeRecipeHandleId(edge.sourceHandle); - const targetHandleNormalized = normalizeRecipeHandleId(edge.targetHandle); + const sourceHandleNormalized = normalizeRecipeHandleId(displayEdge.sourceHandle); + const targetHandleNormalized = normalizeRecipeHandleId(displayEdge.targetHandle); const semanticSourceDefault = source?.kind === "llm" ? getDefaultDataSourceHandle(layoutDirection) @@ -104,11 +125,11 @@ function normalizeEdge( } return { - ...edge, + ...displayEdge, type: semantic ? "semantic" : "canvas", data: semantic - ? { ...(edge.data ?? {}), active: isActiveEdge } - : { ...(edge.data ?? {}), path: "smoothstep", active: isActiveEdge }, + ? { ...(displayEdge.data ?? {}), active: isActiveEdge } + : { ...(displayEdge.data ?? {}), path: "smoothstep", active: isActiveEdge }, sourceHandle, targetHandle, animated: isActiveEdge, @@ -585,7 +606,14 @@ export function deriveDisplayGraph({ return { nodes: [...displayNodes, ...auxNodes], edges: [...edges, ...auxEdges].map((edge) => - normalizeEdge(edge, configs, layoutDirection, activeEdgeIds), + normalizeEdge( + edge, + configs, + layoutDirection, + activeEdgeIds, + runningNodeId, + doneNodeIds, + ), ), }; } diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts b/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts index e0c2045ee0..d36014d4c9 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts @@ -5,6 +5,7 @@ import type { RecipeExecutionStatus, } from "../../execution-types"; import type { NodeConfig } from "../../types"; +import { extractRefs } from "../refs"; const ACTIVE_STATUSES: ReadonlySet = new Set([ "pending", @@ -35,6 +36,47 @@ function isAuxEdge(edge: Edge): boolean { return edge.source.startsWith("aux-") || edge.target.startsWith("aux-"); } +function collectTemplateRefs(config: NodeConfig | null): Set { + if (!config) { + return new Set(); + } + const refs = new Set(); + if (config.kind === "llm") { + for (const ref of extractRefs(config.prompt ?? "")) { + refs.add(ref.trim()); + } + for (const ref of extractRefs(config.system_prompt ?? "")) { + refs.add(ref.trim()); + } + if (typeof config.output_format === "string") { + for (const ref of extractRefs(config.output_format)) { + refs.add(ref.trim()); + } + } + return refs; + } + if (config.kind === "expression") { + for (const ref of extractRefs(config.expr ?? "")) { + refs.add(ref.trim()); + } + } + return refs; +} + +function isReversedRuntimeReferenceEdge(input: { + edge: Edge; + runningNodeId: string; + runningTemplateRefs: Set; + configs: Record; +}): boolean { + const { edge, runningNodeId, runningTemplateRefs, configs } = input; + if (edge.source !== runningNodeId) { + return false; + } + const targetName = configs[edge.target]?.name?.trim() ?? ""; + return Boolean(targetName && runningTemplateRefs.has(targetName)); +} + function hasLiveExecutionSignal(execution: RecipeExecutionRecord): boolean { if (execution.lastEventId !== null) { return true; @@ -121,6 +163,14 @@ export function deriveGraphRuntimeVisualState(input: { const activeEdgeIds = new Set(); if (runningNodeId) { + const runningConfig = configs[runningNodeId] ?? null; + const runningTemplateRefs = collectTemplateRefs(runningConfig); + for (const ref of runningTemplateRefs) { + const refNodeId = nameToNodeId.get(ref); + if (refNodeId && refNodeId !== runningNodeId) { + doneNodeIds.add(refNodeId); + } + } for (const upstreamNodeId of collectUpstreamDoneNodeIds({ rootNodeId: runningNodeId, edges, @@ -129,13 +179,23 @@ export function deriveGraphRuntimeVisualState(input: { doneNodeIds.add(upstreamNodeId); } for (const edge of edges) { - if (edge.target !== runningNodeId) { - continue; - } if (isAuxEdge(edge)) { continue; } - activeEdgeIds.add(edge.id); + if (edge.target === runningNodeId) { + activeEdgeIds.add(edge.id); + continue; + } + if ( + isReversedRuntimeReferenceEdge({ + edge, + runningNodeId, + runningTemplateRefs, + configs, + }) + ) { + activeEdgeIds.add(edge.id); + } } }