snapshot working

backend/server.py

@@ -181,6 +181,18 @@ def create_app(loop: asyncio.AbstractEventLoop) -> web.Application:
 
         return web.Response(text=_dumps({"filename": filename}), content_type="application/json")
 
+    async def download_file(request: web.Request) -> web.Response:
+        """Accept a blob POST and return it with Content-Disposition: attachment."""
+        body = await request.read()
+        filename = request.query.get("filename", "workflow.png")
+        return web.Response(
+            body=body,
+            content_type="application/octet-stream",
+            headers={
+                "Content-Disposition": f'attachment; filename="{filename}"',
+            },
+        )
+
     async def submit_prompt(request: web.Request) -> web.Response:
         body = await request.json()
         prompt = body.get("prompt")
@@ -249,6 +261,7 @@ def create_app(loop: asyncio.AbstractEventLoop) -> web.Application:
     app.router.add_get("/files", list_files)
     app.router.add_get("/browse", browse_dir)
     app.router.add_post("/upload", upload_file)
+    app.router.add_post("/download", download_file)
     app.router.add_post("/prompt", submit_prompt)
     app.router.add_get("/ws", websocket_handler)
 

desktop.py

@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+import base64
 import logging
 import socket
 import threading
 import time
 import urllib.request
+from pathlib import Path
 
 import webview
 from aiohttp import web
@@ -43,6 +45,34 @@ class _Api:
             return result[0]
         return None
 
+    def save_workflow_png(self, data_url: str, default_filename: str = "workflow.png") -> str | None:
+        """Open a native save dialog, write the PNG bytes, and return the saved path."""
+        win = self._window_ref[0]
+        if win is None:
+            return None
+
+        result = win.create_file_dialog(
+            webview.SAVE_DIALOG,
+            save_filename=default_filename,
+            file_types=(
+                "PNG image (*.png)",
+                "All files (*.*)",
+            ),
+        )
+        if not result:
+            return None
+
+        path = Path(result[0] if isinstance(result, (list, tuple)) else result).expanduser()
+        if path.suffix.lower() != ".png":
+            path = path.with_suffix(".png")
+
+        _, _, encoded = data_url.partition(",")
+        if not encoded:
+            raise ValueError("Invalid data URL payload")
+
+        path.write_bytes(base64.b64decode(encoded))
+        return str(path)
+
 
 def _pick_free_port() -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:

frontend/src/App.jsx

@@ -42,6 +42,40 @@ function getOutputSlot(handleId) {
   return parseInt(handleId.split('::')[1], 10);
 }
 
