get-figma-data-tool.ts

import { z } from "zod";
import type { GetFileResponse, GetFileNodesResponse } from "@figma/rest-api-spec";
import { FigmaService } from "~/services/figma.js";
import { simplifyRawFigmaObject, allExtractors } from "~/extractors/index.js";
import yaml from "js-yaml";
import { Logger, writeLogs } from "~/utils/logger.js";
import { calcStringSize } from "~/utils/calc-string-size.js";
import type { CallToolResult } from "@modelcontextprotocol/sdk/types.js";

const parameters = {
  fileKey: z
    .string()
    .describe(
      "The key of the Figma file to fetch, often found in a provided URL like figma.com/(file|design)/<fileKey>/...",
    ),
  nodeId: z
    .string()
    .optional()
    .describe(
      "The ID of the node to fetch, often found as URL parameter node-id=<nodeId>, always use if provided",
    ),
  depth: z
    .number()
    .optional()
    .describe(
      "OPTIONAL. Do NOT use unless explicitly requested by the user. Controls how many levels deep to traverse the node tree.",
    ),
};

const parametersSchema = z.object(parameters);
export type GetFigmaDataParams = z.infer<typeof parametersSchema>;

const sizeLimit = process.env.GET_NODE_SIZE_LIMIT
  ? parseInt(process.env.GET_NODE_SIZE_LIMIT)
  : undefined;
const needLimitPrompt = sizeLimit && sizeLimit > 0;

// Simplified handler function
async function getFigmaData(
  params: GetFigmaDataParams,
  figmaService: FigmaService,
  outputFormat: "yaml" | "json",
): Promise<CallToolResult> {
  try {
    const { fileKey, nodeId, depth } = params;

    Logger.log(
      `Fetching ${depth ? `${depth} layers deep` : "all layers"} of ${
        nodeId ? `node ${nodeId} from file` : `full file`
      } ${fileKey}`,
    );

    // Get raw Figma API response
    let rawApiResponse: GetFileResponse | GetFileNodesResponse;
    if (nodeId) {
      rawApiResponse = await figmaService.getRawNode(fileKey, nodeId, depth);
    } else {
      rawApiResponse = await figmaService.getRawFile(fileKey, depth);
    }

    // Use unified design extraction (handles nodes + components consistently)
    const simplifiedDesign = simplifyRawFigmaObject(rawApiResponse, allExtractors, {
      maxDepth: depth,
    });

    writeLogs("figma-simplified.json", simplifiedDesign);

    Logger.log(
      `Successfully extracted data: ${simplifiedDesign.nodes.length} nodes, ${Object.keys(simplifiedDesign.globalVars.styles).length} styles`,
    );

    const { nodes, globalVars, ...metadata } = simplifiedDesign;
    const result = {
      metadata,
      nodes,
      globalVars,
    };

    Logger.log(`Generating ${outputFormat.toUpperCase()} result from extracted data`);
    const formattedResult =
      outputFormat === "json" ? JSON.stringify(result, null, 2) : yaml.dump(result);

    const size = calcStringSize(formattedResult);

    const overSizeLimit = sizeLimit && size > sizeLimit;
    const depthLimit = !depth || depth > 1;

    Logger.log(`Data Size: ${size} KB`);

    if (overSizeLimit) {
      Logger.log(`Data Size over size limit: ${sizeLimit} KB`);
      if (depthLimit) {
        Logger.log(`send error to client, need to perform truncated reading`);
        return {
          isError: true,
          content: [
            {
              type: "text",
              text: `The data size of file ${fileKey} ${nodeId ? `node ${nodeId}` : ""} is ${size} KB, exceeded the limit of ${sizeLimit} KB, please performing truncated reading`,
            },
          ],
        };
      }
    }

    Logger.log("Sending result to client");
    return {
      content: [{ type: "text", text: formattedResult }],
    };
  } catch (error) {
    const message = error instanceof Error ? error.message : JSON.stringify(error);
    Logger.error(`Error fetching file ${params.fileKey}:`, message);
    return {
      isError: true,
      content: [{ type: "text", text: `Error fetching file: ${message}` }],
    };
  }
}

// Export tool configuration
export const getFigmaDataTool = {
  name: "get_figma_data",
  description: `Get comprehensive Figma file data including layout, content, visuals, and component information
${
  needLimitPrompt
    ? `## Figma Data Size Guidelines
- **Over ${sizeLimit}KB**: with \`depth: 1\` to get structure only, enter the *Pruning Reading Strategy*
- **Under ${sizeLimit}KB**: Get full data without depth parameter

## Figma Data Pruning Reading Strategy

**IMPORTANT: Work incrementally, not comprehensively.**

### Core Principle
Retrieve and implement ONE node at a time. Don't try to understand the entire design upfront.

### Pruning Reading Process
1. **Start Small**: Get shallow data (depth: 1) of the main node to see basic information of itself and children nodes
2. **Pick One**: Choose one from the child nodes to implement completely
3. **Get Full Data**: Retrieve complete data for that one node only
4. **Implement**: Implement specifically according to user needs based on the content of the node
5. **Repeat**: Move to the next node only after the current one is done

### Key Point
**Don't analyze multiple nodes in parallel.** Focus on implementing one complete, working node at a time. This avoids context overload and produces better results.`
    : ""
}`,
  parameters,
  handler: getFigmaData,
} as const;