diff --git a/src/extractors/design-extractor.ts b/src/extractors/design-extractor.ts index b33e7621..e5e16f75 100644 --- a/src/extractors/design-extractor.ts +++ b/src/extractors/design-extractor.ts @@ -14,18 +14,18 @@ import { extractFromDesign } from "./node-walker.js"; /** * Extract a complete SimplifiedDesign from raw Figma API response using extractors. */ -export function simplifyRawFigmaObject( +export async function simplifyRawFigmaObject( apiResponse: GetFileResponse | GetFileNodesResponse, nodeExtractors: ExtractorFn[], options: TraversalOptions = {}, -): SimplifiedDesign { +): Promise { // Extract components, componentSets, and raw nodes from API response const { metadata, rawNodes, components, componentSets, extraStyles } = parseAPIResponse(apiResponse); // Process nodes using the flexible extractor system const globalVars: TraversalContext["globalVars"] = { styles: {}, extraStyles }; - const { nodes: extractedNodes, globalVars: finalGlobalVars } = extractFromDesign( + const { nodes: extractedNodes, globalVars: finalGlobalVars } = await extractFromDesign( rawNodes, nodeExtractors, options, diff --git a/src/extractors/index.ts b/src/extractors/index.ts index 05e6892e..64d785ec 100644 --- a/src/extractors/index.ts +++ b/src/extractors/index.ts @@ -8,7 +8,7 @@ export type { } from "./types.js"; // Core traversal function -export { extractFromDesign } from "./node-walker.js"; +export { extractFromDesign, getNodesProcessed } from "./node-walker.js"; // Design-level extraction (unified nodes + components) export { simplifyRawFigmaObject } from "./design-extractor.js"; diff --git a/src/extractors/node-walker.ts b/src/extractors/node-walker.ts index c3938a2f..01ff994d 100644 --- a/src/extractors/node-walker.ts +++ b/src/extractors/node-walker.ts @@ -9,6 +9,24 @@ import type { SimplifiedNode, } from "./types.js"; +// Yield the event loop every N nodes so heartbeats, SIGINT, and +// other async work can run during large file processing. +// Yield the event loop every N nodes so heartbeats, SIGINT, and +// other async work can run during large file processing. +const YIELD_INTERVAL = 100; +let nodesProcessed = 0; + +export function getNodesProcessed(): number { + return nodesProcessed; +} + +async function maybeYield(): Promise { + nodesProcessed++; + if (nodesProcessed % YIELD_INTERVAL === 0) { + await new Promise((resolve) => setImmediate(resolve)); + } +} + /** * Extract data from Figma nodes using a flexible, single-pass approach. * @@ -18,21 +36,25 @@ import type { * @param globalVars - Global variables for style deduplication * @returns Object containing processed nodes and updated global variables */ -export function extractFromDesign( +export async function extractFromDesign( nodes: FigmaDocumentNode[], extractors: ExtractorFn[], options: TraversalOptions = {}, globalVars: GlobalVars = { styles: {} }, -): { nodes: SimplifiedNode[]; globalVars: GlobalVars } { +): Promise<{ nodes: SimplifiedNode[]; globalVars: GlobalVars }> { const context: TraversalContext = { globalVars, currentDepth: 0, }; - const processedNodes = nodes - .filter((node) => shouldProcessNode(node, options)) - .map((node) => processNodeWithExtractors(node, extractors, context, options)) - .filter((node): node is SimplifiedNode => node !== null); + nodesProcessed = 0; + + const processedNodes: SimplifiedNode[] = []; + for (const node of nodes) { + if (!shouldProcessNode(node, options)) continue; + const result = await processNodeWithExtractors(node, extractors, context, options); + if (result !== null) processedNodes.push(result); + } return { nodes: processedNodes, @@ -43,16 +65,18 @@ export function extractFromDesign( /** * Process a single node with all provided extractors in one pass. */ -function processNodeWithExtractors( +async function processNodeWithExtractors( node: FigmaDocumentNode, extractors: ExtractorFn[], context: TraversalContext, options: TraversalOptions, -): SimplifiedNode | null { +): Promise { if (!shouldProcessNode(node, options)) { return null; } + await maybeYield(); + // Always include base metadata const result: SimplifiedNode = { id: node.id, @@ -75,10 +99,12 @@ function processNodeWithExtractors( // Use the same pattern as the existing parseNode function if (hasValue("children", node) && node.children.length > 0) { - const children = node.children - .filter((child) => shouldProcessNode(child, options)) - .map((child) => processNodeWithExtractors(child, extractors, childContext, options)) - .filter((child): child is SimplifiedNode => child !== null); + const children: SimplifiedNode[] = []; + for (const child of node.children) { + if (!shouldProcessNode(child, options)) continue; + const processed = await processNodeWithExtractors(child, extractors, childContext, options); + if (processed !== null) children.push(processed); + } if (children.length > 0) { // Allow custom logic to modify parent and control which children to include diff --git a/src/mcp/index.ts b/src/mcp/index.ts index 3d26347f..a73dfe72 100644 --- a/src/mcp/index.ts +++ b/src/mcp/index.ts @@ -1,6 +1,7 @@ import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { FigmaService, type FigmaAuthOptions } from "../services/figma.js"; import { Logger } from "../utils/logger.js"; +import type { ToolExtra } from "./progress.js"; import { downloadFigmaImagesTool, getFigmaDataTool, @@ -57,8 +58,8 @@ function registerTools( inputSchema: getFigmaDataTool.parametersSchema, annotations: { readOnlyHint: true }, }, - (params: GetFigmaDataParams) => - getFigmaDataTool.handler(params, figmaService, options.outputFormat), + (params: GetFigmaDataParams, extra: ToolExtra) => + getFigmaDataTool.handler(params, figmaService, options.outputFormat, extra), ); if (!options.skipImageDownloads) { @@ -70,8 +71,8 @@ function registerTools( inputSchema: downloadFigmaImagesTool.parametersSchema, annotations: { openWorldHint: true }, }, - (params: DownloadImagesParams) => - downloadFigmaImagesTool.handler(params, figmaService, options.imageDir), + (params: DownloadImagesParams, extra: ToolExtra) => + downloadFigmaImagesTool.handler(params, figmaService, options.imageDir, extra), ); } } diff --git a/src/mcp/progress.ts b/src/mcp/progress.ts new file mode 100644 index 00000000..a10c7a0c --- /dev/null +++ b/src/mcp/progress.ts @@ -0,0 +1,49 @@ +import type { RequestHandlerExtra } from "@modelcontextprotocol/sdk/shared/protocol.js"; +import type { ServerNotification, ServerRequest } from "@modelcontextprotocol/sdk/types.js"; + +export type ToolExtra = RequestHandlerExtra; + +/** No-ops silently when the client didn't ask for progress (no progressToken). */ +export async function sendProgress( + extra: ToolExtra, + progress: number, + total?: number, + message?: string, +): Promise { + const progressToken = extra._meta?.progressToken; + if (progressToken === undefined) return; + + await extra.sendNotification({ + method: "notifications/progress", + params: { progressToken, progress, total, message }, + }); +} + +/** + * Send periodic progress notifications during a long-running operation. + * Keeps clients with resetTimeoutOnProgress alive during slow I/O like + * Figma API calls that can take up to ~55 seconds. Returns a stop function + * that must be called when the operation completes or errors. + */ +export function startProgressHeartbeat( + extra: ToolExtra, + message: string | (() => string), + intervalMs = 3_000, +): () => void { + const progressToken = extra._meta?.progressToken; + if (progressToken === undefined) return () => {}; + + let tick = 0; + const interval = setInterval(() => { + tick++; + const msg = typeof message === "function" ? message() : message; + extra + .sendNotification({ + method: "notifications/progress", + params: { progressToken, progress: tick, message: msg }, + }) + .catch(() => clearInterval(interval)); + }, intervalMs); + + return () => clearInterval(interval); +} diff --git a/src/mcp/tools/download-figma-images-tool.ts b/src/mcp/tools/download-figma-images-tool.ts index e9dfd1f8..866355fe 100644 --- a/src/mcp/tools/download-figma-images-tool.ts +++ b/src/mcp/tools/download-figma-images-tool.ts @@ -2,6 +2,7 @@ import path from "path"; import { z } from "zod"; import { FigmaService } from "../../services/figma.js"; import { Logger } from "../../utils/logger.js"; +import { sendProgress, startProgressHeartbeat, type ToolExtra } from "../progress.js"; const parameters = { fileKey: z @@ -85,7 +86,8 @@ export type DownloadImagesParams = z.infer; async function downloadFigmaImages( params: DownloadImagesParams, figmaService: FigmaService, - imageDir?: string, + imageDir: string | undefined, + extra: ToolExtra, ) { try { const { fileKey, nodes, localPath, pngScale = 2 } = parametersSchema.parse(params); @@ -109,6 +111,8 @@ async function downloadFigmaImages( }; } + await sendProgress(extra, 0, 3, "Resolving image downloads"); + // Process nodes: collect unique downloads and track which requests they satisfy const downloadItems = []; const downloadToRequests = new Map(); // download index -> requested filenames @@ -171,11 +175,20 @@ async function downloadFigmaImages( } } - const allDownloads = await figmaService.downloadImages(fileKey, resolvedPath, downloadItems, { - pngScale, - }); + await sendProgress(extra, 1, 3, `Resolved ${downloadItems.length} images, downloading`); + const stopHeartbeat = startProgressHeartbeat(extra, "Downloading images"); + + let allDownloads; + try { + allDownloads = await figmaService.downloadImages(fileKey, resolvedPath, downloadItems, { + pngScale, + }); + } finally { + stopHeartbeat(); + } const successCount = allDownloads.filter(Boolean).length; + await sendProgress(extra, 2, 3, `Downloaded ${successCount} images, formatting response`); // Format results with aliases const imagesList = allDownloads diff --git a/src/mcp/tools/get-figma-data-tool.ts b/src/mcp/tools/get-figma-data-tool.ts index e97637a1..3204b2ad 100644 --- a/src/mcp/tools/get-figma-data-tool.ts +++ b/src/mcp/tools/get-figma-data-tool.ts @@ -5,9 +5,11 @@ import { simplifyRawFigmaObject, allExtractors, collapseSvgContainers, + getNodesProcessed, } from "~/extractors/index.js"; import yaml from "js-yaml"; import { Logger, writeLogs } from "~/utils/logger.js"; +import { sendProgress, startProgressHeartbeat, type ToolExtra } from "~/mcp/progress.js"; const parameters = { fileKey: z @@ -42,6 +44,7 @@ async function getFigmaData( params: GetFigmaDataParams, figmaService: FigmaService, outputFormat: "yaml" | "json", + extra: ToolExtra, ) { try { const { fileKey, nodeId: rawNodeId, depth } = parametersSchema.parse(params); @@ -55,19 +58,37 @@ async function getFigmaData( } ${fileKey}`, ); + await sendProgress(extra, 0, 4, "Fetching design data from Figma API"); + const stopHeartbeat = startProgressHeartbeat(extra, "Waiting for Figma API response"); + // 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); + try { + if (nodeId) { + rawApiResponse = await figmaService.getRawNode(fileKey, nodeId, depth); + } else { + rawApiResponse = await figmaService.getRawFile(fileKey, depth); + } + } finally { + stopHeartbeat(); } + await sendProgress(extra, 1, 4, "Fetched design data, simplifying"); + const stopSimplifyHeartbeat = startProgressHeartbeat( + extra, + () => `Simplifying design data (${getNodesProcessed()} nodes processed)`, + ); + // Use unified design extraction (handles nodes + components consistently) - const simplifiedDesign = simplifyRawFigmaObject(rawApiResponse, allExtractors, { - maxDepth: depth, - afterChildren: collapseSvgContainers, - }); + let simplifiedDesign; + try { + simplifiedDesign = await simplifyRawFigmaObject(rawApiResponse, allExtractors, { + maxDepth: depth, + afterChildren: collapseSvgContainers, + }); + } finally { + stopSimplifyHeartbeat(); + } writeLogs("figma-simplified.json", simplifiedDesign); @@ -77,6 +98,8 @@ async function getFigmaData( } styles`, ); + await sendProgress(extra, 2, 4, "Simplified design, serializing response"); + const { nodes, globalVars, ...metadata } = simplifiedDesign; const result = { metadata, @@ -88,6 +111,8 @@ async function getFigmaData( const formattedResult = outputFormat === "json" ? JSON.stringify(result, null, 2) : yaml.dump(result); + await sendProgress(extra, 3, 4, "Serialized, sending response"); + Logger.log("Sending result to client"); return { content: [{ type: "text" as const, text: formattedResult }], diff --git a/src/server.ts b/src/server.ts index fe2c2c36..41d3dea2 100644 --- a/src/server.ts +++ b/src/server.ts @@ -11,6 +11,12 @@ import { ErrorCode } from "@modelcontextprotocol/sdk/types.js"; let httpServer: Server | null = null; +type ActiveConnection = { + transport: StreamableHTTPServerTransport; + server: McpServer; +}; +const activeConnections = new Set(); + /** * Start the MCP server in either stdio or HTTP mode. */ @@ -57,7 +63,10 @@ export async function startHttpServer( Logger.log("Received StreamableHTTP request"); const transport = new StreamableHTTPServerTransport({ sessionIdGenerator: undefined }); const mcpServer = createMcpServer(); + const conn: ActiveConnection = { transport, server: mcpServer }; + activeConnections.add(conn); res.on("close", () => { + activeConnections.delete(conn); transport.close(); mcpServer.close(); }); @@ -114,11 +123,19 @@ export async function stopHttpServer(): Promise { throw new Error("HTTP server is not running"); } + // Gracefully close all active MCP connections before tearing down the server + for (const conn of activeConnections) { + await conn.transport.close(); + await conn.server.close(); + } + activeConnections.clear(); + return new Promise((resolve, reject) => { httpServer!.close((err) => { httpServer = null; if (err) reject(err); else resolve(); }); + httpServer!.closeAllConnections(); }); } diff --git a/src/tests/path-validation.test.ts b/src/tests/path-validation.test.ts index 6abefe7a..b642537b 100644 --- a/src/tests/path-validation.test.ts +++ b/src/tests/path-validation.test.ts @@ -2,11 +2,17 @@ import path from "path"; import { describe, expect, it } from "vitest"; import { downloadFigmaImagesTool } from "~/mcp/tools/download-figma-images-tool.js"; import { downloadFigmaImage } from "~/utils/common.js"; +import type { ToolExtra } from "~/mcp/progress.js"; const stubFigmaService = { downloadImages: () => Promise.resolve([]), } as unknown as Parameters[1]; +const stubExtra = { + sendNotification: () => Promise.resolve(), + signal: AbortSignal.timeout(30_000), +} as unknown as ToolExtra; + const validParams = { fileKey: "abc123", nodes: [{ nodeId: "1:2", fileName: "test.png" }], @@ -21,6 +27,7 @@ describe("download path validation", () => { { ...validParams, localPath: "../../etc" }, stubFigmaService, imageDir, + stubExtra, ); expect(result.isError).toBe(true); @@ -33,6 +40,7 @@ describe("download path validation", () => { { ...validParams, localPath: "/../../etc" }, stubFigmaService, imageDir, + stubExtra, ); expect(result.isError).toBe(true); @@ -44,6 +52,7 @@ describe("download path validation", () => { { ...validParams, localPath: "public/images" }, stubFigmaService, imageDir, + stubExtra, ); expect(result.isError).toBeUndefined(); @@ -58,6 +67,7 @@ describe("download path validation", () => { { ...validParams, localPath: "project/src/static/images/test" }, stubFigmaService, driveRoot, + stubExtra, ); expect(result.isError).toBeUndefined(); @@ -68,6 +78,7 @@ describe("download path validation", () => { { ...validParams, localPath: "/public/images" }, stubFigmaService, imageDir, + stubExtra, ); expect(result.isError).toBeUndefined(); diff --git a/src/tests/tree-walker.test.ts b/src/tests/tree-walker.test.ts new file mode 100644 index 00000000..dde978d4 --- /dev/null +++ b/src/tests/tree-walker.test.ts @@ -0,0 +1,143 @@ +import { describe, expect, it } from "vitest"; +import { extractFromDesign } from "~/extractors/node-walker.js"; +import { allExtractors, collapseSvgContainers } from "~/extractors/built-in.js"; +import { simplifyRawFigmaObject } from "~/extractors/design-extractor.js"; +import type { GetFileResponse } from "@figma/rest-api-spec"; +import type { Node as FigmaNode } from "@figma/rest-api-spec"; + +// Minimal Figma node factory — only the fields the walker actually reads. +// The Figma types are deeply discriminated unions; we cast through unknown +// because tests only need the subset of fields the walker touches. +function makeNode(overrides: Record): FigmaNode { + return { visible: true, ...overrides } as unknown as FigmaNode; +} + +// A small but representative node tree: +// Page +// ├── Frame "Header" (visible) +// │ ├── Text "Title" +// │ └── Rectangle "Bg" (invisible) +// ├── Frame "Body" +// │ └── Frame "Card" +// │ └── Text "Label" +// └── Vector "Icon" (becomes IMAGE-SVG) +const fixtureNodes: FigmaNode[] = [ + makeNode({ + id: "1:1", + name: "Header", + type: "FRAME", + children: [ + makeNode({ id: "1:2", name: "Title", type: "TEXT", characters: "Hello" }), + makeNode({ id: "1:3", name: "Bg", type: "RECTANGLE", visible: false }), + ], + }), + makeNode({ + id: "2:1", + name: "Body", + type: "FRAME", + children: [ + makeNode({ + id: "2:2", + name: "Card", + type: "FRAME", + children: [makeNode({ id: "2:3", name: "Label", type: "TEXT", characters: "World" })], + }), + ], + }), + makeNode({ id: "3:1", name: "Icon", type: "VECTOR" }), +]; + +describe("extractFromDesign", () => { + it("produces correct node structure from a nested tree", async () => { + const { nodes } = await extractFromDesign(fixtureNodes, allExtractors); + + // Top-level: Header, Body, Icon (3 nodes — Bg is invisible, filtered out) + expect(nodes).toHaveLength(3); + expect(nodes.map((n) => n.name)).toEqual(["Header", "Body", "Icon"]); + + // Header has 1 child (Title only — Bg is invisible) + const header = nodes[0]; + expect(header.children).toHaveLength(1); + expect(header.children![0].name).toBe("Title"); + expect(header.children![0].text).toBe("Hello"); + + // Body > Card > Label + const body = nodes[1]; + expect(body.children).toHaveLength(1); + expect(body.children![0].name).toBe("Card"); + expect(body.children![0].children).toHaveLength(1); + expect(body.children![0].children![0].name).toBe("Label"); + expect(body.children![0].children![0].text).toBe("World"); + + // Vector becomes IMAGE-SVG + const icon = nodes[2]; + expect(icon.type).toBe("IMAGE-SVG"); + expect(icon.children).toBeUndefined(); + }); + + it("respects maxDepth option", async () => { + const { nodes } = await extractFromDesign(fixtureNodes, allExtractors, { maxDepth: 1 }); + + // At depth 0 we get top-level nodes, depth 1 gets their direct children, no deeper + const header = nodes.find((n) => n.name === "Header")!; + expect(header.children).toHaveLength(1); + expect(header.children![0].name).toBe("Title"); + + // Body's child "Card" is at depth 1 — it should exist but have no children + const body = nodes.find((n) => n.name === "Body")!; + expect(body.children).toHaveLength(1); + expect(body.children![0].name).toBe("Card"); + expect(body.children![0].children).toBeUndefined(); + }); + + it("accumulates global style variables across nodes", async () => { + const styledNode = makeNode({ + id: "4:1", + name: "Styled", + type: "FRAME", + fills: [{ type: "SOLID", color: { r: 1, g: 0, b: 0, a: 1 }, visible: true }], + }); + + const { globalVars } = await extractFromDesign([styledNode], allExtractors); + + // The fill should be extracted into a global variable + expect(Object.keys(globalVars.styles).length).toBeGreaterThan(0); + }); +}); + +describe("simplifyRawFigmaObject", () => { + it("produces a complete SimplifiedDesign from a mock API response", async () => { + const mockResponse = { + name: "Test File", + document: { + id: "0:0", + name: "Document", + type: "DOCUMENT", + children: fixtureNodes, + visible: true, + }, + components: {}, + componentSets: {}, + styles: {}, + schemaVersion: 0, + version: "1", + role: "owner", + lastModified: "2024-01-01", + thumbnailUrl: "", + editorType: "figma", + } as unknown as GetFileResponse; + + const result = await simplifyRawFigmaObject(mockResponse, allExtractors, { + afterChildren: collapseSvgContainers, + }); + + expect(result.name).toBe("Test File"); + expect(result.nodes).toHaveLength(3); + expect(result.nodes.map((n) => n.name)).toEqual(["Header", "Body", "Icon"]); + + // Verify full depth traversal happened + const label = result.nodes[1].children![0].children![0]; + expect(label.name).toBe("Label"); + expect(label.text).toBe("World"); + }); +});