feat: Add property for preserving css variables in token resolution

Author: jkuelz

src/build/__tests__/css-vars.test.ts

@@ -0,0 +1,137 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+import { test, expect } from 'vitest';
+import { buildThemedComponentsInternal } from '../internal';
+import { Theme, resolveTheme, resolveContext } from '../../shared/theme';
+import { isReferenceToken, generateReferenceTokenDefaults, generateReferenceTokenName } from '../../shared/theme/utils';
+import { createBuildDeclarations } from '../../shared/declaration';
+import { mkdtemp, rm } from 'fs/promises';
+import { tmpdir } from 'os';
+import { join } from 'path';
+
+const testTheme: Theme = {
+  id: 'test',
+  selector: ':root',
+  tokens: {
+    colorPrimary: '#0073bb',
+    colorBackground: '{colorPrimary}',
+    colorPrimary500: '#direct-primary-500', // Different from reference token to test precedence
+    colorSecondary: '{colorNeutral900}', // References a token that exists in both places
+    colorNeutral900: '#direct-neutral', // Different from reference token
+  },
+  modes: {},
+  contexts: {},
+  tokenModeMap: {},
+  referenceTokens: {
+    color: {
+      primary: { 500: '#reference-primary-value' }, // Also defined in tokens - tests precedence
+      neutral: { 900: '#reference-neutral-value' }, // Also defined in tokens - tests precedence
+      success: { 200: '#reference-success-value' }, // Only in reference tokens - tests auto-generation
+    },
+  },
+};
+
+const propertiesMap = {
+  colorPrimary: '--color-primary',
+  colorBackground: '--color-background',
+  colorPrimary500: '--color-primary-500',
+  colorSecondary: '--color-secondary',
+  colorNeutral900: '--color-neutral-900',
+  colorSuccess200: '--color-success-200', // For the reference-only token
+};
+
+test('CSS variable optimization works without errors', async () => {
+  const tempDir = await mkdtemp(join(tmpdir(), 'css-vars-test-'));
+
+  try {
+    await buildThemedComponentsInternal({
+      primary: testTheme,
+      exposed: ['colorPrimary', 'colorBackground'],
+      themeable: ['colorPrimary', 'colorBackground'],
+      variablesMap: {
+        colorPrimary: 'color-primary',
+        colorBackground: 'color-background',
+      },
+      componentsOutputDir: tempDir,
+      scssDir: tempDir,
+      useCssVars: true,
+      skip: ['preset', 'design-tokens'],
+    });
+
+    expect(true).toBe(true);
+  } finally {
+    await rm(tempDir, { recursive: true, force: true });
+  }
+});
+
+test('isReferenceToken identifies reference tokens correctly', () => {
+  expect(isReferenceToken(testTheme, 'colorPrimary500')).toBe(true);
+  expect(isReferenceToken(testTheme, 'colorNeutral900')).toBe(true);
+  expect(isReferenceToken(testTheme, 'colorPrimary')).toBe(false);
+  expect(isReferenceToken(testTheme, 'colorSuccess500')).toBe(false); // Not in referenceTokens
+});
+
+test('generateReferenceTokenName creates correct token names', () => {
+  expect(generateReferenceTokenName('primary', '500')).toBe('colorPrimary500');
+  expect(generateReferenceTokenName('neutral', '900')).toBe('colorNeutral900');
+});
+
+test('generateReferenceTokenDefaults creates CSS variable declarations', () => {
+  const defaults = generateReferenceTokenDefaults(testTheme, propertiesMap);
+
+  expect(defaults).toEqual({
+    '--color-primary-500': '#reference-primary-value',
+    '--color-neutral-900': '#reference-neutral-value',
+    '--color-success-200': '#reference-success-value',
+  });
+});
+
+test('resolveTheme with CSS variables returns CSS var() for reference tokens', () => {
+  const resolved = resolveTheme(testTheme, undefined, {
+    useCssVars: true,
+    propertiesMap,
+  });
+
+  // Direct token takes precedence over reference token
+  expect(resolved.colorPrimary500).toBe('#direct-primary-500');
+  // Token referencing a reference token should get CSS var
+  expect(resolved.colorSecondary).toBe('var(--color-neutral-900)');
+  expect(resolved.colorPrimary).toBe('#0073bb'); // Not a reference token
+});
+
+test('token precedence: direct tokens override reference tokens', () => {
+  const resolved = resolveTheme(testTheme, undefined);
+
+  // Direct token value should be used, not reference token value
+  expect(resolved.colorPrimary500).toBe('#direct-primary-500');
+  expect(resolved.colorNeutral900).toBe('#direct-neutral');
+});
+
+test('resolveContext with CSS variables preserves var() references', () => {
+  const context = {
+    id: 'test-context',
+    selector: '.test',
+    tokens: {
+      colorSecondary: '{colorNeutral900}', // References a token that should become CSS var
+    },
+  };
+
+  const resolved = resolveContext(testTheme, context, undefined, undefined, {
+    useCssVars: true,
+    propertiesMap,
+  });
+
+  // Should resolve to CSS var since colorNeutral900 is a reference token
+  expect(resolved.colorSecondary).toBe('var(--color-neutral-900)');
+});
+
+test('createBuildDeclarations includes reference tokens when useCssVars enabled', () => {
+  const css = createBuildDeclarations(testTheme, [], propertiesMap, (selector) => selector, [], {
+    useCssVars: true,
+    propertiesMap,
+  });
+
+  expect(css).toContain('--color-primary-500');
+  expect(css).toContain('--color-neutral-900');
+  expect(css).not.toContain('--color-success-200'); // Not auto-generated because it's not in used tokens
+});

