refactor: refactor token management and variable collection tools

- Introduced `config.ts` to manage current code generation configuration.
- Created `candidates.ts` for collecting candidate variable IDs from scene nodes.
- Developed `defs.ts` to handle token definitions and resolve tokens by names.
- Added `index.ts` to export candidate variable collection and token definition functions.
- Implemented `indexer.ts` for indexing tokens and canonicalizing variable names.
- Enhanced `worker.ts` to cache plugin code and transform variables efficiently.
- Updated `codegen.ts` with a utility function for worker options.
- Refactored `css.ts` to improve variable function handling and normalization.
- Modified server tools to accurately count token references in responses.
- Adjusted shared tools to clarify token definition schema documentation.

Author: Justineo

packages/extension/mcp/runtime.ts

@@ -14,7 +14,7 @@ import { selection } from '@/ui/state'
 import { handleGetCode as runGetCode } from './tools/code'
 import { handleGetScreenshot as runGetScreenshot } from './tools/screenshot'
 import { handleGetStructure as runGetStructure } from './tools/structure'
-import { handleGetTokenDefs as runGetTokenDefs } from './tools/token-defs'
+import { handleGetTokenDefs as runGetTokenDefs } from './tools/token'
 
 function isSceneNode(node: BaseNode | null): node is SceneNode {
   return !!node && 'visible' in node && 'type' in node

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

@@ -1,16 +1,27 @@
-import type { AssetDescriptor, GetCodeResult, GetTokenDefsResult } from '@tempad-dev/mcp-shared'
+import type {
+  AssetDescriptor,
+  GetCodeResult,
+  GetTokenDefsResult,
+  TokenEntry
+} from '@tempad-dev/mcp-shared'
 
 import { MCP_MAX_PAYLOAD_BYTES } from '@tempad-dev/mcp-shared'
 
-import type { CodegenConfig } from '@/utils/codegen'
-
 import { buildSemanticTree } from '@/mcp/semantic-tree'
-import { activePlugin, options } from '@/ui/state'
+import { activePlugin } from '@/ui/state'
 import { stringifyComponent } from '@/utils/component'
+import {
+  extractVarNames,
+  normalizeCssVarName,
+  replaceVarFunctions,
+  stripFallback,
+  toVarExpr
+} from '@/utils/css'
 
 import type { RenderContext, CodeLanguage } from './render'
 
-import { collectTokenReferences, resolveVariableTokens } from '../token-defs'
+import { currentCodegenConfig } from '../config'
+import { collectCandidateVariableIds, resolveTokenDefsByNames } from '../token'
 import { collectSceneData } from './collect'
 import { buildGetCodeMessage } from './messages'
 import { renderSemanticNode } from './render'
@@ -52,17 +63,7 @@ const COMPACT_TAGS = new Set([
   'summary'
 ])
 
-function filterTokensByNames(input: GetTokenDefsResult, names: Set<string>): GetTokenDefsResult {
-  const tokens: Record<string, GetTokenDefsResult['tokens'][string]> = {}
-  Object.entries(input.tokens).forEach(([name, entry]) => {
-    if (names.has(name)) {
-      tokens[name] = entry
-    }
-  })
-  return { tokens }
-}
-
-function primaryTokenValue(entry: GetTokenDefsResult['tokens'][string]): string | undefined {
+function primaryTokenValue(entry: TokenEntry): string | undefined {
   if (typeof entry.value === 'string') return entry.value
   if (typeof entry.resolvedValue === 'string') return entry.resolvedValue
   if (entry.activeMode && typeof entry.value[entry.activeMode] === 'string') {
@@ -72,6 +73,34 @@ function primaryTokenValue(entry: GetTokenDefsResult['tokens'][string]): string
   return typeof first === 'string' ? first : undefined
 }
 
+function resolveConcreteValue(
+  name: string,
+  allTokens: GetTokenDefsResult,
+  cache: Map<string, string | undefined>,
+  depth = 0
+): string | undefined {
+  if (cache.has(name)) return cache.get(name)
+  if (depth > 20) return undefined
+
+  const entry = allTokens[name]
+  if (!entry) {
+    cache.set(name, undefined)
+    return undefined
+  }
+  const val = primaryTokenValue(entry)
+  if (typeof val !== 'string') {
+    cache.set(name, undefined)
+    return undefined
+  }
+  if (val.startsWith('--') && val !== name) {
+    const resolved = resolveConcreteValue(val, allTokens, cache, depth + 1) ?? val
+    cache.set(name, resolved)
+    return resolved
+  }
+  cache.set(name, val)
+  return val
+}
+
 export async function handleGetCode(
   nodes: SceneNode[],
   preferredLang?: CodeLanguage,
@@ -92,13 +121,13 @@ export async function handleGetCode(
     throw new Error('No renderable nodes found for the current selection.')
   }
 
-  const config = codegenConfig()
+  const config = currentCodegenConfig()
   const pluginCode = activePlugin.value?.code
 
   const assetRegistry = new Map<string, AssetDescriptor>()
   const { nodes: nodeMap, styles, svgs } = await collectSceneData(tree.roots, config, assetRegistry)
 
-  await applyVariableTransforms(styles, {
+  const styleVarNames = await applyVariableTransforms(styles, {
     pluginCode,
     config
   })
@@ -131,33 +160,48 @@ export async function handleGetCode(
   const MAX_CODE_CHARS = Math.floor(MCP_MAX_PAYLOAD_BYTES * 0.6)
   const { markup, message } = buildGetCodeMessage(rawMarkup, MAX_CODE_CHARS, tree.stats)
 
-  const { variableIds } = collectTokenReferences(nodes)
-  const allTokens = await resolveVariableTokens(variableIds, config, pluginCode)
-
-  const usedTokenNames = new Set<string>()
-  // Use a simple regex to capture all var(--name) occurrences, including nested ones.
-  // We don't use CSS_VAR_FUNCTION_RE because it consumes the whole function and might miss nested vars in fallbacks.
-  const regex = /var\(--([a-zA-Z0-9-_]+)/g
-  let match
-  while ((match = regex.exec(markup))) {
-    usedTokenNames.add(`--${match[1]}`)
-  }
+  // Only include tokens actually referenced in the final output.
+  const usedTokenNames = new Set<string>(styleVarNames)
+  extractVarNames(markup).forEach((n) => usedTokenNames.add(n))
 
-  const usedTokens = filterTokensByNames(allTokens, usedTokenNames)
+  const usedTokens = await resolveTokenDefsByNames(usedTokenNames, config, pluginCode, {
+    candidateIds: () => collectCandidateVariableIds(nodes).variableIds
+  })
 
   const codegen = {
     preset: activePlugin.value?.name ?? 'none',
     config
   }
 
   if (resolveTokens) {
-    const tokenMap = new Map(
-      Object.entries(usedTokens.tokens).map(([name, entry]) => [name, primaryTokenValue(entry)])
-    )
-    const resolvedMarkup = markup.replace(/var\((--[a-zA-Z0-9-_]+)\)/g, (match, name) => {
-      const val = tokenMap.get(name)
-      return typeof val === 'string' ? val : match
+    const tokenCache = new Map<string, string | undefined>()
+    const tokenMap = new Map<string, string | undefined>()
+    Object.keys(usedTokens).forEach((name) => {
+      tokenMap.set(name, resolveConcreteValue(name, usedTokens, tokenCache))
     })
+
+    const replaceToken = (input: string): string => {
+      // Always strip fallbacks (even if token isn't resolved).
+      let out = stripFallback(input)
+
+      // Replace CSS var() functions first (supports whitespace/nesting).
+      out = replaceVarFunctions(out, ({ name, full }) => {
+        const trimmed = name.trim()
+        if (!trimmed.startsWith('--')) return full
+        const canonical = `--${normalizeCssVarName(trimmed.slice(2))}`
+        const val = tokenMap.get(canonical)
+        return typeof val === 'string' ? val : toVarExpr(canonical)
+      })
+
+      // replace bare --token (e.g., Tailwind arbitrary value: border-[--foo])
+      out = out.replace(/--[A-Za-z0-9-_]+/g, (m) => {
+        const val = tokenMap.get(m)
+        return typeof val === 'string' ? val : m
+      })
+      return out
+    }
+
+    const resolvedMarkup = replaceToken(markup)
     // If tokens are resolved, we don't need to return the token definitions
     return {
       lang: resolvedLang,
@@ -191,8 +235,3 @@ function normalizeRootString(content: string, fallbackTag: string | undefined, l
     }
   )
 }
-
-function codegenConfig(): CodegenConfig {
-  const { cssUnit, rootFontSize, scale } = options.value
-  return { cssUnit, rootFontSize, scale }
-}

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

@@ -1,15 +1,19 @@
 import type { CodegenConfig } from '@/utils/codegen'
 
 import { runTransformVariableBatch } from '@/mcp/transform-variables/requester'
+import { workerUnitOptions } from '@/utils/codegen'
 import {
-  CSS_VAR_FUNCTION_RE,
-  canonicalizeVariable,
+  canonicalizeVarName,
   expandShorthands,
+  extractVarNames,
   formatHexAlpha,
   normalizeCssVarName,
-  normalizeStyleVariables,
+  preprocessCssValue,
+  replaceVarFunctions,
+  stripFallback,
   normalizeStyleValues,
-  parseBackgroundShorthand
+  parseBackgroundShorthand,
+  toVarExpr
 } from '@/utils/css'
 import { cssToClassNames } from '@/utils/tailwind'
 
@@ -357,36 +361,38 @@ type ApplyVariableOptions = {
 export async function applyVariableTransforms(
   styles: Map<string, Record<string, string>>,
   { pluginCode, config }: ApplyVariableOptions
-): Promise<void> {
+): Promise<Set<string>> {
   const { references, buckets } = collectVariableReferences(styles)
-  if (!references.length) return
+  if (!references.length) return new Set()
 
   const transformResults = await runTransformVariableBatch(
     references.map(({ code, name, value }) => ({ code, name, value })),
-    {
-      useRem: config.cssUnit === 'rem',
-      rootFontSize: config.rootFontSize ?? 16,
-      scale: config.scale ?? 1
-    },
+    workerUnitOptions(config),
     pluginCode
   )
 
   const replacements = transformResults.map((result) => {
-    return canonicalizeVariable(result) || result
+    const noFallback = stripFallback(result)
+    const canonicalName = canonicalizeVarName(noFallback)
+    return canonicalName ? toVarExpr(canonicalName) : noFallback
   })
-
-  const safeRegex = new RegExp(CSS_VAR_FUNCTION_RE.source, CSS_VAR_FUNCTION_RE.flags)
+  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] = bucket.value.replace(safeRegex, (match: string) => {
+    style[bucket.property] = replaceVarFunctions(bucket.value, ({ full }) => {
       const refIndex = bucket.matchIndices[occurrence++]
-      return replacements[refIndex] ?? match
+      return refIndex != null ? (replacements[refIndex] ?? full) : full
     })
   }
+
+  return usedNames
 }
 
 export function stripInertShadows(style: Record<string, string>, node: SceneNode): void {
@@ -418,38 +424,41 @@ function isFillRenderable(fill: Paint | undefined): boolean {
 function collectVariableReferences(styles: Map<string, Record<string, string>>) {
   const references: VariableReferenceInternal[] = []
   const buckets = new Map<string, PropertyBucket>()
-  const regex = new RegExp(CSS_VAR_FUNCTION_RE.source, CSS_VAR_FUNCTION_RE.flags)
 
   for (const [nodeId, style] of styles.entries()) {
-    normalizeStyleVariables(style)
-
     for (const [property, value] of Object.entries(style)) {
-      let match: RegExpExecArray | null
       let hasMatch = false
       const indices: number[] = []
 
-      regex.lastIndex = 0
-      while ((match = regex.exec(value))) {
+      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 [, name, fallback] = match
         const refIndex =
           references.push({
             nodeId,
             property,
-            code: match[0],
-            name: normalizeCssVarName(name),
+            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, matchIndices: indices })
+          buckets.set(key, { nodeId, property, value: normalized, matchIndices: indices })
         }
       }
     }

packages/extension/mcp/tools/config.ts

@@ -0,0 +1,8 @@
+import type { CodegenConfig } from '@/utils/codegen'
+
+import { options } from '@/ui/state'
+
+export function currentCodegenConfig(): CodegenConfig {
+  const { cssUnit, rootFontSize, scale } = options.value
+  return { cssUnit, rootFontSize, scale }
+}