feat: improve get_code performance

Author: Justineo

docs/extension/design.md

@@ -8,43 +8,47 @@ This document describes the implementation design for MCP `get_code` in `package
    - Exactly one visible root node required.
    - Build a visible tree with depth capping.
 
-2. **Collect data**
+2. **Plan assets and plugin overrides**
+   - Plan vector roots (vector-only containers) and mark their descendants for skipping.
+   - If a plugin is enabled, pre-resolve plugin components for instances.
+   - When a plugin returns component/code, skip collecting its descendants and reuse the cached plugin output in render.
+
+3. **Collect data**
    - `collectNodeData()`
-     - `getCSSAsync()` once per node.
+     - `getCSSAsync()` once per node (skipped nodes are excluded).
      - `getStyledTextSegments()` for text nodes.
      - Preprocess styles (clean background, expand shorthands, merge inferred layout, infer resizing, apply overflow).
      - Apply positioning (auto layout absolute, constraints).
      - Replace image fills with uploaded assets.
 
-3. **Normalize variables**
+4. **Normalize variables**
    - Normalize variable names and codeSyntax to a canonical form.
    - Capture variable usage candidates for token detection.
 
-4. **Asset planning and export**
-   - Plan vector/image export at the tree level.
-   - Export SVGs/images only once.
-
 5. **Sanitize styles**
    - Patch known layout issues (negative gap).
    - Ensure layout parents are `position: relative` when children are absolute.
 
 6. **Build layout-only styles**
    - Extract layout-related CSS into a secondary map for SVG external layout.
 
-7. **Render markup**
+7. **Export assets**
+   - Export SVGs/images only once per planned asset.
+
+8. **Render markup**
    - Render nodes into JSX/HTML component tree.
    - Inject SVG content or `<img>` when applicable.
    - Apply Tailwind class conversion.
 
-8. **Token pipeline**
+9. **Token pipeline**
    - Detect token references in output code.
    - Apply plugin transforms to token names.
    - Rewrite code with transformed token names.
    - Resolve tokens to values when requested.
 
-9. **Truncate and finalize output**
-   - Enforce payload size limits.
-   - Emit warnings only for truncation or inferred auto layout.
+10. **Truncate and finalize output**
+    - Enforce payload size limits.
+    - Emit warnings only for truncation or inferred auto layout.
 
 ## Tree and layout semantics
 
@@ -106,6 +110,7 @@ This document describes the implementation design for MCP `get_code` in `package
 - Vector containers can be converted to a single SVG if all leaves are vector-like.
 - SVG nodes only receive external layout styles, not visual paint styles.
 - Placeholder SVG is emitted when export fails or is unavailable, using node width/height.
+- Vector-root descendants are not rendered; their CSS is not collected.
 
 ## Token pipeline (detailed)
 
@@ -127,3 +132,5 @@ This document describes the implementation design for MCP `get_code` in `package
 - Single-pass CSS collection per node.
 - Asset export is planned and executed once.
 - Token detection is string-based and bounded by truncation.
+- Skip CSS collection for vector-root descendants and plugin-rendered subtrees.
+- Variable candidate scan is limited to bound variables and paint references (not inferred variables).

docs/extension/requirements.md

@@ -9,6 +9,7 @@ This document records the source requirements and hard constraints for the MCP `
 - Do not use `renderBounds` diffs for positioning.
 - Do not inject positioning containers on GROUP/BOOLEAN nodes.
 - Keep `getCSSAsync()` at most once per node.
+- If a plugin returns component/code for an instance, do not collect or render its descendants.
 
 ## Scope
 
@@ -88,6 +89,11 @@ Figma `relativeTransform` is relative to the container parent, not to a GROUP/BO
 - When vector export fails or assets are unavailable, preserve layout using a placeholder SVG with node size.
 - Omit `assets` when empty.
 
+## Plugin output
+
+- If a plugin returns component/code for an instance, the instance subtree is not collected.
+- Plugin output is preferred over fallback rendering for that instance.
+
 ## Token handling
 
 - Token detection is based on the final emitted markup (after plugin transform and token rewrites).