src/build/inline-stylesheets.ts

@@ -18,14 +18,16 @@ export function getInlineStylesheets(
   resolution: SpecificResolution,
   variablesMap: Record<string, string>,
   propertiesMap: Record<string, string>,
-  neededTokens: string[]
+  neededTokens: string[],
+  useCssVars?: boolean
 ): InlineStylesheet[] {
   const declarations = createBuildDeclarations(
     primary,
     secondary,
     propertiesMap,
     (selector) => markGlobal(selector),
-    neededTokens
+    neededTokens,
+    useCssVars ? { useCssVars, propertiesMap } : undefined
   );
 
   const declaration = {

src/build/internal.ts

@@ -40,6 +40,8 @@ export interface BuildThemedComponentsInternalParams {
   jsonSchema?: boolean;
   /** Fail the build when SASS deprecation warning occurs **/
   failOnDeprecations?: boolean;
+  /** Use CSS variables instead of resolved values for better runtime theme swapping */
+  useCssVars?: boolean;
 }
 /**
  * Builds themed components and optionally design tokens, if not skipped.
@@ -71,6 +73,7 @@ export async function buildThemedComponentsInternal(params: BuildThemedComponent
     descriptions = {},
     jsonSchema = false,
     failOnDeprecations,
+    useCssVars = false,
   } = params;
 
   if (!skip.includes('design-tokens') && !designTokensOutputDir) {
@@ -86,7 +89,7 @@ export async function buildThemedComponentsInternal(params: BuildThemedComponent
   const styleTask = buildStyles(
     scssDir,
     componentsOutputDir,
-    getInlineStylesheets(primary, secondary, defaults, variablesMap, propertiesMap, neededTokens),
+    getInlineStylesheets(primary, secondary, defaults, variablesMap, propertiesMap, neededTokens, useCssVars),
     { failOnDeprecations }
   );
   const internalTokensTask = createInternalTokenFiles(defaults, propertiesMap, componentsOutputDir);

src/shared/declaration/index.ts

@@ -1,6 +1,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-import { mergeInPlace, Override, Theme } from '../theme';
+import { mergeInPlace, Override, Theme, ResolveOptions } from '../theme';
+import { generateReferenceTokenName } from '../theme/utils';
 import type { PropertiesMap, SelectorCustomizer } from './interfaces';
 import { RuleCreator } from './rule';
 import { SingleThemeCreator } from './single';
@@ -35,12 +36,13 @@ export function createOverrideDeclarations(
   base: Theme,
   override: Override,
   propertiesMap: PropertiesMap,
-  selectorCustomizer: SelectorCustomizer
+  selectorCustomizer: SelectorCustomizer,
+  options?: ResolveOptions
 ): string {
   // create theme containing only modified tokens
   const minimalTheme = createMinimalTheme(base, override);
   const ruleCreator = new RuleCreator(new Selector(selectorCustomizer), new AllPropertyRegistry(propertiesMap));
-  const stylesheetCreator = new SingleThemeCreator(minimalTheme, ruleCreator, base);
+  const stylesheetCreator = new SingleThemeCreator(minimalTheme, ruleCreator, base, options);
   const stylesheet = stylesheetCreator.create();
   return stylesheet.toString();
 }
@@ -50,10 +52,33 @@ export function createBuildDeclarations(
   secondary: Theme[],
   propertiesMap: PropertiesMap,
   selectorCustomizer: SelectorCustomizer,
-  used: string[]
+  used: string[],
+  options?: ResolveOptions
 ): string {
-  const ruleCreator = new RuleCreator(new Selector(selectorCustomizer), new UsedPropertyRegistry(propertiesMap, used));
-  const stylesheetCreator = new MultiThemeCreator([primary, ...secondary], ruleCreator);
+  // When CSS vars are enabled, include reference tokens from all themes in the used list
+  let effectiveUsed = used;
+  if (options?.useCssVars) {
+    const allThemes = [primary, ...secondary];
+    const referenceTokens: string[] = [];
+    allThemes.forEach((theme) => {
+      if (theme.referenceTokens?.color) {
+        Object.entries(theme.referenceTokens.color).forEach(([colorName, palette]) => {
+          if (palette) {
+            Object.keys(palette).forEach((step) => {
+              referenceTokens.push(generateReferenceTokenName(colorName, step));
+            });
+          }
+        });
+      }
+    });
+    effectiveUsed = [...used, ...referenceTokens];
+  }
+
+  const ruleCreator = new RuleCreator(
+    new Selector(selectorCustomizer),
+    new UsedPropertyRegistry(propertiesMap, effectiveUsed)
+  );
+  const stylesheetCreator = new MultiThemeCreator([primary, ...secondary], ruleCreator, options);
   const stylesheet = stylesheetCreator.create();
   const transformer = new MinimalTransformer();
   const minimal = transformer.transform(stylesheet);

src/shared/declaration/multi.ts

@@ -1,7 +1,17 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 import { isGlobalSelector } from '../styles/selector';
-import { defaultsReducer, modeReducer, OptionalState, reduce, resolveContext, resolveTheme, Theme } from '../theme';
+import {
+  defaultsReducer,
+  modeReducer,
+  OptionalState,
+  reduce,
+  resolveContext,
+  resolveTheme,
+  Theme,
+  ResolveOptions,
+} from '../theme';
+import { generateReferenceTokenDefaults } from '../theme/utils';
 import { AbstractCreator } from './abstract';
 import type { StylesheetCreator } from './interfaces';
 import { RuleCreator, SelectorConfig } from './rule';
@@ -16,11 +26,13 @@ import { compact } from './utils';
 export class MultiThemeCreator extends AbstractCreator implements StylesheetCreator {
   themes: Theme[];
   ruleCreator: RuleCreator;
+  options?: ResolveOptions;
 
-  constructor(themes: Theme[], ruleCreator: RuleCreator) {
+  constructor(themes: Theme[], ruleCreator: RuleCreator, options?: ResolveOptions) {
     super();
     this.themes = themes;
     this.ruleCreator = ruleCreator;
+    this.options = options;
   }
 
   create(): Stylesheet {
@@ -38,7 +50,9 @@ export class MultiThemeCreator extends AbstractCreator implements StylesheetCrea
 
     if (!globalThemes.length) {
       // If there is no root theme, all themes are scoped by their root selector. No interference.
-      const stylesheets = this.themes.map((theme) => new SingleThemeCreator(theme, this.ruleCreator).create());
+      const stylesheets = this.themes.map((theme) =>
+        new SingleThemeCreator(theme, this.ruleCreator, undefined, this.options).create()
+      );
       const result = new Stylesheet();
       stylesheets.forEach((stylesheet) => {
         stylesheet.getAllRules().map((rule) => result.appendRuleWithPath(rule, stylesheet.getPath(rule) ?? []));
@@ -48,7 +62,7 @@ export class MultiThemeCreator extends AbstractCreator implements StylesheetCrea
 
     const [globalTheme] = globalThemes;
 
-    const stylesheet = new SingleThemeCreator(globalTheme, this.ruleCreator).create();
+    const stylesheet = new SingleThemeCreator(globalTheme, this.ruleCreator, undefined, this.options).create();
     const secondaries = this.getThemesWithout(globalTheme);
     secondaries.forEach((secondary) => {
       this.appendRulesForSecondary(stylesheet, globalTheme, secondary);
@@ -58,8 +72,15 @@ export class MultiThemeCreator extends AbstractCreator implements StylesheetCrea
   }
 
   appendRulesForSecondary(stylesheet: Stylesheet, primary: Theme, secondary: Theme) {
-    const secondaryResolution = resolveTheme(secondary);
+    const secondaryResolution = resolveTheme(secondary, undefined, this.options);
     const defaults = reduce(secondaryResolution, secondary, defaultsReducer());
+
+    // Add CSS variable declarations for reference tokens when useCssVars is enabled
+    if (this.options?.useCssVars && this.options?.propertiesMap) {
+      const referenceDefaults = generateReferenceTokenDefaults(secondary, this.options.propertiesMap);
+      Object.assign(defaults, referenceDefaults);
+    }
+
     const rootRule = this.ruleCreator.create({ global: [secondary.selector] }, defaults);
     const parentRule = this.findRule(stylesheet, { global: [primary.selector] });
     MultiThemeCreator.appendRuleToStylesheet(stylesheet, rootRule, compact([parentRule]));
@@ -80,7 +101,11 @@ export class MultiThemeCreator extends AbstractCreator implements StylesheetCrea
     });
 
     MultiThemeCreator.forEachContext(secondary, (context) => {
-      const contextResolution = reduce(resolveContext(secondary, context), secondary, defaultsReducer());
+      const contextResolution = reduce(
+        resolveContext(secondary, context, undefined, undefined, this.options),
+        secondary,
+        defaultsReducer()
+      );
       const contextRule = this.ruleCreator.create(
         { global: [secondary.selector], local: [context.selector] },
         contextResolution
@@ -110,7 +135,11 @@ export class MultiThemeCreator extends AbstractCreator implements StylesheetCrea
 
     MultiThemeCreator.forEachContextWithinOptionalModeState(secondary, (context, mode, state) => {
       const optionalState = mode.states[state] as OptionalState;
-      const contextResolution = reduce(resolveContext(secondary, context), secondary, modeReducer(mode, state));
+      const contextResolution = reduce(
+        resolveContext(secondary, context, undefined, undefined, this.options),
+        secondary,
+        modeReducer(mode, state)
+      );
       const contextRule = this.findRule(stylesheet, { global: [secondary.selector], local: [context.selector] });
       const modeRule = this.findRule(stylesheet, {
         global: [secondary.selector, optionalState.selector],