refactor: replace applyVariableTransforms with transform function and clean up variable handling, fix backgrond-position

Author: Justineo

packages/extension/mcp/tools/code/index.ts

@@ -25,7 +25,7 @@ import { collectCandidateVariableIds, resolveTokenDefsByNames } from '../token'
 import { collectSceneData } from './collect'
 import { buildGetCodeMessage } from './messages'
 import { renderSemanticNode } from './render'
-import { applyVariableTransforms } from './style'
+import { transform } from './variables'
 
 // Tags that should render children without extra whitespace/newlines.
 const COMPACT_TAGS = new Set([
@@ -127,7 +127,7 @@ export async function handleGetCode(
   const assetRegistry = new Map<string, AssetDescriptor>()
   const { nodes: nodeMap, styles, svgs } = await collectSceneData(tree.roots, config, assetRegistry)
 
-  const styleVarNames = await applyVariableTransforms(styles, {
+  const styleVarNames = await transform(styles, {
     pluginCode,
     config
   })

packages/extension/mcp/tools/code/style.ts

@@ -1,19 +1,10 @@
 import type { CodegenConfig } from '@/utils/codegen'
 
-import { runTransformVariableBatch } from '@/mcp/transform-variables/requester'
-import { workerUnitOptions } from '@/utils/codegen'
 import {
-  canonicalizeVarName,
   expandShorthands,
-  extractVarNames,
   formatHexAlpha,
-  normalizeCssVarName,
-  preprocessCssValue,
-  replaceVarFunctions,
-  stripFallback,
   normalizeStyleValues,
-  parseBackgroundShorthand,
-  toVarExpr
+  parseBackgroundShorthand
 } from '@/utils/css'
 import { cssToClassNames } from '@/utils/tailwind'
 
@@ -338,63 +329,6 @@ function hasAtomicPadding(s: Record<string, string>) {
   return !!(s['padding-top'] || s['padding-right'] || s['padding-bottom'] || s['padding-left'])
 }
 
-type VariableReferenceInternal = {
-  nodeId: string
-  property: string
-  code: string
-  name: string
-  value?: string
-}
-
-type PropertyBucket = {
-  nodeId: string
-  property: string
-  value: string
-  matchIndices: number[]
-}
-
-type ApplyVariableOptions = {
-  pluginCode?: string
-  config: CodegenConfig
-}
-
-export async function applyVariableTransforms(
-  styles: Map<string, Record<string, string>>,
-  { pluginCode, config }: ApplyVariableOptions
-): Promise<Set<string>> {
-  const { references, buckets } = collectVariableReferences(styles)
-  if (!references.length) return new Set()
-
-  const transformResults = await runTransformVariableBatch(
-    references.map(({ code, name, value }) => ({ code, name, value })),
-    workerUnitOptions(config),
-    pluginCode
-  )
-
-  const replacements = transformResults.map((result) => {
-    const noFallback = stripFallback(result)
-    const canonicalName = canonicalizeVarName(noFallback)
-    return canonicalName ? toVarExpr(canonicalName) : noFallback
-  })
-  const usedNames = new Set<string>()
-  // Keep original reference names to handle transforms that inline literal values.
-  references.forEach(({ name }) => usedNames.add(`--${name}`))
-  replacements.forEach((repl) => extractVarNames(repl).forEach((n) => usedNames.add(n)))
-
-  for (const bucket of buckets.values()) {
-    const style = styles.get(bucket.nodeId)
-    if (!style) continue
-
-    let occurrence = 0
-    style[bucket.property] = replaceVarFunctions(bucket.value, ({ full }) => {
-      const refIndex = bucket.matchIndices[occurrence++]
-      return refIndex != null ? (replacements[refIndex] ?? full) : full
-    })
-  }
-
-  return usedNames
-}
-
 export function stripInertShadows(style: Record<string, string>, node: SceneNode): void {
   if (!style['box-shadow']) return
   if (hasRenderableFill(node)) return
@@ -421,51 +355,6 @@ function isFillRenderable(fill: Paint | undefined): boolean {
   return true
 }
 
-function collectVariableReferences(styles: Map<string, Record<string, string>>) {
-  const references: VariableReferenceInternal[] = []
-  const buckets = new Map<string, PropertyBucket>()
-
-  for (const [nodeId, style] of styles.entries()) {
-    for (const [property, value] of Object.entries(style)) {
-      let hasMatch = false
-      const indices: number[] = []
-
-      const normalized = preprocessCssValue(value)
-      if (normalized !== value) {
-        style[property] = normalized
-      }
-
-      replaceVarFunctions(normalized, ({ full, name, fallback }) => {
-        const trimmed = name.trim()
-        if (!trimmed.startsWith('--')) return full
-
-        hasMatch = true
-        const refIndex =
-          references.push({
-            nodeId,
-            property,
-            code: full,
-            name: normalizeCssVarName(trimmed.slice(2)),
-            value: fallback?.trim()
-          }) - 1
-        indices.push(refIndex)
-        return full
-      })
-
-      if (hasMatch) {
-        const key = `${nodeId}:${property}`
-        const bucket = buckets.get(key)
-        if (bucket) {
-          bucket.matchIndices.push(...indices)
-        } else {
-          buckets.set(key, { nodeId, property, value: normalized, matchIndices: indices })
-        }
-      }
-    }
-  }
-  return { references, buckets }
-}
-
 export function styleToClassNames(
   style: Record<string, string>,
   config: CodegenConfig,

packages/extension/mcp/tools/code/text/render.ts

@@ -5,7 +5,8 @@ import { joinClassNames } from '@/utils/tailwind'
 
 import type { RenderContext } from '../render'
 
-import { applyVariableTransforms, styleToClassNames } from '../style'
+import { styleToClassNames } from '../style'
+import { transform } from '../variables'
 import { buildTextBlocks, formatTextLiteral, getStyledSegments } from './segments'
 import { computeDominantStyle, omitCommon } from './style'
 import {
@@ -53,7 +54,7 @@ export async function renderTextSegments(
     }
   }
 
-  await applyVariableTransforms(runStyleMap, {
+  await transform(runStyleMap, {
     pluginCode: ctx.pluginCode,
     config: ctx.config
   })

packages/extension/mcp/tools/code/variables.ts

@@ -0,0 +1,115 @@
+import type { CodegenConfig } from '@/utils/codegen'
+
+import { runTransformVariableBatch } from '@/mcp/transform-variables/requester'
+import { workerUnitOptions } from '@/utils/codegen'
+import {
+  canonicalizeVarName,
+  extractVarNames,
+  normalizeCssVarName,
+  preprocessCssValue,
+  replaceVarFunctions,
+  stripFallback,
+  toVarExpr
+} from '@/utils/css'
+
+type VariableReferenceInternal = {
+  nodeId: string
+  property: string
+  code: string
+  name: string
+  value?: string
+}
+
+type PropertyBucket = {
+  nodeId: string
+  property: string
+  value: string
+  matchIndices: number[]
+}
+
+type TransformOptions = {
+  pluginCode?: string
+  config: CodegenConfig
+}
+
+export function transform(
+  styles: Map<string, Record<string, string>>,
+  { pluginCode, config }: TransformOptions
+): Promise<Set<string>> {
+  const { references, buckets } = collectRefs(styles)
+  if (!references.length) return Promise.resolve(new Set())
+
+  return runTransformVariableBatch(
+    references.map(({ code, name, value }) => ({ code, name, value })),
+    workerUnitOptions(config),
+    pluginCode
+  ).then((transformResults) => {
+    const replacements = transformResults.map((result) => {
+      const noFallback = stripFallback(result)
+      const canonicalName = canonicalizeVarName(noFallback)
+      return canonicalName ? toVarExpr(canonicalName) : noFallback
+    })
+    const usedNames = new Set<string>()
+    // Keep original reference names to handle transforms that inline literal values.
+    references.forEach(({ name }) => usedNames.add(`--${name}`))
+    replacements.forEach((repl) => extractVarNames(repl).forEach((n) => usedNames.add(n)))
+
+    for (const bucket of buckets.values()) {
+      const style = styles.get(bucket.nodeId)
+      if (!style) continue
+
+      let occurrence = 0
+      style[bucket.property] = replaceVarFunctions(bucket.value, ({ full }) => {
+        const refIndex = bucket.matchIndices[occurrence++]
+        return refIndex != null ? (replacements[refIndex] ?? full) : full
+      })
+    }
+
+    return usedNames
+  })
+}
+
+export function collectRefs(styles: Map<string, Record<string, string>>) {
+  const references: VariableReferenceInternal[] = []
+  const buckets = new Map<string, PropertyBucket>()
+
+  for (const [nodeId, style] of styles.entries()) {
+    for (const [property, value] of Object.entries(style)) {
+      let hasMatch = false
+      const indices: number[] = []
+
+      const normalized = preprocessCssValue(value)
+      if (normalized !== value) {
+        style[property] = normalized
+      }
+
+      replaceVarFunctions(normalized, ({ full, name, fallback }) => {
+        const trimmed = name.trim()
+        if (!trimmed.startsWith('--')) return full
+
+        hasMatch = true
+        const refIndex =
+          references.push({
+            nodeId,
+            property,
+            code: full,
+            name: normalizeCssVarName(trimmed.slice(2)),
+            value: fallback?.trim()
+          }) - 1
+        indices.push(refIndex)
+        return full
+      })
+
+      if (hasMatch) {
+        const key = `${nodeId}:${property}`
+        const bucket = buckets.get(key)
+        if (bucket) {
+          bucket.matchIndices.push(...indices)
+        } else {
+          buckets.set(key, { nodeId, property, value: normalized, matchIndices: indices })
+        }
+      }
+    }
+  }
+  return { references, buckets }
+}