From 56e369a1c0b0bb1d40e819fcdca01a147a64b004 Mon Sep 17 00:00:00 2001
From: Dominik Jain <dominik.jain@outlook.com>
Date: Mon, 9 Feb 2026 15:36:39 +0100
Subject: [PATCH] :sparkles: Add helper functions for token exploration

Extend PenpotUtils with helper functions for token exploration/discovery
and describe them in the system prompt
---
 mcp/packages/plugin/src/PenpotUtils.ts | 86 ++++++++++++++++++++++++++
 mcp/packages/server/data/prompts.yml   | 13 ++--
 2 files changed, 92 insertions(+), 7 deletions(-)

diff --git a/mcp/packages/plugin/src/PenpotUtils.ts b/mcp/packages/plugin/src/PenpotUtils.ts
index d5044e658a..61477a185d 100644
--- a/mcp/packages/plugin/src/PenpotUtils.ts
+++ b/mcp/packages/plugin/src/PenpotUtils.ts
@@ -422,4 +422,90 @@ export class PenpotUtils {
                 throw new Error(`Unsupported export mode: ${mode}`);
         }
     }
+
+    /**
+     * Finds all tokens that match the given name across all token sets.
+     *
+     * @param name - The name of the token to search for (case-sensitive exact match)
+     * @returns An array of all matching tokens (may be empty)
+     */
+    public static findTokensByName(name: string): any[] {
+        const tokens: any[] = [];
+        const tokenCatalog = penpot.library.local.tokens;
+
+        for (const set of tokenCatalog.sets) {
+            for (const token of set.tokens) {
+                if (token.name === name) {
+                    tokens.push(token);
+                }
+            }
+        }
+
+        return tokens;
+    }
+
+    /**
+     * Finds the first token that matches the given name across all token sets.
+     *
+     * @param name - The name of the token to search for (case-sensitive exact match)
+     * @returns The first matching token, or null if not found
+     */
+    public static findTokenByName(name: string): any | null {
+        const tokenCatalog = penpot.library.local.tokens;
+
+        for (const set of tokenCatalog.sets) {
+            for (const token of set.tokens) {
+                if (token.name === name) {
+                    return token;
+                }
+            }
+        }
+
+        return null;
+    }
+
+    /**
+     * Gets the token set that contains the given token.
+     *
+     * @param token - The token whose set to find
+     * @returns The TokenSet containing this token, or null if not found
+     */
+    public static getTokenSet(token: any): any | null {
+        const tokenCatalog = penpot.library.local.tokens;
+
+        for (const set of tokenCatalog.sets) {
+            if (set.tokens.includes(token)) {
+                return set;
+            }
+        }
+
+        return null;
+    }
+
+    /**
+     * Generates an overview of all tokens organized by token set name, token type, and token name.
+     * The result is a nested object structure: {tokenSetName: {tokenType: [tokenName, ...]}}.
+     *
+     * @returns An object mapping token set names to objects that map token types to arrays of token names
+     */
+    public static tokenOverview(): Record<string, Record<string, string[]>> {
+        const overview: Record<string, Record<string, string[]>> = {};
+        const tokenCatalog = penpot.library.local.tokens;
+
+        for (const set of tokenCatalog.sets) {
+            const setOverview: Record<string, string[]> = {};
+
+            for (const token of set.tokens) {
+                const tokenType = token.type;
+                if (!setOverview[tokenType]) {
+                    setOverview[tokenType] = [];
+                }
+                setOverview[tokenType].push(token.name);
+            }
+
+            overview[set.name] = setOverview;
+        }
+
+        return overview;
+    }
 }
diff --git a/mcp/packages/server/data/prompts.yml b/mcp/packages/server/data/prompts.yml
index c79094b473..450848cc74 100644
--- a/mcp/packages/server/data/prompts.yml
+++ b/mcp/packages/server/data/prompts.yml
@@ -177,6 +177,7 @@ initial_instructions: |
       General-purpose utility for analyzing/validating descendants
       Calls evaluator on each descendant; collects non-null/undefined results
       Powerful pattern: evaluator can return corrector functions or diagnostic data
+    * Further functions for specific tasks (described in the sections below)
 
   General pointers for working with Penpot designs:
     * Prefer `penpotUtils` helper functions — avoid reimplementing shape searching.
@@ -291,19 +292,17 @@ initial_instructions: |
     * `type: TokenType`
 
   Discovering tokens:
-    ```javascript
-    for (const set of penpot.library.local.tokens.sets) {
-      const token = set.tokens.find(t => t.name === 'color.base.white');
-      if (token) { /* found it */ }
-    }
-    ```
+    * `penpotUtils.tokenOverview()`: Maps from token set name to a mapping from token type to list of token names
+    * `penpotUtils.findTokenByName(name: string): Token | null`: Finds the first applicable token matching the given name
+    * `penpotUtils.findTokensByName(name: string): Token[]`: Finds all tokens that match the given name across all token sets
+    * `penpotUtils.getTokenSet(token: Token): TokenSet | null`: Gets the token set that contains the given token
 
   Applying tokens:
     * `shape.applyToken(token)` - Apply to shape
     * `token.applyToShapes(shapes)` or `token.applyToSelected()` - Apply from token
     * Application is **asynchronous** (wait for ~100ms to see the effects)
     * After application:
-       - `shape.tokens` returns `{ propertyName: "token.name" }`
+       - `shape.tokens` returns a mapping `{ propertyName: "token.name" }`
        - The properties that the tokens control will reflect the token's resolved value.
 
   Removing tokens: