Unify figure store globally, enable cross-agent figure references

- Replace per-tab IMAGE_STORES with single global FIGURE_STORE shared by code and image agents
- Unify naming: image_N → figure_T{tab}_N across all agents
- Store values as {type, data} dicts consistently
- Restore figure registry from workspace on session reload
- Remove separate 'images' handling in frontend (unified under 'figures')
- Add cross-agent figure reference instructions to command center prompt
- Fix double-wrapping bug in image agent nudge path

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

backend/agents.py

@@ -53,7 +53,10 @@ AGENT_REGISTRY = {
             "- **Code agent**: data analysis, code execution, visualizations, debugging\n"
             "- **Research agent**: ONLY deep multi-source analysis, comparisons, reports\n"
             "- **Image agent**: generating or editing images (ONLY when the user explicitly asks to generate/create an image — never for finding/showing existing photos)\n\n"
-            "When delegating, provide a clear objective, scope boundaries, and expected output format.\n\n"
+            "When delegating, provide a clear objective, scope boundaries, and expected output format.\n"
+            "**Figures are shared across agents.** If a previous agent produced a figure (e.g., figure_T3_1), "
+            "you can pass its reference to another agent — for example, ask the image agent to edit figure_T3_1 "
+            "or the code agent to process it. Just include the reference name in the task description.\n\n"
             "## Task Decomposition — ALWAYS parallelize\n\n"
             "**RULE: When a request mentions multiple distinct entities or topics, "
             "launch a separate agent for each.** Never combine multiple lookups into one agent.\n\n"
@@ -67,8 +70,8 @@ AGENT_REGISTRY = {
             "- Do NOT save/create files unless the user explicitly requests it.\n"
             "- Reuse task_id when a follow-up relates to an existing agent (preserves context and kernel).\n"
             "- Include key findings in YOUR response — don't just say \"see the agent result\".\n"
-            "- **ALWAYS embed figures/images** from sub-agents in your response using their reference tags "
-            "(e.g., <figure_1>, <image_1>). Sub-agent results are collapsed — if you don't embed the figure, the user won't see it.\n"
+            "- **ALWAYS embed figures** from sub-agents in your response using the exact reference tags from the agent result "
+            "(e.g., <figure_T3_1>). Sub-agent results are collapsed — if you don't embed the figure, the user won't see it.\n"
             "- If an agent was aborted by the user, acknowledge it and ask how to proceed — don't re-launch."
         ),
         "tool": None,
@@ -142,17 +145,18 @@ AGENT_REGISTRY = {
             "**Code runs in a remote sandbox, NOT locally.** "
             "Use upload_files before processing user files, download_files to send results back.\n\n"
             "## Guidelines\n\n"
-            "- **Figures**: Call plt.show() — figures are auto-captured as figure_1, figure_2, etc. "
+            "- **Figures**: Call plt.show() — figures are auto-captured with names like figure_T4_1, figure_T4_2, etc. "
+            "The exact names appear in the execution output. "
             "NEVER use both plt.savefig() and plt.show() (creates duplicates). "
-            "To display a figure, embed it in your result text as <figure_1> — do NOT use show_html with an <img> tag.\n"
+            "To display a figure, embed the exact figure name from the output in your result text — do NOT use show_html with an <img> tag.\n"
             "- **Files**: Do NOT save/download unless explicitly requested. Never overwrite without permission.\n"
             "- Execute code incrementally and reflect on output between steps.\n\n"
             "## CRITICAL: You MUST provide a <result> tag\n\n"
             "Keep results SHORT (1-2 sentences). The user can see code and output above.\n"
-            "Use <figure_1> (self-closing) to embed figures.\n\n"
+            "Use the exact figure name from the execution output (e.g., <figure_T4_1>) to embed figures.\n\n"
             "<result>\n"
             "Here's the sine function plot:\n\n"
-            "<figure_1>\n"
+            "<figure_T4_1>\n"
             "</result>\n"
         ),
         "tool": {
@@ -233,12 +237,12 @@ AGENT_REGISTRY = {
         "system_prompt": (
             "You are a creative AI assistant with image tools.\n\n"
             "## Tools\n\n"
-            "- **generate_image(prompt)**: Generate from text. Returns image reference (e.g., 'image_1').\n"
+            "- **generate_image(prompt)**: Generate from text. Returns figure reference (e.g., 'figure_T4_1').\n"
             "- **edit_image(prompt, source)**: Edit/transform an image. Source: URL, file path, or reference.\n"
             "- **read_image(source)**: Load a raster image (PNG, JPEG, GIF, WebP, BMP). "
-            "SVG NOT supported. Returns image reference.\n"
+            "SVG NOT supported. Returns figure reference.\n"
             "- **save_image(source, filename)**: Save an image to the workspace as PNG. "
-            "Source: reference (e.g., 'image_1') or URL.\n\n"
+            "Source: reference (e.g., 'figure_T4_1') or URL.\n\n"
             "## Strategy\n\n"
             "1. If user provides a URL/file, use read_image first to load it\n"
             "2. Use generate_image ONLY when explicitly asked to generate/create an image — "
@@ -246,10 +250,11 @@ AGENT_REGISTRY = {
             "3. Use edit_image to transform existing ones\n"
             "4. Write detailed prompts. Describe what you see and iterate if needed.\n\n"
             "## CRITICAL: You MUST provide a <result> tag\n\n"
-            "Use <image_1> (self-closing) to embed images in your result.\n\n"
+            "Use the exact figure reference from tool output to embed figures in your result.\n"
+            "Figure references are self-closing tags like <figure_T4_2> — do NOT add a closing </figure_T4_2> tag.\n\n"
             "<result>\n"
             "Here's the comic version of your image:\n\n"
-            "<image_2>\n"
+            "<figure_T4_2>\n"
             "</result>\n"
         ),
         "tool": {

backend/code.py

@@ -174,7 +174,7 @@ def download_files_from_sandbox(sbx: Sandbox, files: List[Dict], files_root: str
     return "\n".join(results)
 
 
-def stream_code_execution(client, model: str, messages: List[Dict], sbx: Sandbox, files_root: str = None, extra_params: Optional[Dict] = None, abort_event=None, multimodal: bool = False, tab_id: str = "0"):
+def stream_code_execution(client, model: str, messages: List[Dict], sbx: Sandbox, files_root: str = None, extra_params: Optional[Dict] = None, abort_event=None, multimodal: bool = False, tab_id: str = "0", figure_store: Optional[Dict[str, dict]] = None):
     """
     Stream code execution results
 
@@ -195,7 +195,10 @@ def stream_code_execution(client, model: str, messages: List[Dict], sbx: Sandbox
     done = False
     figure_counter = 0  # Track figure numbers
     figure_prefix = f"figure_T{tab_id}_"
-    figure_data = {}  # Store figure data by name for result rendering
+    # Use shared global store if provided, otherwise create local one
+    if figure_store is None:
+        figure_store = {}
+    figure_data = figure_store  # Alias for clarity in this function
     has_result = False
     debug_call_number = 0
 

backend/image.py

@@ -4,8 +4,8 @@ Image agent backend — multimodal agent with HuggingFace image generation tools
 Uses the same tool-calling loop pattern as agent.py:
   LLM call → parse tool_calls → execute → update history → repeat
 
-Key difference: maintains an image store (Dict[str, str]) mapping names like
-"image_1" to base64 data, so the VLM can reference images across tool calls
+Key difference: maintains a figure store (Dict[str, str]) mapping names like
+"figure_T1_1" to base64 data, so the VLM can reference images across tool calls
 without passing huge base64 strings in arguments.
 """
 import base64
@@ -61,7 +61,7 @@ def resize_image_for_vlm(base64_png: str) -> str:
 MAX_TURNS = 20
 
 
-def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, image_counter: int, default_gen_model: str = None, default_edit_model: str = None, files_root: str = None, image_prefix: str = "image_") -> dict:
+def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, image_counter: int, default_gen_model: str = None, default_edit_model: str = None, files_root: str = None, image_prefix: str = "figure_") -> dict:
     """
     Execute a tool by name and return result dict.
 
@@ -81,7 +81,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
         if base64_png:
             image_counter += 1
             name = f"{image_prefix}{image_counter}"
-            image_store[name] = base64_png
+            image_store[name] = {"type": "png", "data": base64_png}
             return {
                 "content": f"Image generated successfully as '{name}'. The image is attached.",
                 "image": base64_png,
@@ -104,7 +104,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
         # Resolve source: image store reference, URL, or local path
         source_bytes = None
         if source in image_store:
-            source_bytes = base64.b64decode(image_store[source])
+            source_bytes = base64.b64decode(image_store[source]["data"])
         else:
             source_base64 = execute_read_image(source, files_root=files_root)
             if source_base64:
@@ -112,7 +112,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
 
         if source_bytes is None:
             return {
-                "content": f"Could not resolve image source '{source}'. Use a URL or a reference from a previous tool call (e.g., 'image_1').",
+                "content": f"Could not resolve image source '{source}'. Use a URL or a reference from a previous tool call (e.g., 'figure_T1_1').",
                 "display": {"type": "edit_error", "source": source},
                 "image_counter": image_counter,
             }
@@ -122,7 +122,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
         if base64_png:
             image_counter += 1
             name = f"{image_prefix}{image_counter}"
-            image_store[name] = base64_png
+            image_store[name] = {"type": "png", "data": base64_png}
             return {
                 "content": f"Image edited successfully as '{name}'. The image is attached.",
                 "image": base64_png,
@@ -148,7 +148,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
         # Resolve source from image store or URL
         image_data = None
         if source in image_store:
-            image_data = base64.b64decode(image_store[source])
+            image_data = base64.b64decode(image_store[source]["data"])
         else:
             source_base64 = execute_read_image(source, files_root=files_root)
             if source_base64:
@@ -156,7 +156,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
 
         if image_data is None:
             return {
-                "content": f"Could not resolve image source '{source}'. Use a reference (e.g., 'image_1') or a URL.",
+                "content": f"Could not resolve image source '{source}'. Use a reference (e.g., 'figure_T1_1') or a URL.",
                 "display": {"type": "save_error", "source": source},
                 "image_counter": image_counter,
             }
@@ -187,7 +187,7 @@ def execute_tool(tool_name: str, args: dict, hf_token: str, image_store: dict, i
         if base64_png:
             image_counter += 1
             name = f"{image_prefix}{image_counter}"
-            image_store[name] = base64_png
+            image_store[name] = {"type": "png", "data": base64_png}
             return {
                 "content": f"Image loaded successfully as '{name}'. The image is attached.",
                 "image": base64_png,
@@ -227,7 +227,7 @@ def stream_image_execution(
     files_root: str = None,
     multimodal: bool = False,
     tab_id: str = "0",
-    image_store: Optional[Dict[str, str]] = None,
+    image_store: Optional[Dict[str, dict]] = None,
     image_counter: int = 0,
 ):
     """
@@ -239,7 +239,7 @@ def stream_image_execution(
       - tool_start: { tool, args }
       - tool_result: { tool, result, image? }
       - result_preview: { content }
-      - result: { content, images? }
+      - result: { content, figures? }
       - generating: {}
       - retry: { attempt, max_attempts, delay, message }
       - error: { content }
@@ -249,7 +249,7 @@ def stream_image_execution(
 
     turns = 0
     done = False
-    image_prefix = f"image_T{tab_id}_"
+    image_prefix = f"figure_T{tab_id}_"
 
     # Use provided persistent store, or create a local one as fallback
     if image_store is None:
@@ -302,8 +302,8 @@ def stream_image_execution(
 
         # Send result preview
         if result_content:
-            # Include image store so frontend can resolve <image_N> references
-            yield {"type": "result_preview", "content": result_content, "images": image_store}
+            figures = dict(image_store)
+            yield {"type": "result_preview", "content": result_content, "figures": figures}
 
         # --- Handle tool calls ---
         if tool_calls:
@@ -386,7 +386,8 @@ def stream_image_execution(
 
         # Send result if found
         if result_content:
-            yield {"type": "result", "content": result_content, "images": image_store}
+            figures = dict(image_store)
+            yield {"type": "result", "content": result_content, "figures": figures}
             result_sent = True
 
         # Signal between-turn processing
@@ -397,14 +398,16 @@ def stream_image_execution(
     if not result_sent:
         from .agents import nudge_for_result
         nudge_produced_result = False
-        for event in nudge_for_result(client, model, messages, extra_params=extra_params, extra_result_data={"images": image_store}, call_number=debug_call_number):
+        figures = dict(image_store)
+        for event in nudge_for_result(client, model, messages, extra_params=extra_params, extra_result_data={"figures": figures}, call_number=debug_call_number):
             yield event
             if event.get("type") == "result":
                 nudge_produced_result = True
 
-        # Final fallback: synthesize a result with all images
+        # Final fallback: synthesize a result with all figures
         if not nudge_produced_result:
             fallback_parts = [f"<{name}>" for name in image_store]
-            yield {"type": "result", "content": "\n\n".join(fallback_parts), "images": image_store}
+            figures = dict(image_store)
+            yield {"type": "result", "content": "\n\n".join(fallback_parts), "figures": figures}
 
     yield {"type": "done"}

backend/main.py

@@ -174,10 +174,13 @@ SANDBOX_TIMEOUT = 300
 # Structure: {tab_id: [messages...]}
 CONVERSATION_HISTORY: Dict[str, List[Dict]] = {}
 
-# Image stores per tab (persistent across requests so re-entry works without multimodal)
-# Structure: {tab_id: {image_name: base64_png, ...}}
-IMAGE_STORES: Dict[str, Dict[str, str]] = {}
-IMAGE_COUNTERS: Dict[str, int] = {}
+# Figure stores per tab (persistent across requests so re-entry works without multimodal)
+# Structure: {tab_id: {figure_name: base64_png, ...}}
+# Global figure store: all agents write here so cross-agent references work.
+# Keys are namespaced like "figure_T{tab}_{N}" so there are no collisions.
+FIGURE_STORE: Dict[str, dict] = {}
+# Per-tab counters to track the next figure number for each tab
+FIGURE_COUNTERS: Dict[str, int] = {}
 
 # Multi-user isolation
 MULTI_USER = False
@@ -408,13 +411,28 @@ async def _stream_code_agent_inner(messages, endpoint, token, model, e2b_key, se
         system_prompt = get_system_prompt("code", frontend_context)
         full_messages = [{"role": "system", "content": system_prompt}] + messages
 
+        # Ensure per-tab counter exists
+        if tab_id not in FIGURE_COUNTERS:
+            FIGURE_COUNTERS[tab_id] = 0
+
         async for chunk in _stream_sync_generator(
             stream_code_execution, client, model, full_messages, sbx,
             files_root=files_root or FILES_ROOT, extra_params=extra_params,
-            abort_event=abort_event, multimodal=multimodal, tab_id=tab_id
+            abort_event=abort_event, multimodal=multimodal, tab_id=tab_id,
+            figure_store=FIGURE_STORE,
         ):
             yield chunk
 
+        # Derive counter from store keys for this tab's prefix
+        prefix = f"figure_T{tab_id}_"
+        max_counter = 0
+        for name in FIGURE_STORE:
+            if name.startswith(prefix):
+                m = re.search(r'_(\d+)$', name)
+                if m:
+                    max_counter = max(max_counter, int(m.group(1)))
+        FIGURE_COUNTERS[tab_id] = max_counter
+
     except Exception as e:
         import traceback
         error_message = f"Code execution error: {str(e)}\n{traceback.format_exc()}"
@@ -442,7 +460,8 @@ async def _stream_code_agent_inner(messages, endpoint, token, model, e2b_key, se
                 async for chunk in _stream_sync_generator(
                     stream_code_execution, client, model, full_messages, sbx,
                     files_root=files_root or FILES_ROOT, extra_params=extra_params,
-                    abort_event=abort_event, multimodal=multimodal, tab_id=tab_id
+                    abort_event=abort_event, multimodal=multimodal, tab_id=tab_id,
+                    figure_store=FIGURE_STORE,
                 ):
                     yield chunk
 
@@ -646,11 +665,9 @@ async def _stream_image_agent_inner(messages, endpoint, token, model, hf_token,
         yield f"data: {json.dumps({'type': 'error', 'content': 'HuggingFace token required for image generation. Please configure in settings or set HF_TOKEN environment variable.'})}\n\n"
         return
 
-    # Get or create persistent image store for this tab
-    if tab_id not in IMAGE_STORES:
-        IMAGE_STORES[tab_id] = {}
-    if tab_id not in IMAGE_COUNTERS:
-        IMAGE_COUNTERS[tab_id] = 0
+    # Ensure per-tab counter exists
+    if tab_id not in FIGURE_COUNTERS:
+        FIGURE_COUNTERS[tab_id] = 0
 
     try:
         client = OpenAI(base_url=endpoint, api_key=token)
@@ -663,18 +680,20 @@ async def _stream_image_agent_inner(messages, endpoint, token, model, hf_token,
             extra_params=extra_params, abort_event=abort_event,
             files_root=files_root, multimodal=multimodal,
             tab_id=tab_id,
-            image_store=IMAGE_STORES[tab_id],
-            image_counter=IMAGE_COUNTERS[tab_id],
+            image_store=FIGURE_STORE,
+            image_counter=FIGURE_COUNTERS[tab_id],
         ):
             yield chunk
 
-        # Derive counter from store keys (each image_T{id}_{N} has a number)
+        # Derive counter from store keys for this tab's prefix
+        prefix = f"figure_T{tab_id}_"
         max_counter = 0
-        for name in IMAGE_STORES[tab_id]:
-            m = re.search(r'_(\d+)$', name)
-            if m:
-                max_counter = max(max_counter, int(m.group(1)))
-        IMAGE_COUNTERS[tab_id] = max_counter
+        for name in FIGURE_STORE:
+            if name.startswith(prefix):
+                m = re.search(r'_(\d+)$', name)
+                if m:
+                    max_counter = max(max_counter, int(m.group(1)))
+        FIGURE_COUNTERS[tab_id] = max_counter
 
     except Exception as e:
         import traceback
@@ -1320,14 +1339,15 @@ def select_session(session_name: str, user_id: str = '') -> bool:
         keys_to_remove = [k for k in CONVERSATION_HISTORY if k.startswith(prefix)]
         for k in keys_to_remove:
             del CONVERSATION_HISTORY[k]
-        for k in [k for k in IMAGE_STORES if k.startswith(prefix)]:
-            del IMAGE_STORES[k]
-        for k in [k for k in IMAGE_COUNTERS if k.startswith(prefix)]:
-            del IMAGE_COUNTERS[k]
+        # Clear figure store entries belonging to this user's tabs
+        for k in [k for k in FIGURE_STORE if k.startswith(f"figure_T{prefix}")]:
+            del FIGURE_STORE[k]
+        for k in [k for k in FIGURE_COUNTERS if k.startswith(prefix)]:
+            del FIGURE_COUNTERS[k]
     else:
         CONVERSATION_HISTORY.clear()
-        IMAGE_STORES.clear()
-        IMAGE_COUNTERS.clear()
+        FIGURE_STORE.clear()
+        FIGURE_COUNTERS.clear()
 
     return True
 

frontend/streaming.js

@@ -173,7 +173,7 @@ async function streamChatResponse(messages, chatContainer, agentType, tabId) {
                         // Still generating - no action needed
 
                     } else if (data.type === 'result') {
-                        // References are already globally namespaced by the backend (e.g., figure_T3_1, image_T3_1)
+                        // References are globally namespaced by the backend (e.g., figure_T3_1)
                         const resultText = data.content || '';
 
                         // Populate global registry
@@ -184,16 +184,9 @@ async function streamChatResponse(messages, chatContainer, agentType, tabId) {
                                 }
                             }
                         }
-                        if (data.images) {
-                            for (const [name, imgBase64] of Object.entries(data.images)) {
-                                if (new RegExp(`</?${name}>`, 'i').test(resultText)) {
-                                    globalFigureRegistry[name] = { type: 'png', data: imgBase64 };
-                                }
-                            }
-                        }
 
                         // Agent result - update command center widget
-                        updateActionWidgetWithResult(tabId, resultText, data.figures || {}, data.images || {});
+                        updateActionWidgetWithResult(tabId, resultText, data.figures || {});
 
                     } else if (data.type === 'result_preview') {
                         // Show result preview
@@ -219,19 +212,6 @@ async function streamChatResponse(messages, chatContainer, agentType, tabId) {
                             }
                         }
 
-                        // Handle <image_N> references from image agent
-                        if (data.images) {
-                            for (const [imageName, imageBase64] of Object.entries(data.images)) {
-                                const placeholderId = `%%%IMAGE_${imageName}%%%`;
-                                figurePlaceholders[placeholderId] = { type: 'png', data: imageBase64, isGenerated: true };
-
-                                const pairedTag = new RegExp(`<${imageName}></${imageName}>`, 'gi');
-                                previewContent = previewContent.replace(pairedTag, `\n\n${placeholderId}\n\n`);
-                                const singleTag = new RegExp(`</?${imageName}>`, 'gi');
-                                previewContent = previewContent.replace(singleTag, `\n\n${placeholderId}\n\n`);
-                            }
-                        }
-
                         // Process markdown
                         let html = parseMarkdown(previewContent);
 
@@ -647,7 +627,7 @@ async function streamChatResponse(messages, chatContainer, agentType, tabId) {
                         scrollChatToBottom(chatContainer);
 
                         // Propagate error to parent action widget
-                        updateActionWidgetWithResult(tabId, `Error: ${data.content}`, {}, {});
+                        updateActionWidgetWithResult(tabId, `Error: ${data.content}`, {});
                         const errorWidget = actionWidgets[tabId];
                         if (errorWidget) {
                             const doneIndicator = errorWidget.querySelector('.done-indicator');
@@ -679,7 +659,7 @@ async function streamChatResponse(messages, chatContainer, agentType, tabId) {
             chatContainer.appendChild(resultDiv);
 
             // Send abort result to parent action widget (so command center knows it was aborted)
-            updateActionWidgetWithResult(tabId, abortResultText, {}, {});
+            updateActionWidgetWithResult(tabId, abortResultText, {});
 
             // Override the done indicator to show × instead of ✓
             const widget = actionWidgets[tabId];
@@ -946,7 +926,7 @@ function showActionWidget(chatContainer, action, message, targetTabId, taskId =
     actionWidgets[targetTabId] = widget;
 }
 
-async function updateActionWidgetWithResult(tabId, resultContent, figures, images) {
+async function updateActionWidgetWithResult(tabId, resultContent, figures) {
     const widget = actionWidgets[tabId];
     if (!widget) return;
 
@@ -980,19 +960,6 @@ async function updateActionWidgetWithResult(tabId, resultContent, figures, image
         }
     }
 
-    // Handle <image_N> references from image agent
-    if (images) {
-        for (const [imageName, imageBase64] of Object.entries(images)) {
-            const placeholderId = `%%%IMAGE_${imageName}%%%`;
-            figurePlaceholders[placeholderId] = { type: 'png', data: imageBase64 };
-
-            const pairedTag = new RegExp(`<${imageName}></${imageName}>`, 'gi');
-            processedContent = processedContent.replace(pairedTag, `\n\n${placeholderId}\n\n`);
-            const singleTag = new RegExp(`</?${imageName}>`, 'gi');
-            processedContent = processedContent.replace(singleTag, `\n\n${placeholderId}\n\n`);
-        }
-    }
-
     // Process markdown
     let html = parseMarkdown(processedContent);
 

frontend/workspace.js

@@ -48,6 +48,34 @@ function restoreWorkspace(workspace) {
         }
     }
 
+    // Restore globalFigureRegistry from saved code-cell images
+    // so that <figure_T1_1> tags in results resolve after reload
+    for (const tabData of tabs) {
+        for (const msg of (tabData.messages || [])) {
+            if (msg.type === 'code-cell' && msg.images) {
+                for (const img of msg.images) {
+                    if (img.name && img.src) {
+                        // Parse data URL: "data:image/png;base64,..." -> {type: "png", data: "..."}
+                        const m = img.src.match(/^data:image\/(\w+);base64,(.+)$/);
+                        if (m) {
+                            globalFigureRegistry[img.name] = { type: m[1], data: m[2] };
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    // Re-resolve any figure refs in rendered HTML now that the registry is populated
+    if (Object.keys(globalFigureRegistry).length > 0) {
+        document.querySelectorAll('.action-widget-result-section .section-content, .result-preview-content, .result-content').forEach(el => {
+            const resolved = resolveGlobalFigureRefs(el.innerHTML);
+            if (resolved !== el.innerHTML) {
+                el.innerHTML = resolved;
+            }
+        });
+    }
+
     // Switch to the active tab
     if (workspace.activeTabId !== undefined) {
         switchToTab(workspace.activeTabId);