astro-tractstack 2.0.40 → 2.0.42

package/dist/index.js

@@ -2068,9 +2068,15 @@ async function w(t, e, c) {
     },
     {
       src: t(
-        "../templates/src/components/storykeep/widgets/SetupWizard.tsx"
+        "../templates/src/components/storykeep/widgets/HydrateWizard.tsx"
       ),
-      dest: "src/components/storykeep/widgets/SetupWizard.tsx"
+      dest: "src/components/storykeep/widgets/HydrateWizard.tsx"
+    },
+    {
+      src: t(
+        "../templates/src/components/storykeep/widgets/InitWizard.tsx"
+      ),
+      dest: "src/components/storykeep/widgets/InitWizard.tsx"
     },
     {
       src: t("../templates/src/pages/storykeep/init.astro"),

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "astro-tractstack",
-  "version": "2.0.40",
+  "version": "2.0.42",
   "description": "Astro integration for TractStack - redeeming the web from boring experiences",
   "type": "module",
   "main": "dist/index.js",

package/templates/src/components/Header.astro

@@ -34,6 +34,7 @@ const {
 } = Astro.props;
 
 const isHome = slug === brandConfig?.HOME_SLUG;
+console.log(slug, isHome);
 
 const tenantId =
   Astro.locals.tenant?.id || import.meta.env.PUBLIC_TENANTID || 'default';

package/templates/src/components/compositor/Node.tsx

@@ -198,7 +198,10 @@ const getElement = (
                 panelType="storyfragment"
                 ctx={getCtx(props)}
               >
-                <StoryFragmentConfigPanel nodeId={props.nodeId} />
+                <StoryFragmentConfigPanel
+                  nodeId={props.nodeId}
+                  isSandboxMode={props.isSandboxMode || false}
+                />
               </PanelVisibilityWrapper>
               <StoryFragment {...sharedProps} />
             </>

package/templates/src/components/compositor/preview/PanesPreviewGenerator.tsx

@@ -1,6 +1,7 @@
 import { useEffect, useState } from 'react';
 import { transformLivePaneForPreview } from '@/utils/etl';
 import type { NodesContext } from '@/stores/nodes';
+import { TractStackAPI } from '@/utils/api';
 
 export interface PanePreviewRequest {
   id: string;
@@ -25,6 +26,10 @@ export const PanesPreviewGenerator = ({
   onError,
 }: PanesPreviewGeneratorProps) => {
   const [isGenerating, setIsGenerating] = useState(false);
+  const tenantId =
+    window.TRACTSTACK_CONFIG?.tenantId ||
+    import.meta.env.PUBLIC_TENANTID ||
+    'default';
 
   useEffect(() => {
     if (requests.length === 0) return;
@@ -67,27 +72,22 @@ export const PanesPreviewGenerator = ({
           requestMap.set(previewPayload.id, request.id);
         }
 
-        const goBackend =
-          import.meta.env.PUBLIC_GO_BACKEND || 'http://localhost:8080';
-        const response = await fetch(`${goBackend}/api/v1/fragments/preview`, {
-          method: 'POST',
-          headers: {
-            'Content-Type': 'application/json',
-            'X-Tenant-ID': import.meta.env.PUBLIC_TENANTID || 'default',
-          },
-          body: JSON.stringify({ panes: previewPayloads }),
+        const api = new TractStackAPI(tenantId);
+        const response = await api.post('/api/v1/fragments/preview', {
+          panes: previewPayloads,
         });
 
-        if (!response.ok) {
-          throw new Error(`Preview API failed: ${response.status}`);
+        if (!response.success) {
+          throw new Error(response.error || `Preview API failed`);
         }
 
-        const { fragments, errors } = await response.json();
+        // TractStackAPI unwraps the response.data for us
+        const { fragments, errors } = response.data;
 
         const results: PaneFragmentResult[] = [];
 
         for (const [paneId, requestId] of requestMap.entries()) {
-          if (fragments[paneId]) {
+          if (fragments && fragments[paneId]) {
             results.push({
               id: requestId,
               htmlString: fragments[paneId],

package/templates/src/components/edit/SettingsPanel.tsx

@@ -4,14 +4,12 @@ import XMarkIcon from '@heroicons/react/24/outline/XMarkIcon';
 import { settingsPanelStore } from '@/stores/storykeep';
 import { getCtx } from '@/stores/nodes';
 import PanelSwitch from './PanelSwitch';
-import type { BrandConfig } from '@/types/tractstack';
 
 interface SettingsPanelProps {
-  config: BrandConfig;
   availableCodeHooks: string[];
 }
 
-const SettingsPanel = ({ config, availableCodeHooks }: SettingsPanelProps) => {
+const SettingsPanel = ({ availableCodeHooks }: SettingsPanelProps) => {
   const [panelTitle, setPanelTitle] = useState('Settings');
   const signal = useStore(settingsPanelStore);
   const ctx = getCtx();

package/templates/src/components/edit/pane/AddPanePanel_new.tsx

@@ -24,7 +24,7 @@ import { DirectInjectStep } from './steps/DirectInjectStep';
 import BooleanToggle from '@/components/form/BooleanToggle';
 import EnumSelect from '@/components/form/EnumSelect';
 import type { StoryFragmentNode } from '@/types/compositorTypes';
-import { TractStackAPI } from '@/utils/api'; // <--- IMPORT ADDED
+import { TractStackAPI } from '@/utils/api';
 
 type Step =
   | 'initial'
@@ -40,21 +40,17 @@ type InitialChoice = 'library' | 'ai' | 'blank';
 type LayoutChoice = 'standard' | 'grid';
 type ColumnPresetKey = 'left' | 'right';
 
-interface GenerationResponse {
-  success: boolean;
-  data?: { response: string | object };
-  error?: string;
-}
-
 const callAskLemurAPI = async (
   prompt: string,
   context: string,
   expectJson: boolean,
   isSandboxMode: boolean
 ): Promise<string> => {
-  // FIX: Use the centralized API class to ensure correct Tenant ID resolution
-  const api = new TractStackAPI();
-  const tenantId = api.getTenantId(); // Gets correct ID from window config
+  const tenantId =
+    window.TRACTSTACK_CONFIG?.tenantId ||
+    import.meta.env.PUBLIC_TENANTID ||
+    'default';
+  const api = new TractStackAPI(tenantId);
 
   const requestBody = {
     prompt,

package/templates/src/components/edit/pane/AddPanePanel_reuse.tsx

@@ -34,6 +34,10 @@ const AddPaneReUsePanel = ({
   first,
   setMode,
 }: AddPaneReUsePanelProps) => {
+  const tenantId =
+    window.TRACTSTACK_CONFIG?.tenantId ||
+    import.meta.env.PUBLIC_TENANTID ||
+    'default';
   const [previews, setPreviews] = useState<PreviewItem[]>([]);
   const [query, setQuery] = useState('');
   const [availablePanes, setAvailablePanes] = useState<FullContentMapItem[]>(
@@ -128,7 +132,7 @@ const AddPaneReUsePanel = ({
     const fetchFragments = async () => {
       try {
         const paneIds = visiblePreviews.map((preview) => preview.pane.id);
-        const api = new TractStackAPI();
+        const api = new TractStackAPI(tenantId);
 
         const response = await api.post('/api/v1/fragments/panes', { paneIds });
 
@@ -199,7 +203,7 @@ const AddPaneReUsePanel = ({
     if (!selectedPaneId) return;
 
     try {
-      const api = new TractStackAPI();
+      const api = new TractStackAPI(tenantId);
       const response = await api.get(
         `/api/v1/nodes/panes/${selectedPaneId}/template`
       );

package/templates/src/components/edit/pane/PanePanel_path.tsx

@@ -27,7 +27,10 @@ const PaneMagicPathPanel = ({ nodeId, setMode }: PaneMagicPathPanelProps) => {
   const [availableBeliefs, setAvailableBeliefs] = useState<BeliefNode[]>([]);
   const [isLoading, setIsLoading] = useState(true);
   const [error, setError] = useState<string | null>(null);
-  //const [isCreatingBelief, setIsCreatingBelief] = useState(false);
+  const tenantId =
+    window.TRACTSTACK_CONFIG?.tenantId ||
+    import.meta.env.PUBLIC_TENANTID ||
+    'default';
 
   const ctx = getCtx();
   const allNodes = ctx.allNodes.get();
@@ -46,7 +49,6 @@ const PaneMagicPathPanel = ({ nodeId, setMode }: PaneMagicPathPanelProps) => {
       setIsLoading(true);
       const goBackend =
         import.meta.env.PUBLIC_GO_BACKEND || 'http://localhost:8080';
-      const tenantId = import.meta.env.PUBLIC_TENANTID || 'default';
 
       // Get all belief IDs first
       const idsResponse = await fetch(`${goBackend}/api/v1/nodes/beliefs`, {
@@ -156,7 +158,6 @@ const PaneMagicPathPanel = ({ nodeId, setMode }: PaneMagicPathPanelProps) => {
 
         const goBackend =
           import.meta.env.PUBLIC_GO_BACKEND || 'http://localhost:8080';
-        const tenantId = import.meta.env.PUBLIC_TENANTID || 'default';
 
         const response = await fetch(
           `${goBackend}/api/v1/nodes/beliefs/${beliefId}`,

package/templates/src/components/edit/panels/StyleParentPanel.tsx

@@ -73,7 +73,6 @@ const StyleParentPanel = ({
       return;
     }
 
-    // --- STABILIZATION FIX START ---
     let effectiveNode = initialNode;
     if (isMarkdownPaneFragmentNode(initialNode) && initialNode.parentId) {
       const parent = ctx.allNodes.get().get(initialNode.parentId);
@@ -81,7 +80,6 @@ const StyleParentPanel = ({
         effectiveNode = parent as GridLayoutNode;
       }
     }
-    // --- STABILIZATION FIX END ---
 
     const targets: StyleableTarget[] = [];
     const isGrid = isGridLayoutNode(effectiveNode);

package/templates/src/components/edit/state/SaveModal.tsx

@@ -13,22 +13,25 @@ import {
   pendingHomePageSlugStore,
 } from '@/stores/storykeep';
 import { startLoadingAnimation } from '@/utils/helpers';
+import { processClassesForViewports } from '@/utils/compositor/reduceNodesClassNames';
 import type {
   FlatNode,
   BaseNode,
   PaneNode,
   StoryFragmentNode,
   MarkdownPaneFragmentNode,
+  GridLayoutNode,
 } from '@/types/compositorTypes';
 
 type SaveStage =
   | 'PREPARING'
   | 'SAVING_PENDING_FILES'
   | 'PROCESSING_OG_IMAGES'
+  | 'COOKING_NODES'
+  | 'PROCESSING_STYLES'
   | 'SAVING_PANES'
   | 'SAVING_STORY_FRAGMENTS'
   | 'LINKING_FILES'
-  | 'PROCESSING_STYLES'
   | 'UPDATING_HOME_PAGE'
   | 'COMPLETED'
   | 'ERROR';
@@ -46,6 +49,7 @@ interface SaveModalProps {
   isContext: boolean;
   onClose: () => void;
   isSandboxMode?: boolean;
+  hydrate?: boolean;
 }
 
 const PROGRESS_PHASES = {
@@ -56,6 +60,7 @@ const PROGRESS_PHASES = {
 };
 
 const INDETERMINATE_STAGES: SaveStage[] = [
+  'COOKING_NODES',
   'SAVING_PANES',
   'LINKING_FILES',
   'PROCESSING_STYLES',
@@ -109,6 +114,7 @@ export default function SaveModal({
   isContext,
   onClose,
   isSandboxMode = false,
+  hydrate = false,
 }: SaveModalProps) {
   const [stage, setStage] = useState<SaveStage>('PREPARING');
   const [progress, setProgress] = useState(0);
@@ -127,7 +133,10 @@ export default function SaveModal({
   const pendingHomePageSlug = pendingHomePageSlugStore.get();
   const goBackend =
     import.meta.env.PUBLIC_GO_BACKEND || 'http://localhost:8080';
-  const tenantId = import.meta.env.PUBLIC_TENANTID || 'default';
+  const tenantId =
+    window.TRACTSTACK_CONFIG?.tenantId ||
+    import.meta.env.PUBLIC_TENANTID ||
+    'default';
 
   const addDebugMessage = (message: string) => {
     const timestamp = new Date().toLocaleTimeString();
@@ -388,6 +397,203 @@ export default function SaveModal({
           dirtyPanes.length + dirtyStoryFragments.length;
         let completedProcessingSteps = 0;
 
+        // --- NEW COOKING STAGE ---
+        if (allDirtyNodes.length > 0) {
+          setStage('COOKING_NODES');
+          setIsIndeterminateStage(true);
+          addDebugMessage('Cooking nodes for whitelist extraction...');
+
+          const cookingUpdates: BaseNode[] = [];
+
+          allDirtyNodes.forEach((liveNode) => {
+            try {
+              let updatedNode: BaseNode | null = null;
+
+              // Pattern 1: TagElements -> elementCss
+              if (liveNode.nodeType === 'TagElement') {
+                const flatNode = liveNode as FlatNode;
+                const computedCSS = ctx.getNodeClasses(flatNode.id, 'auto', 0);
+                if (flatNode.elementCss !== computedCSS) {
+                  updatedNode = {
+                    ...liveNode,
+                    elementCss: computedCSS,
+                  } as FlatNode;
+                }
+              }
+              // Pattern 2: Markdown Nodes -> parentCss & gridCss
+              else if (liveNode.nodeType === 'Markdown') {
+                const markdownNode = liveNode as MarkdownPaneFragmentNode;
+                let needsUpdate = false;
+                const nextNode = { ...markdownNode };
+
+                // parentCss
+                if (markdownNode.parentClasses) {
+                  const computedParentCss = markdownNode.parentClasses.map(
+                    (_: any, index: number) =>
+                      ctx.getNodeClasses(liveNode.id, 'auto', index)
+                  );
+                  if (
+                    JSON.stringify(markdownNode.parentCss) !==
+                    JSON.stringify(computedParentCss)
+                  ) {
+                    nextNode.parentCss = computedParentCss;
+                    needsUpdate = true;
+                  }
+                }
+
+                // gridCss
+                if (markdownNode.gridClasses) {
+                  const [allClasses] = processClassesForViewports(
+                    markdownNode.gridClasses,
+                    {},
+                    1
+                  );
+                  if (allClasses && allClasses.length > 0) {
+                    const computedGridCss = allClasses[0];
+                    if (markdownNode.gridCss !== computedGridCss) {
+                      nextNode.gridCss = computedGridCss;
+                      needsUpdate = true;
+                    }
+                  }
+                }
+
+                if (needsUpdate) updatedNode = nextNode;
+              }
+              // Pattern 3: GridLayout Nodes -> parentCss & gridCss
+              else if (liveNode.nodeType === 'GridLayoutNode') {
+                const gridNode = liveNode as GridLayoutNode;
+                let needsUpdate = false;
+                const nextNode = { ...gridNode };
+
+                // parentCss
+                if (gridNode.parentClasses) {
+                  const computedParentCss = gridNode.parentClasses.map(
+                    (_: any, index: number) =>
+                      ctx.getNodeClasses(liveNode.id, 'auto', index)
+                  );
+                  if (
+                    JSON.stringify(gridNode.parentCss) !==
+                    JSON.stringify(computedParentCss)
+                  ) {
+                    nextNode.parentCss = computedParentCss.join(` `);
+                    needsUpdate = true;
+                  }
+                }
+
+                // gridCss
+                if (gridNode.gridColumns) {
+                  const { mobile, tablet, desktop } = gridNode.gridColumns;
+                  let computedGridCss = `grid grid-cols-${mobile}`;
+                  if (tablet !== mobile) {
+                    computedGridCss += ` md:grid-cols-${tablet}`;
+                  }
+                  if (desktop !== tablet) {
+                    computedGridCss += ` xl:grid-cols-${desktop}`;
+                  }
+
+                  if (gridNode.gridCss !== computedGridCss) {
+                    nextNode.gridCss = computedGridCss;
+                    needsUpdate = true;
+                  }
+                }
+
+                if (needsUpdate) updatedNode = nextNode;
+              }
+
+              if (updatedNode) {
+                cookingUpdates.push(updatedNode);
+              }
+            } catch (e) {
+              console.warn(`Failed to cook node ${liveNode.id}`, e);
+            }
+          });
+
+          if (cookingUpdates.length > 0) {
+            ctx.modifyNodes(cookingUpdates, {
+              notify: false,
+              recordHistory: false,
+            });
+            addDebugMessage(`Cooked ${cookingUpdates.length} nodes.`);
+          }
+          setIsIndeterminateStage(false);
+        }
+
+        // --- PROCESSING STYLES ---
+        // Moved before SAVING_PANES to ensure the whitelist is generated from the cooked, exhaustive inventory
+        setStage('PROCESSING_STYLES');
+        setIsIndeterminateStage(true);
+        const baseFinalizationProgress =
+          PROGRESS_PHASES.PREPARATION +
+          PROGRESS_PHASES.UPLOADS +
+          PROGRESS_PHASES.PROCESSING;
+        setProgress(
+          baseFinalizationProgress + PROGRESS_PHASES.FINALIZATION / 2
+        );
+        addDebugMessage(`Processing styles...`);
+
+        try {
+          const { dirtyPaneIds, classes: dirtyClasses } =
+            ctx.getDirtyNodesClassData();
+
+          const astroEndpoint = `/api/tailwind`;
+          const astroPayload = { dirtyPaneIds, dirtyClasses };
+          const astroResponse = await fetch(astroEndpoint, {
+            method: 'POST',
+            headers: {
+              'Content-Type': 'application/json',
+              'X-Tenant-ID': tenantId,
+            },
+            credentials: 'include',
+            body: JSON.stringify(astroPayload),
+          });
+
+          if (!astroResponse.ok) {
+            throw new Error(
+              `CSS generation failed! status: ${astroResponse.status}`
+            );
+          }
+
+          const astroResult = await astroResponse.json();
+
+          if (!astroResult.success || !astroResult.generatedCss) {
+            throw new Error('CSS generation failed: no CSS returned');
+          }
+
+          addDebugMessage(
+            `CSS generated: ${astroResult.generatedCss.length} bytes for ${dirtyClasses.length} classes`
+          );
+
+          const goEndpoint = `${goBackend}/api/v1/tailwind/update`;
+          const goPayload = { frontendCss: astroResult.generatedCss };
+          const goResponse = await fetch(goEndpoint, {
+            method: 'POST',
+            headers: {
+              'Content-Type': 'application/json',
+              'X-Tenant-ID': tenantId,
+            },
+            credentials: 'include',
+            body: JSON.stringify(goPayload),
+          });
+
+          if (!goResponse.ok) {
+            throw new Error(`CSS save failed! status: ${goResponse.status}`);
+          }
+
+          const goResult = await goResponse.json();
+          addDebugMessage(
+            `CSS saved successfully: stylesVer ${goResult.stylesVer}`
+          );
+        } catch (error) {
+          const errorMsg =
+            error instanceof Error ? error.message : 'Unknown error';
+          addDebugMessage(`Styles processing failed: ${errorMsg}`);
+          throw new Error(`Failed to process styles: ${errorMsg}`);
+        } finally {
+          setIsIndeterminateStage(false);
+        }
+
+        // --- SAVING PANES ---
+        // Runs after styles to ensure DB gets minimal, correct payload
         if (dirtyPanes.length > 0) {
           setStage('SAVING_PANES');
           setIsIndeterminateStage(true);
@@ -400,6 +606,7 @@ export default function SaveModal({
             transformLivePaneForSave(ctx, paneNode.id, isContext)
           );
 
+          // Update context with minimal strings (idempotent, restoring runtime state)
           bulkPayload.forEach((payload) => {
             payload.optionsPayload.nodes.forEach((transformedNode) => {
               const liveNode = ctx.allNodes.get().get(transformedNode.id);
@@ -453,8 +660,6 @@ export default function SaveModal({
             `Processing ${dirtyPanes.length} panes via -> POST ${endpoint}`
           );
 
-          //console.log(`bulkPayload`, bulkPayload)
-
           try {
             const response = await fetch(endpoint, {
               method: 'POST',
@@ -516,7 +721,7 @@ export default function SaveModal({
               const payload = await transformStoryFragmentForSave(
                 ctx,
                 fragment.id,
-                window.TRACTSTACK_CONFIG?.tenantId || 'default'
+                tenantId
               );
 
               if (uploadedOGPaths[fragment.id]) {
@@ -583,14 +788,11 @@ export default function SaveModal({
           }
         }
 
-        const baseFinalizationProgress =
-          PROGRESS_PHASES.PREPARATION +
-          PROGRESS_PHASES.UPLOADS +
-          PROGRESS_PHASES.PROCESSING;
-
         if (dirtyPanes.length > 0) {
           setStage('LINKING_FILES');
           setIsIndeterminateStage(true);
+          // ... Linking files logic continues ...
+          // Note: Linking files remains after saving panes because it relies on panes existing in DB
           setProgress(baseFinalizationProgress);
           addDebugMessage('Starting file-pane relationship linking...');
 
@@ -644,74 +846,6 @@ export default function SaveModal({
           setIsIndeterminateStage(false);
         }
 
-        setStage('PROCESSING_STYLES');
-        setIsIndeterminateStage(true);
-        setProgress(
-          baseFinalizationProgress + PROGRESS_PHASES.FINALIZATION / 2
-        );
-        addDebugMessage(`Processing styles...`);
-
-        try {
-          const { dirtyPaneIds, classes: dirtyClasses } =
-            ctx.getDirtyNodesClassData();
-
-          const astroEndpoint = `/api/tailwind`;
-          const astroPayload = { dirtyPaneIds, dirtyClasses };
-          const astroResponse = await fetch(astroEndpoint, {
-            method: 'POST',
-            headers: {
-              'Content-Type': 'application/json',
-              'X-Tenant-ID': tenantId,
-            },
-            credentials: 'include',
-            body: JSON.stringify(astroPayload),
-          });
-
-          if (!astroResponse.ok) {
-            throw new Error(
-              `CSS generation failed! status: ${astroResponse.status}`
-            );
-          }
-
-          const astroResult = await astroResponse.json();
-
-          if (!astroResult.success || !astroResult.generatedCss) {
-            throw new Error('CSS generation failed: no CSS returned');
-          }
-
-          addDebugMessage(
-            `CSS generated: ${astroResult.generatedCss.length} bytes for ${dirtyClasses.length} classes`
-          );
-
-          const goEndpoint = `${goBackend}/api/v1/tailwind/update`;
-          const goPayload = { frontendCss: astroResult.generatedCss };
-          const goResponse = await fetch(goEndpoint, {
-            method: 'POST',
-            headers: {
-              'Content-Type': 'application/json',
-              'X-Tenant-ID': tenantId,
-            },
-            credentials: 'include',
-            body: JSON.stringify(goPayload),
-          });
-
-          if (!goResponse.ok) {
-            throw new Error(`CSS save failed! status: ${goResponse.status}`);
-          }
-
-          const goResult = await goResponse.json();
-          addDebugMessage(
-            `CSS saved successfully: stylesVer ${goResult.stylesVer}`
-          );
-        } catch (error) {
-          const errorMsg =
-            error instanceof Error ? error.message : 'Unknown error';
-          addDebugMessage(`Styles processing failed: ${errorMsg}`);
-          throw new Error(`Failed to process styles: ${errorMsg}`);
-        } finally {
-          setIsIndeterminateStage(false);
-        }
-
         if (pendingHomePageSlug) {
           setStage('UPDATING_HOME_PAGE');
           setIsIndeterminateStage(true);
@@ -774,6 +908,30 @@ export default function SaveModal({
           }
         }
 
+        if (hydrate) {
+          addDebugMessage('Finalizing setup (Kill Switch)...');
+          try {
+            const response = await fetch(`${goBackend}/api/v1/setup/complete`, {
+              method: 'POST',
+              headers: {
+                'Content-Type': 'application/json',
+                'X-Tenant-ID': tenantId,
+              },
+              credentials: 'include',
+              body: JSON.stringify({}),
+            });
+
+            if (!response.ok) {
+              throw new Error(`Kill Switch failed: ${response.status}`);
+            }
+            addDebugMessage('Hydration token cleared.');
+          } catch (e) {
+            console.error('Kill switch error:', e);
+            // We don't throw here to ensure the user still gets to the dashboard
+            addDebugMessage('Warning: Failed to clear hydration token.');
+          }
+        }
+
         setStage('COMPLETED');
         setProgress(100);
         addDebugMessage('Save process completed successfully!');
@@ -822,6 +980,9 @@ export default function SaveModal({
       case 'PROCESSING_OG_IMAGES':
         description = `Processing social images${getProgressText()}`;
         break;
+      case 'COOKING_NODES':
+        description = 'Preparing content styles...';
+        break;
       case 'SAVING_PANES':
         description = `${actionText} pane content...`;
         break;
@@ -1031,7 +1192,17 @@ export default function SaveModal({
 
               {(stage === 'COMPLETED' || stage === 'ERROR') && (
                 <div className="flex justify-end gap-2">
-                  {stage === 'COMPLETED' && (
+                  {hydrate && stage === 'COMPLETED' && (
+                    <button
+                      onClick={() =>
+                        (window.location.href = '/storykeep/branding')
+                      }
+                      className={`rounded bg-cyan-600 px-4 py-2 text-white transition-colors hover:bg-cyan-700`}
+                    >
+                      Continue
+                    </button>
+                  )}
+                  {!hydrate && stage === 'COMPLETED' && (
                     <>
                       <a
                         href={visitPageUrl}