Ethanfel/Comfyui-Workflow-Snapshot-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Audit fixes: data-loss, security, performance, UX + new features

Browse Source 
Comprehensive audit pass across the JS frontend and Python backend.

Bugs / correctness:
- Swap & restore now pre-save current state (hash-deduped) so unsaved
  edits aren't lost when swapping/restoring, incl. rapid double-swap
- Unify captureSnapshot/captureNodeSnapshot into _captureCore; node
  captures now update the dedup hash (no duplicate auto-snapshot after)
- Cycle guard in getDisplayPath; Ctrl+S ignores text fields and the
  other-workflow view; tolerant API error parsing; prompt default pre-fill

Security / robustness (backend):
- Validate workflowKey against path traversal (reject ./.. + containment)
- Generic 500 messages (no exception-string leak), logged server-side
- Request body-size cap + migrate record cap
- Atomic writes (temp file + os.replace) on all write paths

Performance / memory:
- /list omits base64 thumbnails (hasThumbnail flag, lazy-loaded client-side)
- LRU-bounded previous-graph cache; persistent (prune+LRU) SVG cache
- Incremental in-place updates for lock/note instead of full list rebuild

UX / docs:
- Busy-op feedback, named-delete confirm, relative timestamps
- README: remove disabled branching feature, fix version badge & storage paths

Features:
- Export / Import snapshots (export route + reuse migrate)
- Storage-usage display (usage route + footer label)
- Pause auto-capture toggle
- Age-based retention (maxAgeDays setting + prune param)

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Ethanfel
2026-06-14 11:04:09 +02:00
parent 8c84d2ff4e
commit 0d1415fca4
4 changed files with 545 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+17 -39
View File
@@ -5,7 +5,7 @@
<p align="center">
<a href="https://registry.comfy.org/publishers/ethanfel/nodes/comfyui-snapshot-manager"><img src="https://img.shields.io/badge/ComfyUI-Registry-blue?logo=data:image/svg%2bxml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCAyNCAyNCI+PHBhdGggZD0iTTEyIDJMMyA3djEwbDkgNSA5LTVWN2wtOS01eiIgZmlsbD0id2hpdGUiLz48L3N2Zz4=" alt="ComfyUI Registry"/></a>
<a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-green" alt="MIT License"/></a>
<img src="https://img.shields.io/badge/version-3.0.0-blue" alt="Version"/>
<img src="https://img.shields.io/badge/version-3.0.1-blue" alt="Version"/>
<img src="https://img.shields.io/badge/ComfyUI-Extension-purple" alt="ComfyUI Extension"/>
</p>
@@ -20,8 +20,9 @@
## Features
- **Auto-capture** — Snapshots are saved automatically as you edit, with configurable debounce
- **Custom naming** — Name your snapshots when taking them manually ("Before merge", "Working v2", etc.)
- **Search & filter** — Quickly find snapshots by name with the filter bar
- **Custom naming** — Name your snapshots when taking them manually ("Before merge", "Working v2", etc.); double-click a snapshot's name in the sidebar to rename it later
- **Notes** — Attach a freeform note to any snapshot with the note (pencil) button; notes are searchable from the filter bar
- **Search & filter** — Quickly find snapshots by name or note with the filter bar
- **Restore or Swap** — Open a snapshot as a new workflow, or replace the current one in-place
- **Workflow browser** — Browse and recover snapshots from any workflow, including renamed or deleted ones
- **Per-workflow storage** — Each workflow has its own independent snapshot history
@@ -35,9 +36,12 @@
- **Ctrl+S shortcut** — Press Ctrl+S (or Cmd+S on Mac) to take a manual snapshot alongside ComfyUI's own save
- **SVG graph previews** — Hover any snapshot for a tooltip preview of the workflow graph; click the eye button for a full-size modal; diff view now shows side-by-side SVG comparison with color-coded highlights (green = added, red = removed, amber = modified)
- **Diff view** — Compare any snapshot against the current workflow (one click) or two snapshots against each other (Shift+click to set base); see added/removed/modified nodes, widget value changes, and rewired connections in a single modal
- **Snapshot branching** — Swap to an old snapshot and edit to fork into a new branch; navigate between branches with `< 1/3 >` arrows at fork points in the sidebar and timeline, like ChatGPT conversation branching
- **Profile manager** — Save and load named sets of workflows as session profiles (like browser tab groups); profiles track which workflows you visited and restore the latest snapshot for each
- **Hide auto-saves** — Toggle button next to the search bar hides auto-save snapshots to reduce clutter while keeping manual, locked, and node-triggered snapshots visible
- **Pause auto-capture** — Toggle automatic capture on/off for the session without leaving the panel (Auto: On/Off button in the filter row)
- **Export / Import** — Download a workflow's full snapshot history as a JSON file and re-import it on another machine or server
- **Storage usage** — The sidebar footer shows total snapshot storage used on the server across all workflows
- **Retention by age** — Optionally auto-delete snapshots older than a configurable number of days (off by default), alongside the per-workflow count limit
- **Lock/pin snapshots** — Protect important snapshots from auto-pruning and "Clear All" with a single click
- **Concurrency-safe** — Lock guard prevents double-click issues during restore
- **Server-side storage** — Snapshots persist on the ComfyUI server's filesystem, accessible from any browser
@@ -178,32 +182,13 @@ Visually inspect any snapshot without restoring or swapping it.
The SVG renderer draws nodes with their stored position, size, and colors. Links are rendered as bezier curves colored by type (blue for IMAGE, orange for CLIP, purple for MODEL, etc.). Collapsed nodes appear as thin title-only strips. Thumbnails (hover tooltips) auto-simplify by hiding labels and slot dots for clarity at small sizes.
### 14. Snapshot Branching
Branching lets you explore multiple variations of a workflow without losing any history — similar to conversation branching in ChatGPT.
**How it works:**
1. Work normally — snapshots chain linearly as you edit
2. **Swap** to an older snapshot and start editing — the next auto-capture forks into a new branch from that point
3. A **`< 1/2 >`** branch navigator appears at every fork point in the sidebar and the timeline
4. Click the arrows to switch between branches — the sidebar and timeline update together
**Details:**
- Each snapshot stores a `parentId` linking it to its predecessor, forming a tree
- Legacy snapshots (from before branching) are automatically chained by timestamp for backwards compatibility
- **Pruning is branch-safe** — ancestors of the current branch tip and fork-point snapshots are never auto-pruned
- **Deleting a fork point** re-parents its children to the deleted snapshot's parent (with a confirmation dialog)
- Switching workflows clears branch navigation state
### 15. Hide Auto-saves
### 14. Hide Auto-saves
Click the **Hide Auto** button next to the search bar to hide all auto-save snapshots ("Auto" and "Initial"). The button turns blue and switches to **Show Auto** when active.
The filter works together with text search — both are applied simultaneously. Branch navigators remain visible regardless of the filter. Manual, locked, node-triggered, and "Current" snapshots are always shown.
The filter works together with text search — both are applied simultaneously. Manual, locked, node-triggered, and "Current" snapshots are always shown.
### 16. Session Profiles
### 15. Session Profiles
Save and load named sets of workflows — like browser tab groups for ComfyUI.
@@ -251,7 +236,7 @@ All settings are available in **ComfyUI Settings > Snapshot Manager**:
2. A **debounce timer** prevents excessive writes
3. The workflow is serialized and **hash-checked** against the last capture (per-workflow) to avoid duplicates
4. The previous graph state is diffed against the current to **detect the change type** (node add/remove, connection, parameter, move, or mixed) — stored as a `changeType` field on the record
5. New snapshots are sent to the **server** and stored as individual JSON files under `data/snapshots/`
5. New snapshots are sent to the **server** and stored as individual JSON files under `<user_dir>/snapshot_manager/snapshots/`
6. The **sidebar panel** and **timeline bar** fetch snapshots from the server and render them with change-type icons
7. **Restore/Swap** loads graph data back into ComfyUI with a lock guard to prevent concurrent operations, and updates the graph cache so the next diff is accurate
@@ -270,26 +255,25 @@ All settings are available in **ComfyUI Settings > Snapshot Manager**:
4. The **timeline** updates: the swapped-to snapshot gets a white ring (active), the auto-saved snapshot gets a green dot (current)
5. Clicking the green dot swaps back; editing the graph clears both markers (the next auto-capture supersedes them)
**Branching:**
**Branching** (implemented but disabled in the current release — the code remains for a future enable):
1. Each snapshot stores a `parentId` pointing to its predecessor
2. `buildSnapshotTree()` constructs parent/child maps from all records — legacy snapshots (no `parentId`) are chained by timestamp automatically
3. `getDisplayPath()` walks the tree from root to tip, following `activeBranchSelections` at each fork point, producing the linear branch view
4. The sidebar and timeline render only the current branch; `< 1/3 >` navigators at fork points switch between branches
5. **Pruning protection**: before pruning, `getAncestorIds()` collects all ancestors of the branch tip; these IDs plus fork-point IDs are sent as `protectedIds` to the server
4. **Pruning protection**: before pruning, `getAncestorIds()` collects all ancestors of the branch tip; these IDs plus fork-point IDs are sent as `protectedIds` to the server
**Profiles:**
1. Session tracking records each visited workflow key with timestamps
2. **Save** creates a JSON file at `<extension_dir>/data/profiles/<id>.json` with the workflow list and active workflow
2. **Save** creates a JSON file at `<user_dir>/snapshot_manager/profiles/<id>.json` with the workflow list and active workflow
3. **Load** fetches the latest snapshot for each workflow in the profile and calls `loadGraphData`
**Storage:** Snapshots are stored as JSON files on the server at `<extension_dir>/data/snapshots/<workflow_key>/<id>.json`. Profiles are stored at `<extension_dir>/data/profiles/<id>.json`. Both persist across browser sessions, ComfyUI restarts, and are accessible from any browser connecting to the same server.
**Storage:** Snapshots are stored as JSON files on the server in ComfyUI's user directory at `<user_dir>/snapshot_manager/snapshots/<workflow_key>/<id>.json`. Profiles are stored at `<user_dir>/snapshot_manager/profiles/<id>.json`. Data from older versions (kept under the extension's own `data/` folder) is migrated here automatically on first load. Both persist across browser sessions, ComfyUI restarts, and are accessible from any browser connecting to the same server.
## FAQ
**Where are snapshots stored?**
On the server's filesystem under `<extension_dir>/data/snapshots/`. Each workflow gets its own directory, and each snapshot is an individual JSON file. They persist across browser sessions and are accessible from any browser connecting to the same ComfyUI server.
On the server's filesystem in ComfyUI's user directory under `snapshot_manager/snapshots/`. Each workflow gets its own directory, and each snapshot is an individual JSON file. They persist across browser sessions and are accessible from any browser connecting to the same ComfyUI server.
**I'm upgrading from v1.x — what happens to my existing snapshots?**
On first load after upgrading, the extension automatically migrates all snapshots from your browser's IndexedDB to the server. Once migration succeeds, the old IndexedDB database is deleted. If migration fails (e.g., server unreachable), your old data is preserved and migration will retry on the next load.
@@ -303,12 +287,6 @@ Each workflow has its own snapshot history. Switching workflows cancels any pend
**I renamed/deleted a workflow — are my snapshots gone?**
No. Snapshots are keyed by the workflow name at capture time. Use the workflow picker to find and restore them under the old name.
**How does branching work?**
When you swap to an old snapshot and then edit, the next capture forks into a new branch. A `< 1/2 >` navigator appears at the fork point — click the arrows to switch branches. The tree structure is computed from `parentId` links on each snapshot. Old snapshots without `parentId` (from before v3.0) are automatically chained by timestamp.
**Can I delete a fork-point snapshot?**
Yes. The extension re-parents its children to the deleted snapshot's parent, preserving the branch structure. A confirmation dialog warns you first.
**What are profiles?**
Profiles save a list of workflows you've visited in a session. Loading a profile restores the latest snapshot for each workflow. They're useful for switching between project contexts — like browser tab groups.
js/snapshot_manager.js
+317 -127
View File
@@ -23,6 +23,7 @@ let debounceMs = 3000;
let autoCaptureEnabled = true;
let captureOnLoad = true;
let maxNodeSnapshots = 5;
let maxAgeDays = 0; // 0 = never age-prune
let showTimeline = false;
const BRANCHING_ENABLED = false;
let branchingDefault = true; // updated by ComfyUI settings onChange
@@ -56,6 +57,18 @@ const sessionWorkflows = new Map(); // workflowKey -> { firstSeen, lastSeen }
// ─── Server API Layer ───────────────────────────────────────────────
async function _respError(resp) {
// Build an Error from a non-OK response, tolerating non-JSON bodies (proxy
// 502s, HTML error pages) that would otherwise throw on resp.json() and mask
// the real status.
let detail = resp.statusText || `HTTP ${resp.status}`;
try {
const err = await resp.json();
if (err && err.error) detail = err.error;
} catch {}
return new Error(detail);
}
async function db_put(record) {
try {
const resp = await api.fetchApi("/snapshot-manager/save", {
@@ -64,8 +77,7 @@ async function db_put(record) {
body: JSON.stringify({ record }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Save failed:`, err);
@@ -82,8 +94,7 @@ async function db_getAllForWorkflow(workflowKey) {
body: JSON.stringify({ workflowKey }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
return await resp.json();
} catch (err) {
@@ -101,8 +112,7 @@ async function db_delete(workflowKey, id) {
body: JSON.stringify({ workflowKey, id }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Delete failed:`, err);
@@ -118,8 +128,7 @@ async function db_deleteAllForWorkflow(workflowKey) {
body: JSON.stringify({ workflowKey }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
return await resp.json();
} catch (err) {
@@ -133,8 +142,7 @@ async function db_getAllWorkflowKeys() {
try {
const resp = await api.fetchApi("/snapshot-manager/workflows");
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
return await resp.json();
} catch (err) {
@@ -151,8 +159,7 @@ async function db_updateMeta(workflowKey, id, fields) {
body: JSON.stringify({ workflowKey, id, fields }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Update meta failed:`, err);
@@ -175,16 +182,46 @@ async function db_getFullRecord(workflowKey, id) {
}
}
async function db_getStorageUsage() {
try {
const resp = await api.fetchApi("/snapshot-manager/usage");
if (!resp.ok) throw await _respError(resp);
return await resp.json();
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Usage scan failed:`, err);
return null;
}
}
async function db_exportWorkflow(workflowKey) {
const resp = await api.fetchApi("/snapshot-manager/export", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ workflowKey }),
});
if (!resp.ok) throw await _respError(resp);
return await resp.json();
}
async function db_importRecords(records) {
const resp = await api.fetchApi("/snapshot-manager/migrate", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ records }),
});
if (!resp.ok) throw await _respError(resp);
return await resp.json();
}
async function pruneSnapshots(workflowKey, protectedIds = []) {
try {
const resp = await api.fetchApi("/snapshot-manager/prune", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ workflowKey, maxSnapshots, source: "regular", protectedIds }),
body: JSON.stringify({ workflowKey, maxSnapshots, source: "regular", protectedIds, maxAgeDays }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Prune failed:`, err);
@@ -196,11 +233,10 @@ async function pruneNodeSnapshots(workflowKey, protectedIds = []) {
const resp = await api.fetchApi("/snapshot-manager/prune", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ workflowKey, maxSnapshots: maxNodeSnapshots, source: "node", protectedIds }),
body: JSON.stringify({ workflowKey, maxSnapshots: maxNodeSnapshots, source: "node", protectedIds, maxAgeDays }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Node prune failed:`, err);
@@ -217,8 +253,7 @@ async function profile_save(profile) {
body: JSON.stringify({ profile }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Profile save failed:`, err);
@@ -231,8 +266,7 @@ async function profile_list() {
try {
const resp = await api.fetchApi("/snapshot-manager/profile/list");
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
return await resp.json();
} catch (err) {
@@ -249,8 +283,7 @@ async function profile_delete(profileId) {
body: JSON.stringify({ id: profileId }),
});
if (!resp.ok) {
const err = await resp.json();
throw new Error(err.error || resp.statusText);
throw await _respError(resp);
}
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Profile delete failed:`, err);
@@ -335,6 +368,51 @@ function quickHash(str) {
return hash;
}
// Bound the in-memory previous-graph copies (used for change-type detection and
// diff summaries). One full graph per workflow can be large, so keep an LRU of
// the few most-recently-touched workflows instead of growing without limit.
const MAX_GRAPH_CACHE = 5;
function setLastGraphData(workflowKey, graphData) {
lastGraphDataMap.delete(workflowKey); // re-insert to mark most-recently-used
lastGraphDataMap.set(workflowKey, graphData);
while (lastGraphDataMap.size > MAX_GRAPH_CACHE) {
const oldest = lastGraphDataMap.keys().next().value;
lastGraphDataMap.delete(oldest);
}
}
// SVG previews are immutable per snapshot, so the cache can persist across
// refreshes. Drop entries for snapshots that no longer exist and cap the size,
// instead of clearing the whole cache on every refresh.
const MAX_SVG_CACHE = 80;
function pruneSvgCache(records) {
const liveIds = new Set(records.map((r) => r.id));
for (const key of [...svgCache.keys()]) {
const id = key.slice(0, key.lastIndexOf(":"));
if (!liveIds.has(id)) svgCache.delete(key);
}
while (svgCache.size > MAX_SVG_CACHE) {
const oldest = svgCache.keys().next().value;
svgCache.delete(oldest);
}
}
// Thumbnails are no longer embedded in /list metadata (only a hasThumbnail
// flag), so lazy-load and cache the base64 image per snapshot on first use.
const thumbnailCache = new Map(); // snapshotId -> dataURL string or null
async function getThumbnailDataUrl(rec) {
if (rec.thumbnail) return `data:image/jpeg;base64,${rec.thumbnail}`;
if (!rec.hasThumbnail) return null;
if (thumbnailCache.has(rec.id)) return thumbnailCache.get(rec.id);
let url = null;
try {
const full = await db_getFullRecord(rec.workflowKey, rec.id);
if (full && full.thumbnail) url = `data:image/jpeg;base64,${full.thumbnail}`;
} catch {}
thumbnailCache.set(rec.id, url);
return url;
}
function getWorkflowKey() {
try {
const wf = app.extensionManager?.workflow?.activeWorkflow;
@@ -1007,11 +1085,14 @@ function getDisplayPath(tree, branchSelections) {
if (!current) return [];
const path = [current];
const visited = new Set([current.id]);
while (true) {
const children = childrenOf.get(current.id);
if (!children || children.length === 0) break;
const selectedIndex = branchSelections.get(current.id) ?? 0;
current = children[Math.min(selectedIndex, children.length - 1)];
if (visited.has(current.id)) break; // safety: cycle detection (parity with getAllBranches)
visited.add(current.id);
path.push(current);
}
return path;
@@ -1087,7 +1168,10 @@ function selectBranchContaining(snapshotId, tree) {
// ─── Restore Lock ───────────────────────────────────────────────────
async function withRestoreLock(fn) {
if (restoreLock) return;
if (restoreLock) {
showToast("Please wait for the current operation to finish", "warn");
return;
}
let resolve;
restoreLock = new Promise((r) => { resolve = r; });
try {
@@ -1135,6 +1219,7 @@ async function showPromptDialog(message, defaultValue = "Manual") {
const result = await app.extensionManager.dialog.prompt({
title: "Snapshot Name",
message,
defaultValue,
});
return result;
} catch {
@@ -1483,10 +1568,10 @@ async function captureSnapshot(label = "Auto") {
if (restoreLock) return false;
if (captureInProgress) return false;
captureInProgress = true;
try { return await _captureSnapshotInner(label); } finally { captureInProgress = false; }
try { return await _captureCore({ label, dedupe: true, skipMove: true }); } finally { captureInProgress = false; }
}
async function _captureSnapshotInner(label) {
async function _captureCore({ label, source = null, thumbnail = null, dedupe = false, skipMove = false }) {
const graphData = getGraphData();
if (!graphData) return false;
@@ -1497,11 +1582,11 @@ async function _captureSnapshotInner(label) {
const workflowKey = getWorkflowKey();
const serialized = JSON.stringify(graphData);
const hash = quickHash(serialized);
if (hash === lastCapturedHashMap.get(workflowKey)) return false;
if (dedupe && hash === lastCapturedHashMap.get(workflowKey)) return false;
const prevGraph = lastGraphDataMap.get(workflowKey);
const changeType = detectChangeType(prevGraph, graphData);
if (changeType === "move") return false;
if (skipMove && changeType === "move") return false;
// Determine parentId for branching
let parentId = null;
@@ -1524,9 +1609,12 @@ async function _captureSnapshotInner(label) {
locked: false,
changeType,
parentId,
...(source ? { source } : {}),
...(captureDiff ? { captureDiff } : {}),
...(thumbnail ? { thumbnail } : {}),
};
const pruneFn = source === "node" ? pruneNodeSnapshots : pruneSnapshots;
try {
await db_put(record);
if (isBranchingEnabled(workflowKey)) {
@@ -1547,16 +1635,16 @@ async function _captureSnapshotInner(label) {
}
}
protectedIds.add(record.id); // protect the just-captured snapshot
await pruneSnapshots(workflowKey, [...protectedIds]);
await pruneFn(workflowKey, [...protectedIds]);
} else {
await pruneSnapshots(workflowKey);
await pruneFn(workflowKey);
}
} catch {
return false;
}
lastCapturedHashMap.set(workflowKey, hash);
lastGraphDataMap.set(workflowKey, graphData);
setLastGraphData(workflowKey, graphData);
lastCapturedIdMap.set(workflowKey, record.id);
pickerDirty = true;
currentSnapshotId = null; // new capture supersedes "current" bookmark
@@ -1573,82 +1661,11 @@ async function _captureSnapshotInner(label) {
async function captureNodeSnapshot(label = "Node Trigger", thumbnail = null) {
if (restoreLock) return false;
const graphData = getGraphData();
if (!graphData) return false;
const nodes = graphData.nodes || [];
if (nodes.length === 0) return false;
const workflowKey = getWorkflowKey();
const prevGraph = lastGraphDataMap.get(workflowKey);
const changeType = detectChangeType(prevGraph, graphData);
// Determine parentId for branching
let parentId = null;
if (isBranchingEnabled(workflowKey)) {
if (activeSnapshotId) {
parentId = activeSnapshotId;
} else if (lastCapturedIdMap.has(workflowKey)) {
parentId = lastCapturedIdMap.get(workflowKey);
}
}
const captureDiff = computeCaptureMetaDiff(prevGraph, graphData);
const record = {
id: generateId(),
workflowKey,
timestamp: Date.now(),
label,
nodeCount: nodes.length,
graphData,
locked: false,
source: "node",
changeType,
parentId,
...(captureDiff ? { captureDiff } : {}),
...(thumbnail ? { thumbnail } : {}),
};
try {
await db_put(record);
if (isBranchingEnabled(workflowKey)) {
// Compute protected IDs: ancestors + fork points + ancestors of locked snapshots
const allRecs = await db_getAllForWorkflow(workflowKey);
const tempTree = buildSnapshotTree(allRecs);
const protectedNodeIds = getAncestorIds(record.id, tempTree.parentOf);
for (const [pid, children] of tempTree.childrenOf) {
if (children.length > 1) protectedNodeIds.add(pid);
}
for (const rec of allRecs) {
if (rec.locked) {
for (const aid of getAncestorIds(rec.id, tempTree.parentOf)) {
protectedNodeIds.add(aid);
}
}
}
protectedNodeIds.add(record.id);
await pruneNodeSnapshots(workflowKey, [...protectedNodeIds]);
} else {
await pruneNodeSnapshots(workflowKey);
}
} catch {
return false;
}
lastGraphDataMap.set(workflowKey, graphData);
lastCapturedIdMap.set(workflowKey, record.id);
pickerDirty = true;
currentSnapshotId = null;
activeSnapshotId = null;
if (sidebarRefresh) {
sidebarRefresh().catch(() => {});
}
if (timelineRefresh) {
timelineRefresh().catch(() => {});
}
return true;
// Node-triggered captures are never deduped or move-skipped (an explicit
// node trigger should always produce a snapshot), but they share the same
// core so the in-memory hash/id state stays consistent afterwards (fixes
// the duplicate-auto-snapshot-after-node-capture bug).
return await _captureCore({ label, source: "node", thumbnail });
}
function scheduleCaptureSnapshot() {
@@ -1671,6 +1688,9 @@ async function restoreSnapshot(record) {
if (!full) { showToast("Failed to load snapshot data", "error"); return; }
record = full;
}
// Preserve any unsaved live edits before replacing the canvas (parity with
// swap). Deduped by hash, so it's a no-op when there is nothing new to save.
await captureSnapshot("Current").catch(() => {});
await withRestoreLock(async () => {
if (!validateSnapshotData(record.graphData)) {
showToast("Invalid snapshot data", "error");
@@ -1680,7 +1700,7 @@ async function restoreSnapshot(record) {
await app.loadGraphData(record.graphData, true, true);
const wfKey = getWorkflowKey();
lastCapturedHashMap.set(wfKey, quickHash(JSON.stringify(record.graphData)));
lastGraphDataMap.set(wfKey, record.graphData);
setLastGraphData(wfKey, record.graphData);
showToast("Snapshot restored", "success");
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Restore failed:`, err);
@@ -1699,13 +1719,14 @@ async function swapSnapshot(record) {
if (!confirmed) return;
}
// Auto-save current state before swapping (so user can get back),
// but skip if the graph is already a saved snapshot (browsing between old ones)
// Auto-save current state before swapping (so the user can get back).
// captureSnapshot() dedupes by hash, so this is a no-op when browsing
// between already-saved snapshots, but it WILL save any unsaved live edits
// made after a previous swap (when activeSnapshotId is still set) — which
// would otherwise be silently discarded by the load below.
const prevCurrentId = currentSnapshotId;
if (!activeSnapshotId) {
const capturedId = await captureSnapshot("Current");
currentSnapshotId = capturedId || prevCurrentId;
}
const capturedId = await captureSnapshot("Current");
currentSnapshotId = capturedId || prevCurrentId;
if (!record.graphData) {
const full = await db_getFullRecord(record.workflowKey, record.id);
@@ -1723,7 +1744,7 @@ async function swapSnapshot(record) {
await app.loadGraphData(record.graphData, true, true, workflow);
const wfKey = getWorkflowKey();
lastCapturedHashMap.set(wfKey, quickHash(JSON.stringify(record.graphData)));
lastGraphDataMap.set(wfKey, record.graphData);
setLastGraphData(wfKey, record.graphData);
activeSnapshotId = record.id;
showToast("Snapshot swapped", "success");
} catch (err) {
@@ -2007,6 +2028,25 @@ const CSS = `
background: #dc2626;
color: #fff;
}
.snap-footer-row {
display: flex;
gap: 6px;
margin-bottom: 6px;
}
.snap-footer-row button {
width: auto;
flex: 1;
}
.snap-footer-row button:hover {
background: var(--p-primary-color, #2563eb);
color: #fff;
}
.snap-usage-label {
font-size: 11px;
color: var(--descrip-text, #888);
text-align: center;
padding: 6px 2px 0;
}
.snap-item-node {
border-left: 3px solid #6d28d9;
}
@@ -2773,6 +2813,29 @@ function formatDate(ts) {
return d.toLocaleDateString([], { month: "short", day: "numeric", year: "numeric" });
}
function formatBytes(n) {
if (n < 1024) return `${n} B`;
if (n < 1024 * 1024) return `${(n / 1024).toFixed(1)} KB`;
if (n < 1024 * 1024 * 1024) return `${(n / (1024 * 1024)).toFixed(1)} MB`;
return `${(n / (1024 * 1024 * 1024)).toFixed(2)} GB`;
}
function formatRelativeTime(ts) {
const diff = Date.now() - ts;
if (diff < 45000) return "just now";
const m = Math.floor(diff / 60000);
if (m < 60) return `${m}m ago`;
const h = Math.floor(m / 60);
if (h < 24) return `${h}h ago`;
const d = Math.floor(h / 24);
if (d < 7) return `${d}d ago`;
const w = Math.floor(d / 7);
if (w < 5) return `${w}w ago`;
const mo = Math.floor(d / 30);
if (mo < 12) return `${mo}mo ago`;
return `${Math.floor(d / 365)}y ago`;
}
function buildBranchNavigator(forkPointId, children, selectedIndex, refreshFn) {
const nav = document.createElement("div");
nav.className = "snap-branch-nav";
@@ -2843,7 +2906,7 @@ async function buildSidebar(el) {
takeBtn.textContent = "Saving...";
try {
const saved = await captureSnapshot(name);
if (saved) showToast("Snapshot saved", "success");
showToast(saved ? "Snapshot saved" : "No changes since last snapshot", saved ? "success" : "info");
} finally {
const isViewingOther = viewingWorkflowKey != null && viewingWorkflowKey !== getWorkflowKey();
takeBtn.disabled = isViewingOther;
@@ -2893,6 +2956,17 @@ async function buildSidebar(el) {
filterItems(searchInput.value.toLowerCase());
});
const pauseBtn = document.createElement("button");
pauseBtn.className = "snap-filter-auto-btn" + (autoCaptureEnabled ? "" : " active");
pauseBtn.textContent = autoCaptureEnabled ? "Auto: On" : "Auto: Off";
pauseBtn.title = "Pause/resume automatic snapshot capture for this session";
pauseBtn.addEventListener("click", () => {
autoCaptureEnabled = !autoCaptureEnabled;
pauseBtn.classList.toggle("active", !autoCaptureEnabled);
pauseBtn.textContent = autoCaptureEnabled ? "Auto: On" : "Auto: Off";
showToast(autoCaptureEnabled ? "Auto-capture resumed" : "Auto-capture paused", "info");
});
const branchToggleBtn = document.createElement("button");
branchToggleBtn.className = "snap-filter-auto-btn" + (isBranchingEnabled() ? " active" : "");
branchToggleBtn.style.display = BRANCHING_ENABLED ? "" : "none";
@@ -2913,6 +2987,7 @@ async function buildSidebar(el) {
searchRow.appendChild(searchInput);
searchRow.appendChild(searchClear);
searchRow.appendChild(autoFilterBtn);
searchRow.appendChild(pauseBtn);
searchRow.appendChild(branchToggleBtn);
// Workflow selector
@@ -3057,6 +3132,75 @@ async function buildSidebar(el) {
const footer = document.createElement("div");
footer.className = "snap-footer";
// Export / Import row
const ioRow = document.createElement("div");
ioRow.className = "snap-footer-row";
const usageLabel = document.createElement("div");
usageLabel.className = "snap-usage-label";
async function updateUsageLabel() {
const usage = await db_getStorageUsage();
if (!usage) { usageLabel.textContent = ""; return; }
usageLabel.textContent = `Storage: ${formatBytes(usage.totalBytes)} · ${usage.workflows.length} workflow(s)`;
usageLabel.title = "Total snapshot storage on the server";
}
const exportBtn = document.createElement("button");
exportBtn.textContent = "Export";
exportBtn.title = "Download all snapshots for this workflow as a JSON file";
exportBtn.addEventListener("click", async () => {
const effKey = getEffectiveWorkflowKey();
try {
const data = await db_exportWorkflow(effKey);
if (!data.records || data.records.length === 0) { showToast("No snapshots to export", "info"); return; }
const blob = new Blob([JSON.stringify(data)], { type: "application/json" });
const url = URL.createObjectURL(blob);
const a = document.createElement("a");
a.href = url;
a.download = `snapshots-${effKey.replace(/[^a-z0-9._-]+/gi, "_").slice(0, 60)}.json`;
document.body.appendChild(a);
a.click();
a.remove();
URL.revokeObjectURL(url);
showToast(`Exported ${data.records.length} snapshot(s)`, "success");
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Export failed:`, err);
showToast("Export failed", "error");
}
});
const importBtn = document.createElement("button");
importBtn.textContent = "Import";
importBtn.title = "Import snapshots from an exported JSON file";
importBtn.addEventListener("click", () => {
const input = document.createElement("input");
input.type = "file";
input.accept = "application/json,.json";
input.addEventListener("change", async () => {
const file = input.files && input.files[0];
if (!file) return;
try {
const parsed = JSON.parse(await file.text());
const records = Array.isArray(parsed) ? parsed : parsed.records;
if (!Array.isArray(records) || records.length === 0) { showToast("No records found in file", "error"); return; }
const { imported } = await db_importRecords(records);
showToast(`Imported ${imported} snapshot(s)`, "success");
pickerDirty = true;
await refresh(true);
if (timelineRefresh) timelineRefresh().catch(() => {});
updateUsageLabel();
} catch (err) {
console.warn(`[${EXTENSION_NAME}] Import failed:`, err);
showToast("Import failed — invalid file", "error");
}
});
input.click();
});
ioRow.appendChild(exportBtn);
ioRow.appendChild(importBtn);
footer.appendChild(ioRow);
const clearBtn = document.createElement("button");
clearBtn.textContent = "Clear All Snapshots";
clearBtn.addEventListener("click", async () => {
@@ -3085,8 +3229,11 @@ async function buildSidebar(el) {
if (timelineRefresh) {
timelineRefresh().catch(() => {});
}
updateUsageLabel();
});
footer.appendChild(clearBtn);
footer.appendChild(usageLabel);
updateUsageLabel();
// ─── Profiles Section ──────────────────────────────────────────
const profilesSection = document.createElement("div");
@@ -3284,7 +3431,6 @@ async function buildSidebar(el) {
}
async function refresh(resetSearch = false) {
svgCache.clear();
// Hide tooltip — items are about to be destroyed so mouseleave won't fire
if (tooltipTimer) { clearTimeout(tooltipTimer); tooltipTimer = null; }
tooltip.classList.remove("visible");
@@ -3293,6 +3439,8 @@ async function buildSidebar(el) {
const isViewingOther = viewingWorkflowKey != null && viewingWorkflowKey !== currentKey;
const allRecords = await db_getAllForWorkflow(effKey);
// Keep immutable SVG previews across refreshes; only drop stale/excess.
pruneSvgCache(allRecords);
let nodeCount = 0;
for (const r of allRecords) if (r.source === "node") nodeCount++;
@@ -3432,7 +3580,8 @@ async function buildSidebar(el) {
const time = document.createElement("div");
time.className = "snap-item-time";
time.textContent = formatTime(rec.timestamp);
time.textContent = formatRelativeTime(rec.timestamp);
time.title = formatTime(rec.timestamp);
const date = document.createElement("div");
date.className = "snap-item-date";
@@ -3482,7 +3631,19 @@ async function buildSidebar(el) {
const newNotes = textarea.value.trim();
rec.notes = newNotes || undefined;
await db_updateMeta(rec.workflowKey, rec.id, { notes: newNotes || null });
await refresh();
textarea.remove();
// Update this item in place instead of rebuilding the list.
noteBtn.className = "snap-btn-note" + (rec.notes ? " has-note" : "");
noteBtn.title = rec.notes ? "Edit note" : "Add note";
if (rec.notes) {
notesDiv.textContent = rec.notes;
notesDiv.title = rec.notes;
notesDiv.style.display = "";
} else {
notesDiv.style.display = "none";
}
const entry = itemEntries.find((e) => e.element === item);
if (entry) entry.notes = rec.notes || ""; // keep search index in sync
};
textarea.addEventListener("keydown", (ev) => {
if (ev.key === "Enter" && ev.ctrlKey) { ev.preventDefault(); textarea.blur(); }
@@ -3497,10 +3658,16 @@ async function buildSidebar(el) {
? '<svg width="13" height="13" viewBox="0 0 16 16" fill="none"><rect x="3" y="7" width="10" height="7" rx="1.5" fill="currentColor"/><path d="M5 7V5a3 3 0 016 0v2" stroke="currentColor" stroke-width="1.5" fill="none"/></svg>'
: '<svg width="13" height="13" viewBox="0 0 16 16" fill="none"><rect x="3" y="7" width="10" height="7" rx="1.5" fill="currentColor"/><path d="M5 7V5a3 3 0 016 0" stroke="currentColor" stroke-width="1.5" fill="none"/></svg>';
lockBtn.title = rec.locked ? "Unlock snapshot" : "Lock snapshot";
const LOCK_ICON_LOCKED = '<svg width="13" height="13" viewBox="0 0 16 16" fill="none"><rect x="3" y="7" width="10" height="7" rx="1.5" fill="currentColor"/><path d="M5 7V5a3 3 0 016 0v2" stroke="currentColor" stroke-width="1.5" fill="none"/></svg>';
const LOCK_ICON_UNLOCKED = '<svg width="13" height="13" viewBox="0 0 16 16" fill="none"><rect x="3" y="7" width="10" height="7" rx="1.5" fill="currentColor"/><path d="M5 7V5a3 3 0 016 0" stroke="currentColor" stroke-width="1.5" fill="none"/></svg>';
lockBtn.addEventListener("click", async () => {
rec.locked = !rec.locked;
await db_updateMeta(rec.workflowKey, rec.id, { locked: rec.locked });
await refresh();
// Update just this item in place — locking is purely visual and
// does not change list membership/order, so skip a full rebuild.
lockBtn.className = rec.locked ? "snap-btn-lock snap-btn-locked" : "snap-btn-lock";
lockBtn.innerHTML = rec.locked ? LOCK_ICON_LOCKED : LOCK_ICON_UNLOCKED;
lockBtn.title = rec.locked ? "Unlock snapshot" : "Lock snapshot";
});
const swapBtn = document.createElement("button");
@@ -3529,6 +3696,11 @@ async function buildSidebar(el) {
if (rec.locked) {
const confirmed = await showConfirmDialog("This snapshot is locked. Delete anyway?");
if (!confirmed) return;
} else if (rec.label && !["Auto", "Initial", "Current"].includes(rec.label)) {
// Confirm before deleting a named/manual/node snapshot; auto
// snapshots are disposable and frequent, so skip the prompt.
const confirmed = await showConfirmDialog(`Delete snapshot "${rec.label}"? This cannot be undone.`);
if (!confirmed) return;
}
// Fork-point deletion: rebuild tree from fresh data, then re-parent children
if (isBranchingEnabled()) {
@@ -3622,9 +3794,11 @@ async function buildSidebar(el) {
item.addEventListener("mouseenter", () => {
tooltipTimer = setTimeout(async () => {
tooltip.innerHTML = "";
if (rec.thumbnail) {
const thumbUrl = await getThumbnailDataUrl(rec);
if (!tooltipTimer) return;
if (thumbUrl) {
const img = document.createElement("img");
img.src = `data:image/jpeg;base64,${rec.thumbnail}`;
img.src = thumbUrl;
img.style.cssText = "max-width:240px;max-height:180px;border-radius:4px;display:block;";
tooltip.appendChild(img);
} else {
@@ -3665,11 +3839,11 @@ async function buildSidebar(el) {
tooltip.classList.remove("visible");
});
if (rec.thumbnail) {
if (rec.thumbnail || rec.hasThumbnail) {
const thumb = document.createElement("img");
thumb.className = "snap-item-thumb";
thumb.src = `data:image/jpeg;base64,${rec.thumbnail}`;
item.appendChild(thumb);
getThumbnailDataUrl(rec).then((url) => { if (url) thumb.src = url; });
}
item.appendChild(info);
item.appendChild(actions);
@@ -3995,6 +4169,17 @@ if (window.__COMFYUI_FRONTEND_VERSION__) {
maxNodeSnapshots = value;
},
},
{
id: "SnapshotManager.maxAgeDays",
name: "Auto-delete snapshots older than (days, 0 = never)",
type: "slider",
defaultValue: 0,
attrs: { min: 0, max: 365, step: 1 },
category: ["Snapshot Manager", "Capture Settings", "Auto-delete age (days)"],
onChange(value) {
maxAgeDays = value;
},
},
{
id: "SnapshotManager.showTimeline",
name: "Show snapshot timeline on canvas",
@@ -4103,7 +4288,12 @@ if (window.__COMFYUI_FRONTEND_VERSION__) {
// Ctrl+S / Cmd+S shortcut for manual snapshot
document.addEventListener("keydown", (e) => {
if ((e.ctrlKey || e.metaKey) && e.key === "s") {
if ((e.ctrlKey || e.metaKey) && (e.key === "s" || e.key === "S")) {
// Ignore while typing in an editable field (node widgets, search,
// note textareas) or while browsing another workflow's history.
const t = e.target;
if (t && (t.tagName === "INPUT" || t.tagName === "TEXTAREA" || t.isContentEditable)) return;
if (viewingWorkflowKey != null && viewingWorkflowKey !== getWorkflowKey()) return;
captureSnapshot("Manual (Ctrl+S)").then((saved) => {
if (saved) showToast("Snapshot saved", "success");
}).catch(() => {});
@@ -4125,7 +4315,7 @@ if (window.__COMFYUI_FRONTEND_VERSION__) {
const graphData = getGraphData();
if (graphData) {
lastCapturedHashMap.set(wfKey, quickHash(JSON.stringify(graphData)));
lastGraphDataMap.set(wfKey, graphData);
setLastGraphData(wfKey, graphData);
}
if (timelineRefresh) timelineRefresh().catch(() => {});
}
snapshot_routes.py
+94 -27
View File
@@ -4,6 +4,8 @@ HTTP route handlers for snapshot storage.
Registers endpoints with PromptServer.instance.routes at import time.
"""
import logging
from aiohttp import web
from server import PromptServer
@@ -11,10 +13,22 @@ from . import snapshot_storage as storage
routes = PromptServer.instance.routes
# Sanity caps to bound disk/memory use from a single request.
_MAX_BODY_BYTES = 128 * 1024 * 1024 # 128 MB per request
_MAX_MIGRATE_RECORDS = 10000
def _too_large(request):
"""True if the request advertises a body larger than the cap."""
cl = request.content_length
return cl is not None and cl > _MAX_BODY_BYTES
@routes.post("/snapshot-manager/save")
async def save_snapshot(request):
try:
if _too_large(request):
return web.json_response({"error": "Request too large"}, status=413)
data = await request.json()
record = data.get("record")
if not record or "id" not in record or "workflowKey" not in record:
@@ -23,8 +37,9 @@ async def save_snapshot(request):
return web.json_response({"ok": True})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/list")
@@ -36,8 +51,11 @@ async def list_snapshots(request):
return web.json_response({"error": "Missing workflowKey"}, status=400)
records = storage.get_all_for_workflow(workflow_key)
return web.json_response(records)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/get")
@@ -54,8 +72,9 @@ async def get_snapshot(request):
return web.json_response(record)
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/update-meta")
@@ -73,8 +92,9 @@ async def update_snapshot_meta(request):
return web.json_response({"ok": True})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/delete")
@@ -89,8 +109,9 @@ async def delete_snapshot(request):
return web.json_response({"ok": True})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/delete-all")
@@ -102,8 +123,11 @@ async def delete_all_snapshots(request):
return web.json_response({"error": "Missing workflowKey"}, status=400)
result = storage.delete_all_for_workflow(workflow_key)
return web.json_response(result)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.get("/snapshot-manager/workflows")
@@ -111,8 +135,34 @@ async def list_workflows(request):
try:
keys = storage.get_all_workflow_keys()
return web.json_response(keys)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.get("/snapshot-manager/usage")
async def storage_usage(request):
try:
return web.json_response(storage.get_storage_usage())
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/export")
async def export_workflow(request):
try:
data = await request.json()
workflow_key = data.get("workflowKey")
if not workflow_key:
return web.json_response({"error": "Missing workflowKey"}, status=400)
records = storage.get_full_records_for_workflow(workflow_key)
return web.json_response({"version": 1, "workflowKey": workflow_key, "records": records})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/prune")
@@ -123,21 +173,33 @@ async def prune_snapshots(request):
max_snapshots = data.get("maxSnapshots")
source = data.get("source")
protected_ids = data.get("protectedIds")
max_age_days = data.get("maxAgeDays")
if not workflow_key or max_snapshots is None:
return web.json_response({"error": "Missing workflowKey or maxSnapshots"}, status=400)
deleted = storage.prune(workflow_key, int(max_snapshots), source=source, protected_ids=protected_ids)
deleted = storage.prune(
workflow_key, int(max_snapshots),
source=source, protected_ids=protected_ids,
max_age_days=int(max_age_days) if max_age_days else None,
)
return web.json_response({"deleted": deleted})
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/migrate")
async def migrate_snapshots(request):
try:
if _too_large(request):
return web.json_response({"error": "Request too large"}, status=413)
data = await request.json()
records = data.get("records")
if not isinstance(records, list):
return web.json_response({"error": "Missing records array"}, status=400)
if len(records) > _MAX_MIGRATE_RECORDS:
return web.json_response({"error": "Too many records"}, status=413)
imported = 0
for record in records:
if "id" in record and "workflowKey" in record:
@@ -146,8 +208,9 @@ async def migrate_snapshots(request):
return web.json_response({"imported": imported})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
# ─── Profile Endpoints ───────────────────────────────────────────────
@@ -163,8 +226,9 @@ async def save_profile(request):
return web.json_response({"ok": True})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.get("/snapshot-manager/profile/list")
@@ -172,8 +236,9 @@ async def list_profiles(request):
try:
profiles = storage.profile_get_all()
return web.json_response(profiles)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/profile/get")
@@ -189,8 +254,9 @@ async def get_profile(request):
return web.json_response(profile)
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
@routes.post("/snapshot-manager/profile/delete")
@@ -204,5 +270,6 @@ async def delete_profile(request):
return web.json_response({"ok": True})
except ValueError as e:
return web.json_response({"error": str(e)}, status=400)
except Exception as e:
return web.json_response({"error": str(e)}, status=500)
except Exception:
logging.exception("[Snapshot Manager] request handler error")
return web.json_response({"error": "Internal server error"}, status=500)
snapshot_storage.py
+117 -15
View File
@@ -13,6 +13,8 @@ operations. Only get_full_record() reads a file from disk after warm-up.
import json
import os
import shutil
import tempfile
import time
import urllib.parse
# ─── Data directory resolution ───────────────────────────────────────
@@ -37,8 +39,15 @@ _cache_warmed = set() # workflow keys already loaded from disk
def _extract_meta(record):
"""Return a lightweight copy of *record* without graphData."""
return {k: v for k, v in record.items() if k != "graphData"}
"""Return a lightweight copy of *record* without graphData or thumbnail.
The (potentially large, base64) thumbnail is replaced by a boolean
``hasThumbnail`` flag; clients lazy-load the image via get_full_record.
"""
meta = {k: v for k, v in record.items() if k not in ("graphData", "thumbnail")}
if record.get("thumbnail"):
meta["hasThumbnail"] = True
return meta
def _ensure_cached(workflow_key):
@@ -65,8 +74,17 @@ def _ensure_cached(workflow_key):
# ─── Helpers ─────────────────────────────────────────────────────────
def _workflow_dir(workflow_key):
if not workflow_key or not isinstance(workflow_key, str):
raise ValueError(f"Invalid workflow key: {workflow_key!r}")
encoded = urllib.parse.quote(workflow_key, safe="")
return os.path.join(_DATA_DIR, encoded)
path = os.path.normpath(os.path.join(_DATA_DIR, encoded))
# Defense in depth: urllib.parse.quote() leaves "." and ".." unescaped, so a
# key like ".." would escape the snapshots root (the "/" in "../.." *is*
# escaped, so escapes are bounded to one level — but block it anyway).
# Require the resolved directory to be a direct child of _DATA_DIR.
if os.path.dirname(path) != os.path.normpath(_DATA_DIR):
raise ValueError(f"Invalid workflow key: {workflow_key!r}")
return path
def _validate_id(snapshot_id):
@@ -74,6 +92,26 @@ def _validate_id(snapshot_id):
raise ValueError(f"Invalid snapshot id: {snapshot_id!r}")
def _atomic_write_json(path, obj):
"""Write *obj* as JSON to *path* atomically (temp file + os.replace).
Prevents a crash or concurrent reader mid-write from observing a
truncated/corrupt file (the old in-place open("w") truncated first).
"""
directory = os.path.dirname(path)
fd, tmp = tempfile.mkstemp(dir=directory, suffix=".tmp")
try:
with os.fdopen(fd, "w", encoding="utf-8") as f:
json.dump(obj, f, separators=(",", ":"))
os.replace(tmp, path)
except BaseException:
try:
os.remove(tmp)
except OSError:
pass
raise
# ─── Public API ──────────────────────────────────────────────────────
def put(record):
@@ -84,8 +122,7 @@ def put(record):
d = _workflow_dir(workflow_key)
os.makedirs(d, exist_ok=True)
path = os.path.join(d, f"{snapshot_id}.json")
with open(path, "w", encoding="utf-8") as f:
json.dump(record, f, separators=(",", ":"))
_atomic_write_json(path, record)
# Update cache only if already warmed; otherwise _ensure_cached will
# pick up the new file from disk on next read.
@@ -132,8 +169,7 @@ def update_meta(workflow_key, snapshot_id, fields):
record.pop(k, None)
else:
record[k] = v
with open(path, "w", encoding="utf-8") as f:
json.dump(record, f, separators=(",", ":"))
_atomic_write_json(path, record)
# Update cache entry
for entry in _cache.get(workflow_key, []):
if entry.get("id") == snapshot_id:
@@ -207,7 +243,10 @@ def get_all_workflow_keys():
if not os.path.isdir(subdir):
continue
workflow_key = urllib.parse.unquote(encoded_name)
entries = _ensure_cached(workflow_key)
try:
entries = _ensure_cached(workflow_key)
except ValueError:
continue # skip stray/legacy dirs whose name is not a valid key
if not entries:
continue
results.append({"workflowKey": workflow_key, "count": len(entries)})
@@ -215,7 +254,55 @@ def get_all_workflow_keys():
return results
def prune(workflow_key, max_snapshots, source=None, protected_ids=None):
def get_storage_usage():
"""Return {totalBytes, workflows: [{workflowKey, bytes, count}]} for all snapshots."""
workflows = []
total = 0
if os.path.isdir(_DATA_DIR):
for encoded_name in os.listdir(_DATA_DIR):
subdir = os.path.join(_DATA_DIR, encoded_name)
if not os.path.isdir(subdir):
continue
size = 0
count = 0
for fname in os.listdir(subdir):
if not fname.endswith(".json"):
continue
try:
size += os.path.getsize(os.path.join(subdir, fname))
count += 1
except OSError:
continue
if count == 0:
continue
total += size
workflows.append({
"workflowKey": urllib.parse.unquote(encoded_name),
"bytes": size,
"count": count,
})
workflows.sort(key=lambda w: w["bytes"], reverse=True)
return {"totalBytes": total, "workflows": workflows}
def get_full_records_for_workflow(workflow_key):
"""Return all full snapshot records (with graphData) for a workflow, for export."""
d = _workflow_dir(workflow_key)
records = []
if os.path.isdir(d):
for fname in os.listdir(d):
if not fname.endswith(".json"):
continue
try:
with open(os.path.join(d, fname), "r", encoding="utf-8") as f:
records.append(json.load(f))
except (json.JSONDecodeError, OSError):
continue
records.sort(key=lambda r: r.get("timestamp", 0))
return records
def prune(workflow_key, max_snapshots, source=None, protected_ids=None, max_age_days=None):
"""Delete oldest unlocked snapshots beyond limit. Returns count deleted.
source filtering:
@@ -225,6 +312,9 @@ def prune(workflow_key, max_snapshots, source=None, protected_ids=None):
protected_ids: set/list of snapshot IDs that must not be pruned
(e.g. ancestors of active branch tip, fork-point snapshots).
max_age_days: when > 0, also delete unlocked/unprotected snapshots older
than this many days, regardless of the count limit.
"""
_protected = set(protected_ids) if protected_ids else set()
entries = _ensure_cached(workflow_key)
@@ -234,9 +324,23 @@ def prune(workflow_key, max_snapshots, source=None, protected_ids=None):
candidates = [r for r in entries if not r.get("locked") and r.get("source") != "node" and r.get("id") not in _protected]
else:
candidates = [r for r in entries if not r.get("locked") and r.get("id") not in _protected]
if len(candidates) <= max_snapshots:
delete_ids = set()
to_delete = []
# Oldest-beyond-count get deleted...
if len(candidates) > max_snapshots:
for rec in candidates[: len(candidates) - max_snapshots]:
to_delete.append(rec)
delete_ids.add(rec["id"])
# ...as do any candidates older than the age cutoff (locked/protected
# snapshots were already excluded from candidates above).
if max_age_days and max_age_days > 0:
cutoff = time.time() * 1000 - max_age_days * 86400000
for rec in candidates:
if rec["id"] not in delete_ids and rec.get("timestamp", 0) < cutoff:
to_delete.append(rec)
delete_ids.add(rec["id"])
if not to_delete:
return 0
to_delete = candidates[: len(candidates) - max_snapshots]
d = _workflow_dir(workflow_key)
deleted = 0
delete_ids = set()
@@ -304,8 +408,7 @@ def profile_put(profile):
_validate_id(pid)
_ensure_profiles_dir()
path = os.path.join(_PROFILES_DIR, f"{pid}.json")
with open(path, "w", encoding="utf-8") as f:
json.dump(profile, f, separators=(",", ":"))
_atomic_write_json(path, profile)
_invalidate_profile_cache()
@@ -349,8 +452,7 @@ def profile_update(profile_id, fields):
profile.pop(k, None)
else:
profile[k] = v
with open(path, "w", encoding="utf-8") as f:
json.dump(profile, f, separators=(",", ":"))
_atomic_write_json(path, profile)
_invalidate_profile_cache()
return True