implement variables in functions

src/App.tsx

@@ -270,7 +270,7 @@ export default function App() {
           <ContextMenuContent className="w-48">
             <ContextMenuGroup>
               <ContextMenuSub>
-                  <ContextMenuSubTrigger>Create node</ContextMenuSubTrigger>
+                  <ContextMenuSubTrigger>Create Node</ContextMenuSubTrigger>
                   <ContextMenuSubContent className="w-44">
                     <ContextMenuGroup>
                       <ContextMenuItem onSelect={() => createNode('config')}>

src/components/graph/BaseNode.tsx

@@ -173,7 +173,7 @@ export function BaseNodeFooter({ className, ...props }: ComponentProps<"div">) {
     <div
       data-slot="base-node-footer"
       className={cn(
-        "shrink-0 flex flex-col items-center gap-y-2 border-t px-3 pt-2 pb-3",
+        "shrink-0 flex flex-col items-center gap-y-2 border-t p-1",
         className,
       )}
       {...props}

src/components/graph/FunctionNode.tsx

@@ -1,4 +1,4 @@
-import React, { memo, useCallback, useContext, useMemo } from 'react'
+import React, { memo, useCallback, useContext, useMemo, useRef } from 'react'
 import CodeMirror from '@uiw/react-codemirror'
 import { javascript } from '@codemirror/lang-javascript'
 import FlowContext from '../../lib/flowContext'
@@ -14,7 +14,8 @@ import { InputHandle, OutputHandle } from './NodeHandles'
 import { NodeFooterEdgeIndicators } from './NodeFooterEdgeIndicators'
 import { NodeHeaderTitle } from './NodeHeaderTitle'
 import { NodeMenubar } from './NodeMenubar'
-import { Code2 } from 'lucide-react'
+import { MenubarItem, MenubarSub, MenubarSubContent, MenubarSubTrigger } from '../ui/menubar'
+import { Code2, Variable } from 'lucide-react'
 
 type Props = {
     id: string
@@ -28,6 +29,17 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
     const { theme } = useTheme()
     const ctx = useContext(FlowContext)
     const setNodes = ctx?.setNodes
+    const editorRef = useRef<unknown>(null)
+
+    const edges = ctx?.edges ?? []
+    const nodes = ctx?.nodes ?? []
+    const incomingEdges = useMemo(() => edges.filter((e: any) => e.target === id), [edges, id])
+    const incomingIds = useMemo(() => incomingEdges.map((e: any) => e.source).sort(), [incomingEdges])
+    const connectedVariableNodes = useMemo(
+        () => (nodes as any[]).filter((n: any) => incomingIds.includes(n.id) && n.type === 'variable'),
+        [nodes, incomingIds]
+    )
+    const hasConnectedVariables = connectedVariableNodes.length > 0
 
     const onChange = useCallback(
         (val: string) => {
@@ -40,6 +52,34 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
         [id, setNodes]
     )
 
+    const insertAt = useCallback(
+        (insertText: string, mode: 'prepend' | 'append' | 'cursor') => {
+            const ref = editorRef.current as { view: { state: { doc: { length: number }; selection: { main: { from: number } } }; dispatch: (arg: { changes: { from: number; to: number; insert: string } }) => void } } | null
+            if (ref?.view) {
+                const view = ref.view
+                const doc = view.state.doc
+                const len = doc.length
+                let from: number
+                if (mode === 'prepend') from = 0
+                else if (mode === 'append') from = len
+                else from = view.state.selection.main.from
+                view.dispatch({ changes: { from, to: from, insert: insertText } })
+                onChange(view.state.doc.toString())
+                return
+            }
+            if (mode === 'prepend') onChange(insertText + bodyValue)
+            else onChange(bodyValue + insertText)
+        },
+        [onChange, bodyValue]
+    )
+
+    const insertVariableAtCursor = useCallback(
+        (variableNode: any) => {
+            insertAt(variableNode.id, 'cursor')
+        },
+        [insertAt]
+    )
+
     const extensions = useMemo(() => [javascript()], [])
     const [editorHeight, editorContainerRef] = useResizeHeight(120)
     const dimensions =
@@ -53,10 +93,37 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
 
             <BaseNodeContent>
                 <div className="shrink-0 w-full">
-                    <NodeMenubar nodeId={id} nodeType="function" />
+                    <NodeMenubar
+                        nodeId={id}
+                        nodeType="function"
+                        editInputsContent={
+                            hasConnectedVariables ? (
+                                <>
+                                    {connectedVariableNodes.map((n: any) => (
+                                        <MenubarSub key={n.id}>
+                                            <MenubarSubTrigger className="text-xs flex items-center gap-2">
+                                                <Variable className="size-3.5 shrink-0" />
+                                                {n.id}
+                                            </MenubarSubTrigger>
+                                            <MenubarSubContent>
+                                                <MenubarItem
+                                                    className="text-xs"
+                                                    onClick={() => insertVariableAtCursor(n)}
+                                                >
+                                                    Insert at cursor
+                                                </MenubarItem>
+                                            </MenubarSubContent>
+                                        </MenubarSub>
+                                    ))}
+                                </>
+                            ) : undefined
+                        }
+                    />
                 </div>
                 <div ref={editorContainerRef} className="min-h-0 flex-1 w-full nodrag nopan overflow-hidden border-t border-input">
                     <CodeMirror
+                        // @ts-expect-error ref is { view, state, editor }; package ref type not in our node_modules
+                        ref={editorRef}
                         value={bodyValue}
                         height={`${editorHeight}px`}
                         theme={theme}
@@ -72,9 +139,6 @@ export const FunctionNode = memo(function FunctionNode({ id, data, width, height
                 <NodeFooterEdgeIndicators nodeId={id} nodeType="function">
                     {bodyValue ? `${bodyValue.length} chars` : 'none'}
                 </NodeFooterEdgeIndicators>
-                <p className="text-[10px] text-muted-foreground mt-1">
-                    In config: <code className="rounded bg-muted px-0.5">{'{{ x | '}{id}{' }}'}</code> or <code>{'{{ x | '}{id}{'(a, b, key=val) }}'}</code>. Use <code>function(num, x, y, kwargs) &#123; ... &#125;</code> — kwargs has keyword args.
-                </p>
             </BaseNodeFooter>
         </BaseNode>
     )

src/components/graph/RenderingNode.tsx

@@ -234,16 +234,28 @@ export const RenderingNode = memo(function RenderingNode({ id, width, height }:
                     },
                 }
 
-                // Context: only variables. Function nodes are registered as Nunjucks custom filters (see below).
+                // Context: variables connected to configs, plus variables connected to functions that feed configs (so they can be injected as constants).
                 const nunjucksContext = Object.create(null) as Record<string, unknown>
+                const setVarInContext = (src: any) => {
+                    const v = src.data?.value
+                    const str = v === undefined || v === null ? '' : String(v)
+                    nunjucksContext[src.id] =
+                        v === undefined || v === null ? '' : typeof v === 'boolean' || typeof v === 'number' ? v : str
+                }
                 for (const e of edges) {
                     if (!configIdsUsed.has(e.target)) continue
                     const src = nodes.find((n: any) => n.id === e.source)
-                    if (src?.type === 'variable') {
+                    if (src?.type === 'variable') setVarInContext(src)
-                        const v = src.data?.value
+                }
-                        const str = v === undefined || v === null ? '' : String(v)
+                for (const e of edges) {
-                        nunjucksContext[src.id] =
+                    if (!configIdsUsed.has(e.target)) continue
-                            v === undefined || v === null ? '' : typeof v === 'boolean' || typeof v === 'number' ? v : str
+                    const src = nodes.find((n: any) => n.id === e.source)
+                    if (src?.type === 'function') {
+                        for (const e2 of edges) {
+                            if (e2.target !== src.id) continue
+                            const vNode = nodes.find((n: any) => n.id === e2.source)
+                            if (vNode?.type === 'variable') setVarInContext(vNode)
+                        }
                     }
                 }
 
@@ -280,12 +292,31 @@ export const RenderingNode = memo(function RenderingNode({ id, width, height }:
                 const isPlainObject = (v: unknown): v is Record<string, unknown> =>
                     typeof v === 'object' && v !== null && !Array.isArray(v)
 
+                // For each function node that feeds a config: which variable node ids are connected to that function?
+                const functionConnectedVariableIds = Object.create(null) as Record<string, string[]>
+                for (const e of edges) {
+                    if (!configIdsUsed.has(e.target)) continue
+                    const src = nodes.find((n: any) => n.id === e.source)
+                    if (src?.type === 'function') {
+                        const fid = src.id
+                        for (const e2 of edges) {
+                            if (e2.target !== fid) continue
+                            const vNode = nodes.find((n: any) => n.id === e2.source)
+                            if (vNode?.type === 'variable') {
+                                if (!functionConnectedVariableIds[fid]) functionConnectedVariableIds[fid] = []
+                                functionConnectedVariableIds[fid].push(vNode.id)
+                            }
+                        }
+                    }
+                }
+
                 for (const e of edges) {
                     if (!configIdsUsed.has(e.target)) continue
                     const src = nodes.find((n: any) => n.id === e.source)
                     if (src?.type === 'function') {
                         const body = src.data?.body ?? 'return args[0];'
                         const parsed = parseFunctionSignature(body)
+                        const connectedVarIds = new Set<string>(functionConnectedVariableIds[src.id] ?? [])
                         env.addFilter(
                             src.id,
                             (value: unknown, ...args: unknown[]) => {
@@ -301,10 +332,17 @@ export const RenderingNode = memo(function RenderingNode({ id, width, height }:
                                     const lastParam = paramNames[paramNames.length - 1]
                                     const invocationArgs = paramNames.map((name, i) => {
                                         if (name === lastParam && lastParam === 'kwargs') return kwargs
+                                        // Connected variables are available as constants: use variable value from context
+                                        if (connectedVarIds.has(name) && name in nunjucksContext)
+                                            return nunjucksContext[name]
                                         return positionals[i]
                                     })
-                                    const fn = new Function(...paramNames, innerBody)
+                                    // Inject connected variables as extra params so they're in scope in the body (e.g. (value) => value + var_001)
-                                    invoke = () => fn(...invocationArgs)
+                                    const extraVarIds = [...connectedVarIds].filter((vid) => !paramNames.includes(vid))
+                                    const allParamNames = [...paramNames, ...extraVarIds]
+                                    const allArgs = [...invocationArgs, ...extraVarIds.map((vid) => nunjucksContext[vid])]
+                                    const fn = new Function(...allParamNames, innerBody)
+                                    invoke = () => fn(...allArgs)
                                 } else {
                                     const fn = new Function('args', body)
                                     invoke = () => fn(positionals)