diff --git a/editor/contexts/shopify-context.tsx b/editor/contexts/shopify-context.tsx
index 66c7f61..e08e80a 100644
--- a/editor/contexts/shopify-context.tsx
+++ b/editor/contexts/shopify-context.tsx
@@ -106,12 +106,6 @@ export const ShopifyProvider: React.FC<{
   token: string;
   children: React.ReactNode;
 }> = ({ domain, token, children }) => {
-  console.log('[ShopifyProvider] creds', {
-    domain,
-    hasToken: !!token,
-    tokenPreview: token ? `${token.slice(0, 4)}…${token.slice(-4)}` : null,
-  });
-
   // Sync creds into the module-level store synchronously so any render-time
   // call (incl. SSR) reads the right domain/token.
   setShopifyCredentials({ domain, token });
diff --git a/src/App.tsx b/src/App.tsx
index 68bf0e5..0a5322b 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -1,4 +1,4 @@
-import { useCallback, useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useRef, useState } from "react";
 import {
   App as ReactEditorApp,
   blocksPlugin,
@@ -36,6 +36,11 @@ export default function App() {
   const pages = schemaJson as Pages;
   const [currentPath, setCurrentPath] = useState<string>(readPathname);
   const [isPublishing, setIsPublishing] = useState(false);
+  const latestDataRef = useRef<any>(null);
+
+  const handleChange = useCallback((data: any) => {
+    latestDataRef.current = data;
+  }, []);
 
   useEffect(() => {
     const onPop = () => setCurrentPath(readPathname());
@@ -52,25 +57,8 @@ export default function App() {
     [],
   );
 
-  const plugins = useMemo(
-    () => [
-      aiPlugin({
-        api: "https://www.frontend-ai.com/cloud/api/chat",
-        headers: { "X-Api-Key": AI_API_KEY },
-        getCurrentRoute: () => ({ path: window.location.pathname }),
-      }),
-      createTailwindCdnPlugin(),
-      blocksPlugin(),
-      outlinePlugin(),
-      mediaPlugin({
-        adapter: frontendAiMediaAdapter,
-        showSearch: true,
-      }),
-    ],
-    [],
-  );
-
   const handlePublish = useCallback((data: any, route?: string) => {
+    console.log(data);
     setIsPublishing(true);
     if (typeof window !== "undefined" && window.parent !== window) {
       window.parent.postMessage(
@@ -81,6 +69,30 @@ export default function App() {
     setTimeout(() => setIsPublishing(false), 1000);
   }, []);
 
+  const plugins = useMemo(
+    () => [
+      aiPlugin({
+        api: "https://www.frontend-ai.com/cloud/api/chat",
+        headers: { "X-Api-Key": AI_API_KEY },
+        getCurrentRoute: () => ({ path: window.location.pathname }),
+        onFinish: (output: any) => {
+          if (output?.finishReason === "stop") {
+            console.log(output);
+            handlePublish(latestDataRef.current, currentPath);
+          }
+        },
+      }),
+      createTailwindCdnPlugin(),
+      blocksPlugin(),
+      outlinePlugin(),
+      mediaPlugin({
+        adapter: frontendAiMediaAdapter,
+        showSearch: true,
+      }),
+    ],
+    [pages, currentPath, handlePublish],
+  );
+
   return (
     <div className="h-screen w-screen">
       <ShopifyProvider domain={SHOPIFY_DOMAIN} token={STOREFRONT_TOKEN}>
@@ -91,6 +103,24 @@ export default function App() {
           plugins={plugins}
           iframe={{ enabled: true }}
           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>
+                <button
+                  type="button"
+                  disabled={isPublishing}
+                  onClick={() => {
+                    handlePublish(latestDataRef.current, 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"
+                >
+                  {isPublishing ? "Saving..." : "Save"}
+                </button>
+              </div>
+            ),
+          }}
         />
       </ShopifyProvider>
       {isPublishing && (