feat: add progress notifications and async tree walker (#305)

* refactor!: switch HTTP transport to stateless mode

Each request creates its own McpServer and transport with
sessionIdGenerator: undefined. Removes the session registry,
SSE transport, /messages endpoint, and progress notification
interval. StreamableHTTP is served at both /mcp and /sse for
backward compatibility. GET and DELETE return 405.

* docs: update transport references after stateless HTTP refactor

* fix: catch async errors in Express 4 route handler

Express 4 does not catch rejected promises from async handlers.
Without a try/catch, a failure in connect() or handleRequest()
causes an unhandled rejection that crashes the process.

* chore: upgrade Express to v5

Aligns with the MCP SDK which already depends on Express 5.
Express 5 natively catches async route handler rejections,
so the manual try/catch added in the previous commit is no
longer needed.

* fix: add per-request cleanup and JSON-RPC error middleware

Close transport and McpServer when the response ends, matching
the SDK's recommended pattern for stateless servers. Add Express
error-handling middleware that returns a JSON-RPC error response
instead of Express's default HTML 500.

* fix: use SDK's createMcpExpressApp for DNS rebinding protection

Replace raw express() with the SDK's createMcpExpressApp(), which
applies localhost DNS rebinding protection middleware automatically.
This also handles express.json() globally, removing the need for
per-route body parsing.

* feat: add progress notifications to tool handlers

Thread the SDK's RequestHandlerExtra through registerTool wiring
to tool handlers. Add sendProgress helper that sends
notifications/progress when the client provides a progressToken.

get_figma_data reports progress after fetching from the Figma API
and after processing design data (2 phases).

download_figma_images reports progress after resolving download
items and after completing all downloads (2 phases).

* refactor: extract sendProgress helper into mcp/progress.ts

Break the circular dependency where tools imported from their
own registration module (mcp/index.ts -> tools -> mcp/index.ts).

* fix: send initial progress notification before slow I/O

Progress was only sent after the Figma API fetch completed,
which defeats the purpose — the timeout fires during the fetch.
Now sends progress at 0/3 before the slow operation starts, so
clients with resetTimeoutOnProgress get an immediate signal.

* feat: add progress heartbeat during long Figma API calls

Figma API calls can take up to ~55 seconds. A single progress
notification before the call isn't enough if the client timeout
is shorter. startProgressHeartbeat sends periodic notifications
every 5 seconds during slow I/O, keeping clients with
resetTimeoutOnProgress alive. The heartbeat stops as soon as
the operation completes or errors.

* fix: add progress between synchronous processing phases

The tree walker and YAML serializer are synchronous and block
the event loop, preventing heartbeat intervals from firing.
Send progress notifications between each phase so the client
gets a fresh timeout window before each synchronous step:
API fetch -> simplify -> serialize -> return.

* feat: make tree walker async to unblock event loop

The synchronous tree walk blocks the event loop for large Figma
files (75MB+), preventing progress heartbeats from firing, SIGINT
from being handled, and the HTTP server from shutting down cleanly.

Make extractFromDesign and simplifyRawFigmaObject async. The tree
walker yields the event loop every 500 nodes via setImmediate,
allowing heartbeats and signal handlers to run during processing.

Also fix progress total (3 -> 4), add diagnostic logging to
sendProgress, and force-close keep-alive connections on shutdown.

* fix: run heartbeat during simplification phase

The heartbeat was only active during the Figma API call but not
during the tree walk simplification, which is the actual slow
part for large files. Now that the tree walker yields the event
loop, the heartbeat can fire during simplification.

* feat: show node count in heartbeat during simplification

Report how many nodes have been processed in each heartbeat
message (e.g. "Simplifying design data (3500 nodes processed)").
The heartbeat message function is now evaluated dynamically on
each tick.

* fix: tighten yield interval and heartbeat frequency

Reduce yield interval from 500 to 100 nodes — later nodes in
large files are deeper and more complex, so 500 nodes could
take longer than the heartbeat interval. Reduce heartbeat from
5s to 3s for more margin against 10s client timeouts.

* fix: gracefully close MCP connections before server shutdown

Track active per-request transport/server pairs. On shutdown,
close each transport and server cleanly before terminating HTTP
connections. This gives clients a proper stream close instead of
an abrupt disconnect.

* chore: remove diagnostic logging from sendProgress

* test: add tree walker regression tests

Test extractFromDesign and simplifyRawFigmaObject with a
representative node tree covering visibility filtering, depth
limits, child recursion, VECTOR->IMAGE-SVG type mapping, and
global style variable accumulation. Verified against both the
original sync and new async implementations to confirm the
async conversion produces identical output.

Author: GLips

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