feat: add progress notifications and async tree walker

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<SimplifiedDesign> {
   // 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,

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";

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<void> {
+  nodesProcessed++;
+  if (nodesProcessed % YIELD_INTERVAL === 0) {
+    await new Promise<void>((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<SimplifiedNode | null> {
   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

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),
     );
   }
 }

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<ServerRequest, ServerNotification>;
+
+/** 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<void> {
+  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);
+}

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<typeof parametersSchema>;
 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<number, string[]>(); // 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

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 }],

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<ActiveConnection>();
+
 /**
  * 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<void> {
     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();
   });
 }