From 2cc9981ef9847e3e95be8f87b09497063129502e Mon Sep 17 00:00:00 2001
From: Shine1i <wasimysdev@gmail.com>
Date: Sun, 22 Feb 2026 04:07:01 +0100
Subject: [PATCH] refactor: enhance variable handling with structured entries
 and UI updates for badges

---
 .../components/inline/inline-expression.tsx   | 14 +++++---
 .../components/recipe-graph-aux-node.tsx      | 14 +++++---
 .../dialogs/shared/available-variables.tsx    | 14 +++++---
 .../features/recipe-studio/utils/variables.ts | 35 ++++++++++++++-----
 4 files changed, 54 insertions(+), 23 deletions(-)

diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx
index 6f754e04d6..94c5470f49 100644
--- a/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx
+++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-expression.tsx
@@ -10,7 +10,7 @@ import {
 import type { ReactElement } from "react";
 import { useRecipeStudioStore } from "../../stores/recipe-studio";
 import type { ExpressionConfig, ExpressionDtype } from "../../types";
-import { getAvailableVariables } from "../../utils/variables";
+import { getAvailableVariableEntries } from "../../utils/variables";
 import { InlineField } from "./inline-field";
 
 type InlineExpressionProps = {
@@ -25,7 +25,7 @@ export function InlineExpression({
   onUpdate,
 }: InlineExpressionProps): ReactElement {
   const configs = useRecipeStudioStore((state) => state.configs);
-  const vars = getAvailableVariables(configs, config.id);
+  const vars = getAvailableVariableEntries(configs, config.id);
 
   return (
     <div className="space-y-3">
@@ -64,11 +64,15 @@ export function InlineExpression({
           <div className="flex flex-wrap gap-1">
             {vars.map((v) => (
               <Badge
-                key={v}
+                key={`${v.source}:${v.name}`}
                 variant="secondary"
-                className="corner-squircle h-4 px-1.5 font-mono text-[10px]"
+                className={
+                  v.source === "seed"
+                    ? "corner-squircle h-4 border-blue-500/25 bg-blue-500/10 px-1.5 font-mono text-[10px] text-blue-700 dark:text-blue-300"
+                    : "corner-squircle h-4 px-1.5 font-mono text-[10px]"
+                }
               >
-                {v}
+                {v.name}
               </Badge>
             ))}
           </div>
diff --git a/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx b/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx
index 75af2931d2..4de17ccfc0 100644
--- a/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx
+++ b/studio/frontend/src/features/recipe-studio/components/recipe-graph-aux-node.tsx
@@ -15,7 +15,7 @@ import { MAX_NODE_WIDTH, MIN_NODE_WIDTH } from "../constants";
 import { useRecipeStudioStore } from "../stores/recipe-studio";
 import type { LayoutDirection, LlmConfig, Score, ScoreOption } from "../types";
 import { HANDLE_IDS } from "../utils/handles";
-import { getAvailableVariables } from "../utils/variables";
+import { getAvailableVariableEntries } from "../utils/variables";
 import { BaseNode, BaseNodeContent, BaseNodeHeader, BaseNodeHeaderTitle } from "./rf-ui/base-node";
 
 type PromptField = "prompt" | "system_prompt";
@@ -61,7 +61,7 @@ function updateOptionAt(
 
 function AuxVariableBadges({ llmId }: { llmId: string }): ReactElement | null {
   const configs = useRecipeStudioStore((state) => state.configs);
-  const vars = getAvailableVariables(configs, llmId);
+  const vars = getAvailableVariableEntries(configs, llmId);
   if (vars.length === 0) return null;
   return (
     <div className="space-y-1">
@@ -69,11 +69,15 @@ function AuxVariableBadges({ llmId }: { llmId: string }): ReactElement | null {
       <div className="flex flex-wrap gap-1">
         {vars.map((v) => (
           <Badge
-            key={v}
+            key={`${v.source}:${v.name}`}
             variant="secondary"
-            className="corner-squircle h-4 px-1.5 font-mono text-[10px]"
+            className={
+              v.source === "seed"
+                ? "corner-squircle h-4 border-blue-500/25 bg-blue-500/10 px-1.5 font-mono text-[10px] text-blue-700 dark:text-blue-300"
+                : "corner-squircle h-4 px-1.5 font-mono text-[10px]"
+            }
           >
-            {v}
+            {v.name}
           </Badge>
         ))}
       </div>
diff --git a/studio/frontend/src/features/recipe-studio/dialogs/shared/available-variables.tsx b/studio/frontend/src/features/recipe-studio/dialogs/shared/available-variables.tsx
index 04d8ffe356..dea66a1442 100644
--- a/studio/frontend/src/features/recipe-studio/dialogs/shared/available-variables.tsx
+++ b/studio/frontend/src/features/recipe-studio/dialogs/shared/available-variables.tsx
@@ -1,7 +1,7 @@
 import { Badge } from "@/components/ui/badge";
 import type { ReactElement } from "react";
 import { useRecipeStudioStore } from "../../stores/recipe-studio";
-import { getAvailableVariables } from "../../utils/variables";
+import { getAvailableVariableEntries } from "../../utils/variables";
 
 type AvailableVariablesProps = {
   configId: string;
@@ -11,7 +11,7 @@ export function AvailableVariables({
   configId,
 }: AvailableVariablesProps): ReactElement | null {
   const configs = useRecipeStudioStore((state) => state.configs);
-  const vars = getAvailableVariables(configs, configId);
+  const vars = getAvailableVariableEntries(configs, configId);
 
   if (vars.length === 0) return null;
 
@@ -23,11 +23,15 @@ export function AvailableVariables({
       <div className="flex flex-wrap gap-1.5">
         {vars.map((v) => (
           <Badge
-            key={v}
+            key={`${v.source}:${v.name}`}
             variant="secondary"
-            className="corner-squircle font-mono text-[11px]"
+            className={
+              v.source === "seed"
+                ? "corner-squircle border-blue-500/25 bg-blue-500/10 font-mono text-[11px] text-blue-700 dark:text-blue-300"
+                : "corner-squircle font-mono text-[11px]"
+            }
           >
-            {`{{ ${v} }}`}
+            {`{{ ${v.name} }}`}
           </Badge>
         ))}
       </div>
diff --git a/studio/frontend/src/features/recipe-studio/utils/variables.ts b/studio/frontend/src/features/recipe-studio/utils/variables.ts
index f428e8c5ca..a59cd4453c 100644
--- a/studio/frontend/src/features/recipe-studio/utils/variables.ts
+++ b/studio/frontend/src/features/recipe-studio/utils/variables.ts
@@ -1,5 +1,12 @@
 import type { NodeConfig } from "../types";
 
+export type AvailableVariableSource = "column" | "seed";
+
+export type AvailableVariableEntry = {
+  name: string;
+  source: AvailableVariableSource;
+};
+
 function getStructuredRefs(llmName: string, outputFormat: string): string[] {
   try {
     const schema = JSON.parse(outputFormat);
@@ -12,11 +19,11 @@ function getStructuredRefs(llmName: string, outputFormat: string): string[] {
   }
 }
 
-export function getAvailableVariables(
+export function getAvailableVariableEntries(
   configs: Record<string, NodeConfig>,
   currentId: string,
-): string[] {
-  const vars: string[] = [];
+): AvailableVariableEntry[] {
+  const vars: AvailableVariableEntry[] = [];
 
   for (const config of Object.values(configs)) {
     if (config.id === currentId) {
@@ -27,12 +34,12 @@ export function getAvailableVariables(
     }
 
     if (config.kind === "sampler") {
-      vars.push(config.name);
+      vars.push({ name: config.name, source: "column" });
       continue;
     }
 
     if (config.kind === "expression") {
-      vars.push(config.name);
+      vars.push({ name: config.name, source: "column" });
       continue;
     }
 
@@ -40,7 +47,7 @@ export function getAvailableVariables(
       for (const col of config.seed_columns ?? []) {
         const name = col.trim();
         if (!name) continue;
-        vars.push(name);
+        vars.push({ name, source: "seed" });
       }
       continue;
     }
@@ -49,12 +56,24 @@ export function getAvailableVariables(
       continue;
     }
 
-    vars.push(config.name);
+    vars.push({ name: config.name, source: "column" });
     if (config.llm_type !== "structured" || !config.output_format) {
       continue;
     }
-    vars.push(...getStructuredRefs(config.name, config.output_format));
+    vars.push(
+      ...getStructuredRefs(config.name, config.output_format).map((name) => ({
+        name,
+        source: "column" as const,
+      })),
+    );
   }
 
   return vars;
 }
+
+export function getAvailableVariables(
+  configs: Record<string, NodeConfig>,
+  currentId: string,
+): string[] {
+  return getAvailableVariableEntries(configs, currentId).map((entry) => entry.name);
+}