feat: save variable and style references for Studio projects

.changeset/theme-refs-round-trip.md

@@ -0,0 +1,5 @@
+---
+"@tokens-studio/figma-plugin": patch
+---
+
+Added support for saving variable and style references for Studio projects

packages/tokens-studio-for-figma/src/app/store/middlewares/tokenState.ts

@@ -1,5 +1,6 @@
 import { StorageProviderType } from '@/constants/StorageProviderType';
 import { updateThemeGroupsInTokensStudio } from '@/storage/tokensStudio/updateThemeGroupsInTokensStudio';
+import { updateThemeRefsViaRestApi, snapshotThemeRefs } from '@/utils/tokensStudio/updateThemeRefsViaRestApi';
 
 const actionsToTriggerUpdateInTokensStudio = [
   'tokenState/assignVariableIdsToCurrentTheme',
@@ -14,8 +15,32 @@ const actionsToTriggerUpdateInTokensStudio = [
   'tokenState/disconnectStyleFromTheme',
 ];
 
+const refOnlyActions = [
+  'tokenState/assignVariableIdsToCurrentTheme',
+  'tokenState/assignVariableIdsToTheme',
+  'tokenState/assignStyleIdsToCurrentTheme',
+  'tokenState/assignStyleIdsToTheme',
+  'tokenState/disconnectVariableFromTheme',
+  'tokenState/disconnectStyleFromTheme',
+  'tokenState/renameVariableIdsToTheme',
+  'tokenState/renameStyleIdsToCurrentTheme',
+  'tokenState/renameVariableNamesToThemes',
+  'tokenState/renameStyleNamesToCurrentTheme',
+  'tokenState/removeVariableNamesFromThemes',
+  'tokenState/removeStyleNamesFromThemes',
+  'tokenState/removeStyleIdsFromThemes',
+];
+
 export const tokenStateMiddleware = (store) => (next) => (action) => {
   const prevState = store.getState();
+
+  // Capture refs snapshot BEFORE the action runs
+  const isOAuthRefAction = prevState.uiState.api?.provider === StorageProviderType.TOKENS_STUDIO_OAUTH
+    && refOnlyActions.includes(action.type);
+  const prevThemeRefs = isOAuthRefAction
+    ? snapshotThemeRefs(prevState.tokenState.themes)
+    : undefined;
+
   next(action);
   const nextState = store.getState();
 
@@ -30,4 +55,11 @@ export const tokenStateMiddleware = (store) => (next) => (action) => {
       dispatch: store.dispatch,
     });
   }
+
+  if (prevThemeRefs) {
+    updateThemeRefsViaRestApi({
+      prevThemeRefs,
+      rootState: nextState,
+    }).catch((err) => console.error('[tokenStateMiddleware] REST ref sync failed:', err));
+  }
 };

packages/tokens-studio-for-figma/src/storage/schemas/themeObjectSchema.ts

@@ -14,4 +14,6 @@ export const themeObjectSchema = z.object({
   $figmaVariableReferences: z.record(z.string()).optional(),
   $figmaCollectionId: z.string().optional(),
   $figmaModeId: z.string().optional(),
+  $themeGroupId: z.string().optional(),
+  $themeOptionId: z.string().optional(),
 });

packages/tokens-studio-for-figma/src/types/ThemeObject.ts

@@ -15,4 +15,7 @@ export type ThemeObject = {
   $figmaCollectionId?: string;
   $figmaModeId?: string;
   $figmaVariableReferences?: Record<string, string>;
+  // Server IDs for REST API round-trip (Studio-on-Rails)
+  $themeGroupId?: string;
+  $themeOptionId?: string;
 };

packages/tokens-studio-for-figma/src/utils/tokensStudio/fetchBranchesListRest.ts

@@ -25,6 +25,33 @@ export function parseBranchesFromResponse(branchesData: any): RestBranch[] {
   return branches;
 }
 
+export async function resolveChangeSetId(
+  authToken: string,
+  apiBaseUrl: string,
+  projectId: string,
+  branchName: string,
+): Promise<string | null> {
+  const headers = {
+    Authorization: `Bearer ${authToken}`,
+    'Content-Type': 'application/json',
+  };
+
+  try {
+    const branchesRes = await fetch(`${apiBaseUrl}/api/v1/projects/${projectId}/branches`, { headers });
+    if (!branchesRes.ok) throw new Error(`Failed to fetch branches: ${branchesRes.statusText}`);
+    const branchesData = await branchesRes.json();
+
+    const branches = parseBranchesFromResponse(branchesData);
+    const branch = branches.find((b) => b.name === branchName)
+      || branches.find((b) => b.is_default);
+
+    return branch?.change_set_id || null;
+  } catch (error) {
+    console.error('Error resolving change_set_id:', error);
+    return null;
+  }
+}
+
 export async function fetchBranchesListRest(
   authToken: string,
   apiBaseUrl: string,

packages/tokens-studio-for-figma/src/utils/tokensStudio/fetchProjectDataRest.ts

@@ -20,7 +20,6 @@ interface RestThemeOption {
   theme_group_id: string;
   selected_token_sets: Record<string, string>;
   figmaStyleReferences?: Record<string, string>;
-  figmaVariableReferences?: Record<string, string>;
   figmaCollectionId?: string;
   figmaModeId?: string;
 }
@@ -219,7 +218,6 @@ export async function fetchProjectDataRest(
           theme_group_id: item.relationships?.theme_group?.data?.id || '',
           selected_token_sets: item.attributes?.selected_token_sets || {},
           figmaStyleReferences: item.attributes?.figma_style_references,
-          figmaVariableReferences: item.attributes?.figma_variable_references,
           figmaCollectionId: item.attributes?.figma_collection_id,
           figmaModeId: item.attributes?.figma_mode_id,
         });
