diff --git a/src/App.tsx b/src/App.tsx
index e7557bc..b63d908 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -1,5 +1,5 @@
 import { useCallback, useMemo, useRef, useState } from "react";
-import { App as ReactEditorApp } from "@reacteditor/core";
+import { App as ReactEditorApp, createUseEditor } from "@reacteditor/core";
 import "@reacteditor/core/react-editor.css";
 import createTailwindCdnPlugin from "@reacteditor/plugin-tailwind-cdn";
 import { mediaPlugin } from "@reacteditor/plugin-media";
@@ -16,6 +16,8 @@ const AI_API_KEY = (import.meta.env.VITE_API_KEY as string | undefined) ?? "";
 
 type Pages = Record<string, { root: any; content: any[] }>;
 
+const useEditor = createUseEditor();
+
 const SHOPIFY_DOMAIN =
   (import.meta.env.VITE_SHOPIFY_DOMAIN as string | undefined) ?? "mock.shop";
 const STOREFRONT_TOKEN =
@@ -78,21 +80,22 @@ export default function App() {
           onPublish={handlePublish}
           onChange={handleChange}
           overrides={{
-            header: ({ children }: any) => (
-              <div className="flex items-center justify-between gap-2 border-b border-gray-200 bg-white px-4 py-2">
-                <div className="flex flex-1 items-center gap-2">{children}</div>
+            headerActions: () => {
+              const appState = useEditor((s: any) => s.appState);
+              return (
                 <button
                   type="button"
                   disabled={isPublishing}
                   onClick={() => {
-                    handlePublish(latestDataRef.current, currentPath);
+                    handlePublish(appState.data, currentPath);
                   }}
-                  className="inline-flex items-center justify-center rounded-md bg-black px-4 py-1.5 text-sm font-medium text-white hover:bg-gray-800 disabled:cursor-not-allowed disabled:opacity-60"
+                  className="inline-flex items-center justify-center gap-2 rounded-md bg-black px-4 py-1.5 text-sm font-medium text-white hover:bg-gray-800 disabled:cursor-not-allowed disabled:opacity-60"
                 >
+                  {isPublishing && <Loader size={14} />}
                   {isPublishing ? "Saving..." : "Save"}
                 </button>
-              </div>
-            ),
+              );
+            },
           }}
         />
       </ShopifyProvider>