Fix Figma host message handling

.github/workflows/ci.yml

@@ -40,6 +40,9 @@ jobs:
       - name: Build Figma plugin
         run: bun run build:figma
 
+      - name: Test Figma plugin messaging
+        run: bun run --filter './packages/figma' test
+
       - name: Build Gutenberg integration
         run: bun run build:gutenberg
 

packages/figma/package.json

@@ -14,6 +14,7 @@
 		"main": "esbuild main/code.ts --bundle --outfile=dist/code.js",
 		"editor": "vite build --config vite.editor.config.ts",
 		"editor:watch": "vite build --config vite.editor.config.ts --watch",
+		"test": "bun test ./tests",
 		"ui": "tsc && vite build --minify esbuild",
 		"build": "bun run editor && concurrently -n main,ui \"bun run main\" \"bun run ui\"",
 		"dev": "bun run editor && concurrently -n editor,main,ui \"bun run editor:watch\" \"bun run main -- --watch\" \"bun run ui -- --watch\""

packages/figma/src/app.tsx

@@ -12,7 +12,12 @@ import "./style.scss";
 import { ColorVariable, Declaration, Preset, SimpleVariable } from "./types";
 import { generateFluidTypographyObjects } from "./typography/getFluidTypeVariables";
 import { footerLinks } from "./utils/footer";
-import { isMessageFromEditor, postMessageToIframe, postMessageToParent } from "./utils/frameMessaging";
+import {
+	getPluginMessage,
+	isMessageFromEditor,
+	postMessageToIframe,
+	postMessageToParent,
+} from "./utils/frameMessaging";
 
 // Sync variables to Figma - must be outside component to avoid stale closures
 function syncVariables(presetData: Preset, colorVariables: ColorVariable[]) {
@@ -116,7 +121,6 @@ function App() {
 			// Raw messages belong to the bundled editor iframe. Figma host messages
 			// arrive wrapped in event.data.pluginMessage.
 			if (event.data?.type && !isMessageFromEditor(event)) return;
-			if (event.data?.pluginMessage && event.source !== parent) return;
 
 			if (event.data.type === "figma-reopen") {
 				setPreset(null);
@@ -145,7 +149,12 @@ function App() {
 
 			// Handle messages from Figma main code (code.ts)
 			// This handler must be in app.tsx because SelectProjectSection unmounts after project loads
-			const pluginMessage = event.data?.pluginMessage;
+			const pluginMessage = getPluginMessage<{
+				type?: string;
+				error?: string;
+				preset?: Preset;
+				projectId?: string;
+			}>(event);
 			if (pluginMessage?.type === "wordpress-response") {
 				postMessageToIframe("cf-figma-wordpress-response", pluginMessage);
 			}

packages/figma/src/components/SelectProjectSection.tsx

@@ -7,7 +7,12 @@ import { generateSpacingObjects } from "../spacing/getFluidSpacingVariables";
 import { ColorVariable, Declaration, Preset, SimpleVariable } from "../types";
 import { generateFluidTypographyObjects } from "../typography/getFluidTypeVariables";
 import { devLog } from "../utils";
-import { isMessageFromEditor, postMessageToIframe, postMessageToParent } from "../utils/frameMessaging";
+import {
+	getPluginMessage,
+	isMessageFromEditor,
+	postMessageToIframe,
+	postMessageToParent,
+} from "../utils/frameMessaging";
 import { Card } from "./Card";
 
 interface SelectProjectSection {
@@ -174,9 +179,14 @@ export const SelectProjectSection = memo<SelectProjectSection>(({ handleLoadedPr
 				};
 			}>,
 		) => {
-			if (event.data?.pluginMessage && event.source !== parent) return;
+			const pluginMessage = getPluginMessage<{
+				type?: string;
+				error?: string;
+				projectId?: string;
+				preset?: Preset;
+			}>(event);
+			if (!pluginMessage) return;
 
-			const pluginMessage = event?.data?.pluginMessage;
 			devLog("pluginMessage", pluginMessage);
 			switch (pluginMessage?.type) {
 				case "import-project": {
@@ -193,7 +203,7 @@ export const SelectProjectSection = memo<SelectProjectSection>(({ handleLoadedPr
 					break;
 				}
 				case "get-project-id": {
-					const receivedApiKey = event.data.pluginMessage?.projectId;
+					const receivedApiKey = pluginMessage.projectId;
 
 					if (receivedApiKey) {
 						postMessageToIframe("cf-figma-set-api-key", { apiKey: receivedApiKey });
@@ -203,7 +213,7 @@ export const SelectProjectSection = memo<SelectProjectSection>(({ handleLoadedPr
 					break;
 				}
 				case "get-project-locally": {
-					const preset = event.data.pluginMessage?.preset;
+					const preset = pluginMessage.preset;
 					if (preset) {
 						setLocalPreset(preset);
 					}

packages/figma/src/utils/frameMessaging.ts

@@ -8,6 +8,12 @@ export const isMessageFromEditor = (event: MessageEvent) => {
 	return Boolean(iframe?.contentWindow && event.source === iframe.contentWindow);
 };
 
+export const getPluginMessage = <T extends Record<string, unknown>>(event: MessageEvent): T | null => {
+	const pluginMessage = event.data?.pluginMessage;
+
+	return pluginMessage && typeof pluginMessage === "object" ? (pluginMessage as T) : null;
+};
+
 export const postMessageToParent = (message: Record<string, unknown>) => {
 	parent.postMessage({ pluginMessage: message }, "*");
 };

packages/figma/tests/frameMessaging.test.ts

@@ -0,0 +1,29 @@
+import { describe, expect, test } from "bun:test";
+import { getPluginMessage } from "../src/utils/frameMessaging";
+
+describe("getPluginMessage", () => {
+	test("accepts Figma host messages when event.source is null", () => {
+		const pluginMessage = { type: "import-project", projectId: "connection-key" };
+		const event = {
+			data: { pluginMessage },
+			source: null,
+		} as unknown as MessageEvent;
+
+		expect(getPluginMessage(event)).toEqual(pluginMessage);
+	});
+
+	test("accepts Figma host messages without assuming the parent source", () => {
+		const pluginMessage = { type: "import-project-error", error: "Failed to fetch preset" };
+		const event = {
+			data: { pluginMessage },
+			source: {} as MessageEventSource,
+		} as MessageEvent;
+
+		expect(getPluginMessage(event)).toEqual(pluginMessage);
+	});
+
+	test("rejects raw editor and malformed messages", () => {
+		expect(getPluginMessage({ data: { type: "cf-figma-ready" } } as MessageEvent)).toBeNull();
+		expect(getPluginMessage({ data: { pluginMessage: "invalid" } } as MessageEvent)).toBeNull();
+	});
+});