@@ -232,12 +230,13 @@ export async function fetchProjectDataRest(
       optionsByGroupId[opt.theme_group_id].push(opt);
     });
 
-    // Parse theme groups
+    // Parse theme groups — variable refs live at the group level
     const themes: ThemeObjectsList = [];
     if (themeGroupsData.data && Array.isArray(themeGroupsData.data)) {
       themeGroupsData.data.forEach((group: any) => {
         const { id } = group;
         const name = group.attributes?.name || id;
+        const groupVariableRefs: Record<string, string> | undefined = group.attributes?.figma_variable_references;
         const options = optionsByGroupId[id] || [];
 
         options.forEach((opt) => {
@@ -255,10 +254,12 @@ export async function fetchProjectDataRest(
             name: opt.name,
             group: name,
             selectedTokenSets: selectedTokenSetsByName as any,
+            $themeGroupId: id,
+            $themeOptionId: opt.id,
           };
 
           if (opt.figmaStyleReferences !== undefined) themeObj.$figmaStyleReferences = opt.figmaStyleReferences;
-          if (opt.figmaVariableReferences !== undefined) themeObj.$figmaVariableReferences = opt.figmaVariableReferences;
+          if (groupVariableRefs !== undefined) themeObj.$figmaVariableReferences = { ...groupVariableRefs };
           if (opt.figmaCollectionId !== undefined) themeObj.$figmaCollectionId = opt.figmaCollectionId;
           if (opt.figmaModeId !== undefined) themeObj.$figmaModeId = opt.figmaModeId;
 

packages/tokens-studio-for-figma/src/utils/tokensStudio/pushThemeRefsRest.test.ts

@@ -0,0 +1,102 @@
+import { patchThemeGroupVariableRefs, patchThemeOptionStyleRefs } from './pushThemeRefsRest';
+
+describe('pushThemeRefsRest', () => {
+  const mockFetch = jest.fn();
+  const originalFetch = global.fetch;
+
+  beforeEach(() => {
+    global.fetch = mockFetch;
+    mockFetch.mockReset();
+  });
+
+  afterAll(() => {
+    global.fetch = originalFetch;
+  });
+
+  const authToken = 'test-token';
+  const apiBaseUrl = 'https://studio.example.com';
+  const projectId = 'proj-123';
+  const changeSetId = 'cs-456';
+
+  describe('patchThemeGroupVariableRefs', () => {
+    it('should skip PATCH when refs is undefined (preserve semantics)', async () => {
+      await patchThemeGroupVariableRefs(authToken, apiBaseUrl, projectId, changeSetId, 'group-1', undefined);
+      expect(mockFetch).not.toHaveBeenCalled();
+    });
+
+    it('should PATCH with empty object to clear refs', async () => {
+      mockFetch.mockResolvedValue({ ok: true });
+
+      await patchThemeGroupVariableRefs(authToken, apiBaseUrl, projectId, changeSetId, 'group-1', {});
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        `${apiBaseUrl}/api/v1/projects/${projectId}/theme_groups/group-1?change_set_id=${encodeURIComponent(changeSetId)}`,
+        expect.objectContaining({
+          method: 'PATCH',
+          headers: expect.objectContaining({
+            Authorization: `Bearer ${authToken}`,
+            'Content-Type': 'application/json',
+          }),
+          body: JSON.stringify({ theme_group: { figma_variable_references: {} } }),
+        }),
+      );
+    });
+
+    it('should PATCH with refs map', async () => {
+      mockFetch.mockResolvedValue({ ok: true });
+      const refs = { 'color.brand.500': 'VariableID:123:456' };
+
+      await patchThemeGroupVariableRefs(authToken, apiBaseUrl, projectId, changeSetId, 'group-1', refs);
+
+      const body = JSON.parse(mockFetch.mock.calls[0][1].body);
+      expect(body.theme_group.figma_variable_references).toEqual(refs);
+    });
+
+    it('should throw on non-OK response', async () => {
+      mockFetch.mockResolvedValue({ ok: false, status: 422, text: () => Promise.resolve('Unprocessable') });
+
+      await expect(
+        patchThemeGroupVariableRefs(authToken, apiBaseUrl, projectId, changeSetId, 'group-1', {}),
+      ).rejects.toThrow('Failed to PATCH theme_group variable refs (422)');
+    });
+  });
+
+  describe('patchThemeOptionStyleRefs', () => {
+    it('should skip PATCH when refs is undefined (preserve semantics)', async () => {
+      await patchThemeOptionStyleRefs(authToken, apiBaseUrl, projectId, changeSetId, 'opt-1', undefined);
+      expect(mockFetch).not.toHaveBeenCalled();
+    });
+
+    it('should PATCH with empty object to clear refs', async () => {
+      mockFetch.mockResolvedValue({ ok: true });
+
+      await patchThemeOptionStyleRefs(authToken, apiBaseUrl, projectId, changeSetId, 'opt-1', {});
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        `${apiBaseUrl}/api/v1/projects/${projectId}/theme_options/opt-1?change_set_id=${encodeURIComponent(changeSetId)}`,
+        expect.objectContaining({
+          method: 'PATCH',
+          body: JSON.stringify({ theme_option: { figma_style_references: {} } }),
+        }),
+      );
+    });
+
+    it('should PATCH with style refs verbatim (S: prefix included)', async () => {
+      mockFetch.mockResolvedValue({ ok: true });
+      const refs = { 'hint.text': 'S:81e4c301abc,' };
+
+      await patchThemeOptionStyleRefs(authToken, apiBaseUrl, projectId, changeSetId, 'opt-1', refs);
+
+      const body = JSON.parse(mockFetch.mock.calls[0][1].body);
+      expect(body.theme_option.figma_style_references).toEqual(refs);
+    });
+
+    it('should throw on non-OK response', async () => {
+      mockFetch.mockResolvedValue({ ok: false, status: 500, text: () => Promise.resolve('Internal') });
+
+      await expect(
+        patchThemeOptionStyleRefs(authToken, apiBaseUrl, projectId, changeSetId, 'opt-1', {}),
+      ).rejects.toThrow('Failed to PATCH theme_option style refs (500)');
+    });
+  });
+});

packages/tokens-studio-for-figma/src/utils/tokensStudio/pushThemeRefsRest.ts

@@ -0,0 +1,73 @@
+/**
+ * REST PATCH helpers for syncing $figmaVariableReferences and $figmaStyleReferences
+ * with Studio-on-Rails theme_groups and theme_options endpoints.
+ *
+ * Semantics:
+ * - undefined → omit key from payload (preserve existing refs on server)
+ * - {}       → explicitly clear all refs
+ * - never send null
+ */
+
+export async function patchThemeGroupVariableRefs(
+  authToken: string,
+  apiBaseUrl: string,
+  projectId: string,
+  changeSetId: string,
+  themeGroupId: string,
+  refs: Record<string, string> | undefined,
+): Promise<void> {
+  if (refs === undefined) return;
+
+  const url = `${apiBaseUrl}/api/v1/projects/${projectId}/theme_groups/${themeGroupId}?change_set_id=${encodeURIComponent(changeSetId)}`;
+  const body = {
+    theme_group: {
+      figma_variable_references: refs,
+    },
+  };
+
+  const res = await fetch(url, {
+    method: 'PATCH',
+    headers: {
+      Authorization: `Bearer ${authToken}`,
+      'Content-Type': 'application/json',
+    },
+    body: JSON.stringify(body),
+  });
+
+  if (!res.ok) {
+    const errorText = await res.text().catch(() => '');
+    throw new Error(`Failed to PATCH theme_group variable refs (${res.status}): ${errorText}`);
+  }
+}
+
+export async function patchThemeOptionStyleRefs(
+  authToken: string,
+  apiBaseUrl: string,
+  projectId: string,
+  changeSetId: string,
+  themeOptionId: string,
+  refs: Record<string, string> | undefined,
+): Promise<void> {
+  if (refs === undefined) return;
+
+  const url = `${apiBaseUrl}/api/v1/projects/${projectId}/theme_options/${themeOptionId}?change_set_id=${encodeURIComponent(changeSetId)}`;
+  const body = {
+    theme_option: {
+      figma_style_references: refs,
+    },
+  };
+
+  const res = await fetch(url, {
+    method: 'PATCH',
+    headers: {
+      Authorization: `Bearer ${authToken}`,
+      'Content-Type': 'application/json',
+    },
+    body: JSON.stringify(body),
+  });
+
+  if (!res.ok) {
+    const errorText = await res.text().catch(() => '');
+    throw new Error(`Failed to PATCH theme_option style refs (${res.status}): ${errorText}`);
+  }
+}