@@ -105,9 +111,12 @@ Figma `relativeTransform` is relative to the container parent, not to a GROUP/BO
 
 - Only emit `warnings` for truncation and inferred auto layout.
 - Other degradations should be logged to console with `[tempad-dev]` prefix.
+- The tool may log high-level timing info to console for performance diagnostics.
 
 ## Performance
 
 - `getCSSAsync` must be called at most once per node.
 - `getStyledTextSegments` only for text nodes.
 - Avoid repeated vector export calls; plan and export once per tree.
+- Skip style collection for vector-root descendants (they are not rendered).
+- Variable candidate scanning uses bound variables and paint references; inferred variables are not required.

packages/extension/mcp/tools/code/assets/plan.ts

@@ -2,6 +2,7 @@ import type { NodeSnapshot, VisibleTree } from '../model'
 
 export type AssetPlan = {
   vectorRoots: Set<string>
+  skippedIds: Set<string>
 }
 
 export function planAssets(tree: VisibleTree): AssetPlan {
@@ -32,7 +33,7 @@ export function planAssets(tree: VisibleTree): AssetPlan {
     }
   }
 
-  return { vectorRoots }
+  return { vectorRoots, skippedIds: skipped }
 }
 
 function computeVectorInfo(

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

@@ -15,12 +15,14 @@ import { REQUESTED_SEGMENT_FIELDS } from './text/types'
 export async function collectNodeData(
   tree: VisibleTree,
   config: CodegenConfig,
-  assetRegistry: Map<string, AssetDescriptor>
+  assetRegistry: Map<string, AssetDescriptor>,
+  skipIds?: Set<string>
 ): Promise<CollectedData> {
   const styles = new Map<string, Record<string, string>>()
   const textSegments = new Map<string, StyledTextSegment[] | null>()
 
   for (const id of tree.order) {
+    if (skipIds?.has(id)) continue
     const snapshot = tree.nodes.get(id)
     if (!snapshot) continue
     const node = snapshot.node

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

@@ -2,10 +2,13 @@ import type { AssetDescriptor, GetCodeResult } from '@tempad-dev/mcp-shared'
 
 import { MCP_MAX_PAYLOAD_BYTES } from '@tempad-dev/mcp-shared'
 
+import type { CodegenConfig } from '@/utils/codegen'
+
 import { activePlugin } from '@/ui/state'
 import { stringifyComponent } from '@/utils/component'
 import { simplifyColorMixToRgba } from '@/utils/css'
 
+import type { VisibleTree } from './model'
 import type { CodeLanguage, RenderContext } from './render'
 
 import { currentCodegenConfig } from '../config'
@@ -15,6 +18,7 @@ import { planAssets } from './assets/plan'
 import { collectNodeData } from './collect'
 import { buildGetCodeWarnings, truncateCode } from './messages'
 import { renderTree } from './render'
+import { resolvePluginComponent, type PluginComponent } from './render/plugin'
 import { sanitizeStyles } from './sanitize'
 import { buildLayoutStyles } from './styles'
 import {
@@ -71,6 +75,14 @@ export async function handleGetCode(
   preferredLang?: CodeLanguage,
   resolveTokens?: boolean
 ): Promise<GetCodeResult> {
+  const now = () => (typeof performance !== 'undefined' ? performance.now() : Date.now())
+  const startedAt = now()
+  const timings: Array<[string, number]> = []
+  const stamp = (label: string, start: number) => {
+    const elapsed = Math.round((now() - start) * 10) / 10
+    timings.push([label, elapsed])
+  }
+
   if (nodes.length !== 1) {
     throw new Error('Select exactly one node or provide a single root node id.')
   }
@@ -80,7 +92,9 @@ export async function handleGetCode(
     throw new Error('The selected node is not visible.')
   }
 
+  let t = now()
   const tree = buildVisibleTree(nodes)
+  stamp('tree', t)
   const rootId = tree.rootIds[0]
   if (!rootId) {
     throw new Error('No renderable nodes found for the current selection.')
@@ -94,20 +108,52 @@ export async function handleGetCode(
   const config = currentCodegenConfig()
   const pluginCode = activePlugin.value?.code
 
-  const mappings = buildVariableMappings(nodes)
+  t = now()
+  const variableCache = new Map<string, Variable | null>()
+  const mappings = buildVariableMappings(nodes, variableCache)
+  stamp('vars', t)
+
+  t = now()
+  const plan = planAssets(tree)
+  stamp('plan-assets', t)
+
+  let pluginComponents: Map<string, PluginComponent | null> | undefined
+  const pluginSkipped = new Set<string>()
+  if (pluginCode) {
+    pluginComponents = await collectPluginComponents(tree, config, pluginCode, preferredLang)
+    if (pluginComponents.size) {
+      for (const [id, component] of pluginComponents.entries()) {
+        if (!component) continue
+        const snapshot = tree.nodes.get(id)
+        if (!snapshot) continue
+        snapshot.children.forEach((childId) => skipDescendants(childId, tree, pluginSkipped))
+      }
+    }
+  }
 
   const assetRegistry = new Map<string, AssetDescriptor>()
-  const collected = await collectNodeData(tree, config, assetRegistry)
+  const skipIds =
+    plan.skippedIds.size || pluginSkipped.size
+      ? new Set<string>([...plan.skippedIds, ...pluginSkipped])
+      : plan.skippedIds
+  t = now()
+  const collected = await collectNodeData(tree, config, assetRegistry, skipIds)
+  stamp('collect', t)
 
   // Normalize codeSyntax outputs (e.g. "$kui-space-0") before Tailwind conversion.
-  const usedCandidateIds = normalizeStyleVars(collected.styles, mappings)
-
-  const plan = planAssets(tree)
+  t = now()
+  const usedCandidateIds = normalizeStyleVars(collected.styles, mappings, variableCache)
+  stamp('normalize-vars', t)
 
   // Post-process styles (negative gap, auto-relative, etc.) after var normalization.
+  t = now()
   sanitizeStyles(tree, collected.styles, plan.vectorRoots)
   const layoutStyles = buildLayoutStyles(collected.styles, plan.vectorRoots)
+  stamp('layout', t)
+
+  t = now()
   const svgs = await exportVectorAssets(tree, plan, config, assetRegistry)
+  stamp('export-assets', t)
 
   const nodeMap = new Map<string, SceneNode>()
   collected.nodes.forEach((snap, id) => nodeMap.set(id, snap.node))
@@ -118,12 +164,15 @@ export async function handleGetCode(
     nodes: nodeMap,
     svgs,
     textSegments: collected.textSegments,
+    pluginComponents,
     pluginCode,
     config,
     preferredLang
   }
 
+  t = now()
   const componentTree = await renderTree(rootId, tree, ctx)
+  stamp('render', t)
 
   if (!componentTree) {
     throw new Error('Unable to build markup for the current selection.')
@@ -132,52 +181,70 @@ export async function handleGetCode(
   const resolvedLang = preferredLang ?? ctx.detectedLang ?? 'jsx'
   const rootTag = collected.nodes.get(rootId)?.tag
 
+  t = now()
   const rawMarkup =
     typeof componentTree === 'string'
       ? normalizeRootString(componentTree, rootTag, rootId, resolvedLang)
       : stringifyComponent(componentTree, {
           lang: resolvedLang,
           isInline: (tag) => COMPACT_TAGS.has(tag)
         })
+  stamp('stringify', t)
 
+  t = now()
   const MAX_CODE_CHARS = Math.floor(MCP_MAX_PAYLOAD_BYTES * 0.6)
   let { code, truncated } = truncateCode(rawMarkup, MAX_CODE_CHARS)
+  stamp('truncate', t)
 
   // Token pipeline: detect -> transform -> rewrite -> detect
   const candidateIds = usedCandidateIds.size
     ? new Set<string>([...mappings.variableIds, ...usedCandidateIds])
     : mappings.variableIds
-  const variableCache = new Map<string, Variable | null>()
+  t = now()
   const sourceIndex = buildSourceNameIndex(candidateIds, variableCache)
   const sourceNames = new Set(sourceIndex.keys())
   const usedNamesRaw = extractTokenNames(code, sourceNames)
+  stamp('tokens:detect', t)
 
+  t = now()
   const { rewriteMap, finalBridge } = await applyPluginTransformToNames(
     usedNamesRaw,
     sourceIndex,
     pluginCode,
     config
   )
 
+  let hasRenames = false
+  for (const [key, next] of rewriteMap) {
+    if (key !== next) {
+      hasRenames = true
+      break
+    }
+  }
+
   code = rewriteTokenNamesInCode(code, rewriteMap)
   if (code.length > MAX_CODE_CHARS) {
     code = code.slice(0, MAX_CODE_CHARS)
     truncated = true
   }
 
-  const usedNamesFinal = extractTokenNames(code, sourceNames)
-  const finalBridgeFiltered = filterBridge(finalBridge, usedNamesFinal)
+  const usedNamesFinal = hasRenames ? extractTokenNames(code, sourceNames) : usedNamesRaw
+  const finalBridgeFiltered = hasRenames ? filterBridge(finalBridge, usedNamesFinal) : finalBridge
+  stamp('tokens:rewrite', t)
 
+  t = now()
   const { usedTokens, tokensByCanonical, canonicalById } = await buildUsedTokens(
     usedNamesFinal,
     finalBridgeFiltered,
     config,
     pluginCode,
     variableCache
   )
+  stamp('tokens:used', t)
 
   let resolvedTokens: Record<string, string> | undefined
   if (resolveTokens) {
+    t = now()
     const resolvedByFinal = buildResolvedTokenMap(
       finalBridgeFiltered,
       tokensByCanonical,
@@ -190,6 +257,7 @@ export async function handleGetCode(
       truncated = true
     }
     resolvedTokens = mapResolvedTokens(resolvedByFinal)
+    stamp('tokens:resolve', t)
   }
 
   const warnings = buildGetCodeWarnings(code, MAX_CODE_CHARS, truncated)
@@ -209,6 +277,14 @@ export async function handleGetCode(
         }
       : undefined
 
+  const elapsed = Math.round((now() - startedAt) * 10) / 10
+  console.info(`[tempad-dev] get_code total ${elapsed}ms`)
+  if (timings.length) {
+    const detail = timings.map(([label, ms]) => `${label}=${ms}ms`).join(' ')
+    const info = `nodes=${tree.order.length} text=${collected.textSegments.size} vectors=${plan.vectorRoots.size} assets=${assetRegistry.size}`
+    console.info(`[tempad-dev] get_code timings ${detail} (${info})`)
+  }
+
   return {
     lang: resolvedLang,
     code,
@@ -219,6 +295,31 @@ export async function handleGetCode(
   }
 }
 
+async function collectPluginComponents(
+  tree: VisibleTree,
+  config: CodegenConfig,
+  pluginCode: string,
+  preferredLang?: CodeLanguage
+): Promise<Map<string, PluginComponent | null>> {
+  const out = new Map<string, PluginComponent | null>()
+  for (const id of tree.order) {
+    const snapshot = tree.nodes.get(id)
+    if (!snapshot) continue
+    if (snapshot.node.type !== 'INSTANCE') continue
+    const component = await resolvePluginComponent(snapshot.node, config, pluginCode, preferredLang)
+    out.set(id, component)
+  }
+  return out
+}
+
+function skipDescendants(id: string, tree: VisibleTree, skipped: Set<string>): void {
+  const node = tree.nodes.get(id)
+  if (!node) return
+  if (skipped.has(id)) return
+  skipped.add(id)
+  node.children.forEach((childId) => skipDescendants(childId, tree, skipped))
+}
+
 function normalizeRootString(
   content: string,
   fallbackTag: string | undefined,