+function blobToDataUrl(blob) {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+    reader.onloadend = () => resolve(reader.result);
+    reader.onerror = () => reject(reader.error || new Error('Failed to read file'));
+    reader.readAsDataURL(blob);
+  });
+}
+
+function serializeWorkflowState(nodes, edges) {
+  return {
+    version: 1,
+    nodes: nodes.map((node) => ({
+      id: node.id,
+      type: node.type || 'custom',
+      position: node.position,
+      dragHandle: node.dragHandle || '.drag-handle',
+      data: {
+        label: node.data?.label || node.data?.className || 'Node',
+        className: node.data?.className || '',
+        widgetValues: node.data?.widgetValues || {},
+      },
+    })),
+    edges: edges.map((edge) => ({
+      id: edge.id,
+      source: edge.source,
+      sourceHandle: edge.sourceHandle,
+      target: edge.target,
+      targetHandle: edge.targetHandle,
+      style: edge.style,
+    })),
+  };
+}
+
 // ── Graph serialisation → backend prompt format ───────────────────────
 
 function serializeGraph(nodes, edges) {
@@ -442,8 +476,12 @@ function Flow() {
     const defs = nodeDefsRef.current;
     const hydrated = loadedNodes.map((n) => ({
       ...n,
+      type: n.type || 'custom',
+      dragHandle: n.dragHandle || '.drag-handle',
       data: {
         ...n.data,
+        label: n.data?.label || n.data?.className || 'Node',
+        widgetValues: n.data?.widgetValues || {},
         definition: defs[n.data.className] || n.data.definition,
         previewImage: null, tableRows: null, meshData: null, overlay: null,
       },
@@ -479,11 +517,7 @@ function Flow() {
     });
     if (!blob) throw new Error('Capture returned empty');
 
-    const currentNodes = allNodes.map((n) => ({
+    const workflow = serializeWorkflowState(allNodes, reactFlow.getEdges());
-      ...n,
-      data: { ...n.data, previewImage: null, tableRows: null, meshData: null, overlay: null },
-    }));
-    const workflow = { version: 1, nodes: currentNodes, edges: reactFlow.getEdges() };
     return embedWorkflow(blob, workflow);
   }, [reactFlow]);
 
@@ -492,32 +526,62 @@ function Flow() {
     try {
       const finalBlob = await getWorkflowBlob();
 
-      if (window.showSaveFilePicker) {
+      if (window.pywebview?.api?.save_workflow_png) {
-        const handle = await window.showSaveFilePicker({
+        const dataUrl = await blobToDataUrl(finalBlob);
-          suggestedName: 'workflow.png',
+        const savedPath = await window.pywebview.api.save_workflow_png(dataUrl, 'workflow.png');
-          types: [{ description: 'PNG Image', accept: { 'image/png': ['.png'] } }],
+        if (!savedPath) {
-        });
+          setStatus({ text: 'Save cancelled.', level: 'info' });
-        const writable = await handle.createWritable();
+          return;
-        await writable.write(finalBlob);
+        }
-        await writable.close();
+        setStatus({ text: `Workflow saved to ${savedPath}.`, level: 'info' });
-      } else {
+        return;
-        // Fallback: programmatic download
-        const a = document.createElement('a');
-        a.href = URL.createObjectURL(finalBlob);
-        a.download = 'workflow.png';
-        document.body.appendChild(a);
-        a.click();
-        document.body.removeChild(a);
-        URL.revokeObjectURL(a.href);
       }
 
-      setStatus({ text: 'Workflow saved.', level: 'info' });
+      if ('showSaveFilePicker' in window) {
-    } catch (err) {
+        try {
-      if (err.name === 'AbortError') {
+          const handle = await window.showSaveFilePicker({
-        setStatus({ text: 'Save cancelled.', level: 'info' });
+            suggestedName: 'workflow.png',
-      } else {
+            types: [
-        setStatus({ text: 'Save failed: ' + err.message, level: 'error' });
+              {
+                description: 'PNG image',
+                accept: { 'image/png': ['.png'] },
+              },
+            ],
+          });
+          const writable = await handle.createWritable();
+          await writable.write(finalBlob);
+          await writable.close();
+          setStatus({ text: 'Workflow saved as workflow.png.', level: 'info' });
+          return;
+        } catch (err) {
+          if (err?.name === 'AbortError') {
+            setStatus({ text: 'Save cancelled.', level: 'info' });
+            return;
+          }
+          throw err;
+        }
       }
+
+      // Final fallback: trigger a browser download and tell the user where it went.
+      const resp = await fetch('/download?filename=workflow.png', {
+        method: 'POST',
+        body: finalBlob,
+      });
+      const dlBlob = await resp.blob();
+      const a = document.createElement('a');
+      a.href = URL.createObjectURL(dlBlob);
+      a.download = 'workflow.png';
+      document.body.appendChild(a);
+      a.click();
+      document.body.removeChild(a);
+      setTimeout(() => URL.revokeObjectURL(a.href), 1000);
+
+      setStatus({
+        text: 'Workflow downloaded as workflow.png to your browser default downloads folder.',
+        level: 'info',
+      });
+    } catch (err) {
+      setStatus({ text: 'Save failed: ' + err.message, level: 'error' });
     }
   }, [getWorkflowBlob]);
 
@@ -545,7 +609,8 @@ function Flow() {
       if (!file) return;
       try {
         let data;
-        if (file.name.endsWith('.png') || file.type === 'image/png') {
+        const lowerName = file.name.toLowerCase();
+        if (lowerName.endsWith('.png') || file.type === 'image/png') {
           data = await extractWorkflow(file);
           if (!data) {
             setStatus({ text: 'No workflow data found in image.', level: 'error' });
@@ -571,7 +636,8 @@ function Flow() {
     event.preventDefault();
 
     const file = files[0];
-    if (file.type !== 'image/png') return;
+    const lowerName = file.name.toLowerCase();
+    if (file.type !== 'image/png' && !lowerName.endsWith('.png')) return;
 
     try {
       const data = await extractWorkflow(file);

frontend/src/pngMetadata.js

@@ -1,9 +1,9 @@
 /**
- * PNG tEXt chunk utilities for embedding/extracting workflow metadata.
+ * PNG text chunk utilities for embedding/extracting workflow metadata.
  *
  * PNG files are composed of chunks: [4-byte length][4-byte type][data][4-byte CRC].
- * We add a tEXt chunk with key "workflow" containing the JSON-serialised graph,
+ * We add an iTXt chunk with key "workflow" containing the JSON-serialised graph,
- * inserted just before the IEND chunk.
+ * inserted just before the IEND chunk. We still read legacy tEXt chunks.
  */
 
 // ── CRC32 (PNG uses CRC-32/ISO 3309) ────────────────────────────────
@@ -43,10 +43,64 @@ function chunkType(data, offset) {
   );
 }
 
+function readUint32(data, offset) {
+  return new DataView(data.buffer, data.byteOffset + offset, 4).getUint32(0);
+}
+
+function buildChunk(type, payload) {
+  const encoder = new TextEncoder();
+  const typeBytes = encoder.encode(type);
+  const forCrc = new Uint8Array(4 + payload.length);
+  forCrc.set(typeBytes, 0);
+  forCrc.set(payload, 4);
+
+  const chunk = new Uint8Array(12 + payload.length);
+  const view = new DataView(chunk.buffer);
+  view.setUint32(0, payload.length);
+  chunk.set(typeBytes, 4);
+  chunk.set(payload, 8);
+  view.setUint32(8 + payload.length, crc32(forCrc));
+  return chunk;
+}
+
+function parseTextChunk(type, chunkData) {
+  const decoder = new TextDecoder();
+  const keywordEnd = chunkData.indexOf(0);
+  if (keywordEnd === -1) return null;
+
+  const keyword = decoder.decode(chunkData.subarray(0, keywordEnd));
+  if (keyword !== 'workflow') return null;
+
+  if (type === 'tEXt') {
+    return JSON.parse(decoder.decode(chunkData.subarray(keywordEnd + 1)));
+  }
+
+  if (type !== 'iTXt') return null;
+
+  const compressionFlagIdx = keywordEnd + 1;
+  const compressionMethodIdx = keywordEnd + 2;
+  if (compressionMethodIdx >= chunkData.length) return null;
+
+  const compressionFlag = chunkData[compressionFlagIdx];
+  if (compressionFlag !== 0) {
+    throw new Error('Compressed PNG workflow metadata is not supported');
+  }
+
+  let offset = compressionMethodIdx + 1;
+  const languageEnd = chunkData.indexOf(0, offset);
+  if (languageEnd === -1) return null;
+
+  offset = languageEnd + 1;
+  const translatedEnd = chunkData.indexOf(0, offset);
+  if (translatedEnd === -1) return null;
+
+  return JSON.parse(decoder.decode(chunkData.subarray(translatedEnd + 1)));
+}
+
 // ── Public API ───────────────────────────────────────────────────────
 
 /**
- * Embed a workflow object into a PNG blob as a tEXt chunk.
+ * Embed a workflow object into a PNG blob as an iTXt chunk.
  * Returns a new Blob with the metadata inserted before IEND.
  */
 export async function embedWorkflow(pngBlob, workflow) {
@@ -55,33 +109,22 @@ export async function embedWorkflow(pngBlob, workflow) {
 
   const encoder = new TextEncoder();
 
-  // Build tEXt payload: keyword \0 text
+  // Build iTXt payload:
+  // keyword \0 compression-flag compression-method language-tag \0 translated-keyword \0 text
   const key = encoder.encode('workflow');
   const val = encoder.encode(JSON.stringify(workflow));
-  const payload = new Uint8Array(key.length + 1 + val.length);
+  const payload = new Uint8Array(key.length + 5 + val.length);
   payload.set(key, 0);
-  // payload[key.length] is already 0 (null separator)
+  payload.set(val, key.length + 5);
-  payload.set(val, key.length + 1);
+  const chunk = buildChunk('iTXt', payload);
-
-  // CRC covers type + payload
-  const typeBytes = encoder.encode('tEXt');
-  const forCrc = new Uint8Array(4 + payload.length);
-  forCrc.set(typeBytes, 0);
-  forCrc.set(payload, 4);
-
-  // Assemble chunk: length(4) + type(4) + payload + crc(4)
-  const chunk = new Uint8Array(12 + payload.length);
-  const view = new DataView(chunk.buffer);
-  view.setUint32(0, payload.length);
-  chunk.set(typeBytes, 4);
-  chunk.set(payload, 8);
-  view.setUint32(8 + payload.length, crc32(forCrc));
 
   // Locate IEND
   let pos = 8;
   let iendPos = data.length;
   while (pos < data.length) {
-    const len = new DataView(data.buffer, pos, 4).getUint32(0);
+    if (pos + 8 > data.length) break;
+    const len = readUint32(data, pos);
+    if (pos + 12 + len > data.length) break;
     if (chunkType(data, pos) === 'IEND') { iendPos = pos; break; }
     pos += 12 + len;
   }
@@ -96,34 +139,30 @@ export async function embedWorkflow(pngBlob, workflow) {
 }
 
 /**
- * Extract the workflow object from a PNG blob's tEXt chunks.
+ * Extract the workflow object from a PNG blob's iTXt/tEXt chunks.
  * Returns the parsed object, or null if no "workflow" key is found.
  */
 export async function extractWorkflow(pngBlob) {
   const data = new Uint8Array(await pngBlob.arrayBuffer());
   if (!isPng(data)) return null;
 
-  const decoder = new TextDecoder();
   let pos = 8;
+  let found = null;
 
   while (pos + 8 <= data.length) {
-    const len = new DataView(data.buffer, pos, 4).getUint32(0);
+    const len = readUint32(data, pos);
+    if (pos + 12 + len > data.length) break;
     const type = chunkType(data, pos);
 
-    if (type === 'tEXt' && pos + 8 + len <= data.length) {
+    if (type === 'tEXt' || type === 'iTXt') {
       const chunkData = data.subarray(pos + 8, pos + 8 + len);
-      const nullIdx = chunkData.indexOf(0);
+      const parsed = parseTextChunk(type, chunkData);
-      if (nullIdx !== -1) {
+      if (parsed) found = parsed;
-        const k = decoder.decode(chunkData.subarray(0, nullIdx));
-        if (k === 'workflow') {
-          return JSON.parse(decoder.decode(chunkData.subarray(nullIdx + 1)));
-        }
-      }
     }
 
     if (type === 'IEND') break;
     pos += 12 + len;
   }
 
-  return null;
+  return found;
 }

frontend/vite.config.js

@@ -10,6 +10,7 @@ export default defineConfig({
       '/files': 'http://127.0.0.1:8188',
       '/browse': 'http://127.0.0.1:8188',
       '/upload': 'http://127.0.0.1:8188',
+      '/download': 'http://127.0.0.1:8188',
       '/prompt': 'http://127.0.0.1:8188',
       '/ws': {
         target: 'http://127.0.0.1:8188',