@@ -12,7 +12,14 @@ import { useWorkflowStore } from '@/stores/workflows/workflow/store'
1212import 'prismjs/components/prism-javascript'
1313import 'prismjs/themes/prism.css'
1414
15+ import {
16+ isLikelyReferenceSegment ,
17+ SYSTEM_REFERENCE_PREFIXES ,
18+ splitReferenceSegment ,
19+ } from '@/lib/workflows/references'
1520import type { LoopType , ParallelType } from '@/lib/workflows/types'
21+ import { useAccessibleReferencePrefixes } from '@/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-accessible-reference-prefixes'
22+ import { normalizeBlockName } from '@/stores/workflows/utils'
1623
1724type IterationType = 'loop' | 'parallel'
1825
@@ -130,6 +137,88 @@ export function IterationBadges({ nodeId, data, iterationType }: IterationBadges
130137 collaborativeUpdateIterationCount,
131138 collaborativeUpdateIterationCollection,
132139 } = useCollaborativeWorkflow ( )
140+ const accessiblePrefixes = useAccessibleReferencePrefixes ( nodeId )
141+
142+ const shouldHighlightReference = useCallback (
143+ ( part : string ) : boolean => {
144+ if ( ! part . startsWith ( '<' ) || ! part . endsWith ( '>' ) ) {
145+ return false
146+ }
147+
148+ if ( ! isLikelyReferenceSegment ( part ) ) {
149+ return false
150+ }
151+
152+ const split = splitReferenceSegment ( part )
153+ if ( ! split ) {
154+ return false
155+ }
156+
157+ const reference = split . reference
158+
159+ if ( ! accessiblePrefixes ) {
160+ return true
161+ }
162+
163+ const inner = reference . slice ( 1 , - 1 )
164+ const [ prefix ] = inner . split ( '.' )
165+ const normalizedPrefix = normalizeBlockName ( prefix )
166+
167+ if ( SYSTEM_REFERENCE_PREFIXES . has ( normalizedPrefix ) ) {
168+ return true
169+ }
170+
171+ return accessiblePrefixes . has ( normalizedPrefix )
172+ } ,
173+ [ accessiblePrefixes ]
174+ )
175+
176+ const highlightWithReferences = useCallback (
177+ ( code : string ) : string => {
178+ const placeholders : Array < {
179+ placeholder : string
180+ original : string
181+ type : 'var' | 'env'
182+ } > = [ ]
183+
184+ let processedCode = code
185+
186+ processedCode = processedCode . replace ( / \{ \{ ( [ ^ } ] + ) \} \} / g, ( match ) => {
187+ const placeholder = `__ENV_VAR_${ placeholders . length } __`
188+ placeholders . push ( { placeholder, original : match , type : 'env' } )
189+ return placeholder
190+ } )
191+
192+ processedCode = processedCode . replace ( / < [ ^ > ] + > / g, ( match ) => {
193+ if ( shouldHighlightReference ( match ) ) {
194+ const placeholder = `__VAR_REF_${ placeholders . length } __`
195+ placeholders . push ( { placeholder, original : match , type : 'var' } )
196+ return placeholder
197+ }
198+ return match
199+ } )
200+
201+ let highlightedCode = highlight ( processedCode , languages . javascript , 'javascript' )
202+
203+ placeholders . forEach ( ( { placeholder, original, type } ) => {
204+ if ( type === 'env' ) {
205+ highlightedCode = highlightedCode . replace (
206+ placeholder ,
207+ `<span class="text-blue-500">${ original } </span>`
208+ )
209+ } else {
210+ const escaped = original . replace ( / < / g, '<' ) . replace ( / > / g, '>' )
211+ highlightedCode = highlightedCode . replace (
212+ placeholder ,
213+ `<span class="text-blue-500">${ escaped } </span>`
214+ )
215+ }
216+ } )
217+
218+ return highlightedCode
219+ } ,
220+ [ shouldHighlightReference ]
221+ )
133222
134223 // Handle type change
135224 const handleTypeChange = useCallback (
@@ -325,7 +414,7 @@ export function IterationBadges({ nodeId, data, iterationType }: IterationBadges
325414 < Editor
326415 value = { conditionString }
327416 onValueChange = { handleEditorChange }
328- highlight = { ( code ) => highlight ( code , languages . javascript , 'javascript' ) }
417+ highlight = { highlightWithReferences }
329418 padding = { 0 }
330419 style = { {
331420 fontFamily : 'monospace' ,
@@ -363,7 +452,7 @@ export function IterationBadges({ nodeId, data, iterationType }: IterationBadges
363452 < Editor
364453 value = { editorValue }
365454 onValueChange = { handleEditorChange }
366- highlight = { ( code ) => highlight ( code , languages . javascript , 'javascript' ) }
455+ highlight = { highlightWithReferences }
367456 padding = { 0 }
368457 style = { {
369458 fontFamily : 'monospace' ,
0 commit comments