@@ -13,13 +13,16 @@ import { checkTagTrigger, TagDropdown } from '@/components/ui/tag-dropdown'
1313import { CodeLanguage } from '@/lib/execution/languages'
1414import { createLogger } from '@/lib/logs/console/logger'
1515import { cn } from '@/lib/utils'
16+ import { isLikelyReferenceSegment , SYSTEM_REFERENCE_PREFIXES } from '@/lib/workflows/references'
1617import { WandPromptBar } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/wand-prompt-bar/wand-prompt-bar'
1718import { useSubBlockValue } from '@/app/workspace/[workspaceId]/w/[workflowId]/components/workflow-block/components/sub-block/hooks/use-sub-block-value'
19+ import { useAccessibleReferencePrefixes } from '@/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-accessible-reference-prefixes'
1820import { useWand } from '@/app/workspace/[workspaceId]/w/[workflowId]/hooks/use-wand'
1921import type { GenerationType } from '@/blocks/types'
2022import { useCollaborativeWorkflow } from '@/hooks/use-collaborative-workflow'
2123import { useTagSelection } from '@/hooks/use-tag-selection'
2224import { useSubBlockStore } from '@/stores/workflows/subblock/store'
25+ import { normalizeBlockName } from '@/stores/workflows/utils'
2326
2427const logger = createLogger ( 'Code' )
2528
@@ -99,6 +102,8 @@ export function Code({
99102 const [ activeSourceBlockId , setActiveSourceBlockId ] = useState < string | null > ( null )
100103 const [ visualLineHeights , setVisualLineHeights ] = useState < number [ ] > ( [ ] )
101104
105+ const accessiblePrefixes = useAccessibleReferencePrefixes ( blockId )
106+
102107 const collapsedStateKey = `${ subBlockId } _collapsed`
103108 const isCollapsed =
104109 ( useSubBlockStore ( ( state ) => state . getValue ( blockId , collapsedStateKey ) ) as boolean ) ?? false
@@ -354,6 +359,30 @@ IMPORTANT FORMATTING RULES:
354359 } , 0 )
355360 }
356361
362+ const shouldHighlightReference = ( part : string ) : boolean => {
363+ if ( ! part . startsWith ( '<' ) || ! part . endsWith ( '>' ) ) {
364+ return false
365+ }
366+
367+ if ( ! isLikelyReferenceSegment ( part ) ) {
368+ return false
369+ }
370+
371+ if ( ! accessiblePrefixes ) {
372+ return true
373+ }
374+
375+ const inner = part . slice ( 1 , - 1 )
376+ const [ prefix ] = inner . split ( '.' )
377+ const normalizedPrefix = normalizeBlockName ( prefix )
378+
379+ if ( SYSTEM_REFERENCE_PREFIXES . has ( normalizedPrefix ) ) {
380+ return true
381+ }
382+
383+ return accessiblePrefixes . has ( normalizedPrefix )
384+ }
385+
357386 const renderLineNumbers = ( ) : ReactElement [ ] => {
358387 const numbers : ReactElement [ ] = [ ]
359388 let lineNumber = 1
@@ -490,13 +519,51 @@ IMPORTANT FORMATTING RULES:
490519 e . preventDefault ( )
491520 }
492521 } }
493- highlight = { ( codeToHighlight ) =>
494- highlight (
495- codeToHighlight ,
496- languages [ effectiveLanguage === 'python' ? 'python' : 'javascript' ] ,
497- effectiveLanguage === 'python' ? 'python' : 'javascript'
498- )
499- }
522+ highlight = { ( codeToHighlight ) => {
523+ const placeholders : { placeholder : string ; original : string ; type : 'var' | 'env' } [ ] =
524+ [ ]
525+ let processedCode = codeToHighlight
526+
527+ // Replace environment variables with placeholders
528+ processedCode = processedCode . replace ( / \{ \{ ( [ ^ } ] + ) \} \} / g, ( match ) => {
529+ const placeholder = `__ENV_VAR_${ placeholders . length } __`
530+ placeholders . push ( { placeholder, original : match , type : 'env' } )
531+ return placeholder
532+ } )
533+
534+ // Replace variable references with placeholders
535+ processedCode = processedCode . replace ( / < ( [ ^ > ] + ) > / g, ( match ) => {
536+ if ( shouldHighlightReference ( match ) ) {
537+ const placeholder = `__VAR_REF_${ placeholders . length } __`
538+ placeholders . push ( { placeholder, original : match , type : 'var' } )
539+ return placeholder
540+ }
541+ return match
542+ } )
543+
544+ // Apply Prism syntax highlighting
545+ const lang = effectiveLanguage === 'python' ? 'python' : 'javascript'
546+ let highlightedCode = highlight ( processedCode , languages [ lang ] , lang )
547+
548+ // Restore and highlight the placeholders
549+ placeholders . forEach ( ( { placeholder, original, type } ) => {
550+ if ( type === 'env' ) {
551+ highlightedCode = highlightedCode . replace (
552+ placeholder ,
553+ `<span class="text-blue-500">${ original } </span>`
554+ )
555+ } else if ( type === 'var' ) {
556+ // Escape the < and > for display
557+ const escaped = original . replace ( / < / g, '<' ) . replace ( / > / g, '>' )
558+ highlightedCode = highlightedCode . replace (
559+ placeholder ,
560+ `<span class="text-blue-500">${ escaped } </span>`
561+ )
562+ }
563+ } )
564+
565+ return highlightedCode
566+ } }
500567 padding = { 12 }
501568 style = { {
502569 fontFamily : 'inherit' ,
0 commit comments