From cc328abe6c38acd6dc3a0531dbae8875c1093261 Mon Sep 17 00:00:00 2001
From: Quang Tran <16215255+trmquang93@users.noreply.github.com>
Date: Fri, 17 Apr 2026 06:28:47 +0700
Subject: [PATCH 1/2] feat: add get_screen_code MCP tool and hasCode flag

Expose the stored source HTML of code-generated screens through a
dedicated MCP tool, and surface a hasCode indicator in list_screens so
agents can discover which screens are round-trippable without pulling
the full screen payload.
---
 mcp-server/src/tools/screen-tools.js | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/mcp-server/src/tools/screen-tools.js b/mcp-server/src/tools/screen-tools.js
index 72b3a33..59ed42e 100644
--- a/mcp-server/src/tools/screen-tools.js
+++ b/mcp-server/src/tools/screen-tools.js
@@ -97,6 +97,17 @@ export const screenTools = [
       required: ["screenId"],
     },
   },
+  {
+    name: "get_screen_code",
+    description: "Get the source HTML for a screen that was created with code (via create_screen, batch_create_screens, update_screen_image, or Figma paste). Returns hasCode:false when the screen has no stored source (e.g., uploaded image or blank screen).",
+    inputSchema: {
+      type: "object",
+      properties: {
+        screenId: { type: "string", description: "ID of the screen" },
+      },
+      required: ["screenId"],
+    },
+  },
   {
     name: "update_screen_image",
     description: "Re-render a screen's image from new HTML content. Use inline styles only (no <style> tags).",
@@ -226,6 +237,7 @@ export async function handleScreenTool(name, args, state, renderer) {
           y: s.y,
           hotspotCount: (s.hotspots || []).length,
           hasImage: !!s.imageData,
+          hasCode: !!s.sourceHtml,
           description: s.description || "",
           status: s.status || "new",
           tbd: s.tbd || false,
@@ -278,6 +290,19 @@ export async function handleScreenTool(name, args, state, renderer) {
       return { __contentBlocks: content };
     }
 
+    case "get_screen_code": {
+      const screen = state.getScreen(args.screenId);
+      if (!screen) throw new Error(`Screen not found: ${args.screenId}`);
+      const hasCode = !!screen.sourceHtml;
+      return {
+        screenId: screen.id,
+        name: screen.name,
+        hasCode,
+        sourceHtml: screen.sourceHtml || null,
+        figmaSource: screen.figmaSource || null,
+      };
+    }
+
     case "update_screen_image": {
       const screen = state.getScreen(args.screenId);
       if (!screen) throw new Error(`Screen not found: ${args.screenId}`);

From 0cda71b2ae312badffe8e8fef498b723804b5d73 Mon Sep 17 00:00:00 2001
From: Quang Tran <16215255+trmquang93@users.noreply.github.com>
Date: Fri, 17 Apr 2026 06:31:17 +0700
Subject: [PATCH 2/2] chore: simplify get_screen_code handler

Inline the one-use hasCode local and drop redundant `|| null` fallbacks
on sourceHtml/figmaSource (both fields are initialized to null, so the
coalesce never fires).
---
 mcp-server/src/tools/screen-tools.js | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/mcp-server/src/tools/screen-tools.js b/mcp-server/src/tools/screen-tools.js
index 59ed42e..f9cd264 100644
--- a/mcp-server/src/tools/screen-tools.js
+++ b/mcp-server/src/tools/screen-tools.js
@@ -293,13 +293,12 @@ export async function handleScreenTool(name, args, state, renderer) {
     case "get_screen_code": {
       const screen = state.getScreen(args.screenId);
       if (!screen) throw new Error(`Screen not found: ${args.screenId}`);
-      const hasCode = !!screen.sourceHtml;
       return {
         screenId: screen.id,
         name: screen.name,
-        hasCode,
-        sourceHtml: screen.sourceHtml || null,
-        figmaSource: screen.figmaSource || null,
+        hasCode: !!screen.sourceHtml,
+        sourceHtml: screen.sourceHtml,
+        figmaSource: screen.figmaSource,
       };
     }