diff --git a/src/components/ai-optimized-editor.tsx b/src/components/ai-optimized-editor.tsx
index 3e40f94..9e4b11e 100644
--- a/src/components/ai-optimized-editor.tsx
+++ b/src/components/ai-optimized-editor.tsx
@@ -4,25 +4,24 @@ import React, { useState, useEffect } from "react";
 import { DiffEditor } from "@monaco-editor/react";
 import { useMonacoTheme } from "@/hooks/use-monaco-theme";
 import { optimizeCode } from "@/app/actions/ai-improve";
-import { AIOptimizeButton } from "@/features/problems/code/components/toolbar/actions/AIOptimizeButton";
+import { useProblemEditorStore } from "@/stores/problem-editor";
+// import {LanguageServerConfig} from "@/generated/client";
+// import type {editor} from "monaco-editor";
 
-interface AIEditorWrapperProps {
-    language: string;
-    originalCode: string;
-    onReset: () => void;
-}
+export const AIEditorWrapper = (
+) => {
+    const {
+        language,
+        value: originalCode,
+        // setUseAIEditor,
+        setLoading,
+    } = useProblemEditorStore();
 
-export const AIEditorWrapper = ({ language, originalCode, onReset }: AIEditorWrapperProps) => {
     const [optimizedCode, setOptimizedCode] = useState<string>("");
-    const [loading, setLoading] = useState(false);
     const { theme } = useMonacoTheme();
 
-    // 自动在组件首次挂载后执行 AI 优化
     useEffect(() => {
-        if (!optimizedCode) {
-            handleOptimize();
-        }
-        // eslint-disable-next-line react-hooks/exhaustive-deps
+        handleOptimize();
     }, []);
 
     const handleOptimize = async () => {
@@ -42,27 +41,8 @@ export const AIEditorWrapper = ({ language, originalCode, onReset }: AIEditorWra
         }
     };
 
-    const handleClick = () => {
-        if (optimizedCode) {
-            // 已有优化，点击返回
-            setOptimizedCode("");
-            onReset();
-        } else {
-            // 手动触发优化（如果需要）
-            handleOptimize();
-        }
-    };
-
     return (
         <div className="w-full h-[80vh] flex flex-col gap-4">
-            <div>
-                <AIOptimizeButton
-                    loading={loading}
-                    hasOptimized={!!optimizedCode}
-                    onClick={handleClick}
-                />
-            </div>
-
             {optimizedCode && (
                 <div className="flex-1">
                     <DiffEditor
@@ -71,7 +51,11 @@ export const AIEditorWrapper = ({ language, originalCode, onReset }: AIEditorWra
                         modified={optimizedCode}
                         height="100%"
                         theme={theme}
-                        options={{ readOnly: true, renderSideBySide: true, automaticLayout: true }}
+                        options={{
+                            readOnly: true,
+                            renderSideBySide: true,
+                            automaticLayout: true,
+                        }}
                     />
                 </div>
             )}
diff --git a/src/components/problem-editor.tsx b/src/components/problem-editor.tsx
index a591203..3428c93 100644
--- a/src/components/problem-editor.tsx
+++ b/src/components/problem-editor.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { useState, useEffect } from "react";
+import { useEffect } from "react";
 import { CoreEditor } from "@/components/core-editor";
 import { useProblemEditorStore } from "@/stores/problem-editor";
 import type { LanguageServerConfig, Template } from "@/generated/client";
@@ -26,9 +26,10 @@ export const ProblemEditor = ({
         setEditor,
         setLspWebSocket,
         setMarkers,
+        useAIEditor,
+        // setUseAIEditor
     } = useProblemEditorStore();
 
-    const [useAIEditor, setUseAIEditor] = useState(false);
 
     useEffect(() => {
         setProblem(problemId, templates);
@@ -38,12 +39,12 @@ export const ProblemEditor = ({
         <div className="w-full h-[85vh] relative">
             {!useAIEditor ? (
                 <>
-                    <button
-                        className="absolute right-4 top-4 bg-blue-600 text-white px-3 py-1 rounded z-10"
-                        onClick={() => setUseAIEditor(true)}
-                    >
-                        AI 优化代码
-                    </button>
+                    {/*<button*/}
+                    {/*    className="absolute right-4 top-4 bg-blue-600 text-white px-3 py-1 rounded z-10"*/}
+                    {/*    onClick={() => setUseAIEditor(true)}*/}
+                    {/*>*/}
+                    {/*    AI 优化代码*/}
+                    {/*</button>*/}
                     <CoreEditor
                         language={language}
                         value={value}
@@ -57,11 +58,7 @@ export const ProblemEditor = ({
                     />
                 </>
             ) : (
-                <AIEditorWrapper
-                    language={language}
-                    originalCode={value}
-                    onReset={() => setUseAIEditor(false)}
-                />
+                <AIEditorWrapper/>
             )}
         </div>
     );
diff --git a/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx b/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx
new file mode 100644
index 0000000..08b88a4
--- /dev/null
+++ b/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx
@@ -0,0 +1,42 @@
+"use client";
+
+import { TooltipButton } from "@/components/tooltip-button";
+import { Wand2Icon, LoaderCircleIcon, Undo2Icon } from "lucide-react";
+import { useProblemEditorStore } from "@/stores/problem-editor";
+
+export const AIOptimizeButton = () => {
+    const { useAIEditor, setUseAIEditor, loading } = useProblemEditorStore();
+
+    const handleClick = () => {
+        if (!loading) {
+            setUseAIEditor(!useAIEditor);
+        }
+    };
+
+    const tooltipContent = loading
+        ? "AI 正在优化中…"
+        : useAIEditor
+            ? "返回原始编辑器"
+            : "使用 AI 优化代码";
+
+    return (
+        <TooltipButton
+            tooltipContent={tooltipContent}
+            onClick={handleClick}
+            disabled={loading}
+        >
+            {loading ? (
+                <LoaderCircleIcon
+                    className="opacity-60 animate-spin"
+                    size={16}
+                    strokeWidth={2}
+                    aria-hidden="true"
+                />
+            ) : useAIEditor ? (
+                <Undo2Icon size={16} strokeWidth={2} aria-hidden="true" />
+            ) : (
+                <Wand2Icon size={16} strokeWidth={2} aria-hidden="true" />
+            )}
+        </TooltipButton>
+    );
+};
diff --git a/src/features/problems/code/components/toolbar/code-toolbar.tsx b/src/features/problems/code/components/toolbar/code-toolbar.tsx
index 989b91a..dc9e12a 100644
--- a/src/features/problems/code/components/toolbar/code-toolbar.tsx
+++ b/src/features/problems/code/components/toolbar/code-toolbar.tsx
@@ -9,6 +9,7 @@ import {
 } from "@/features/problems/code/components/toolbar";
 import { AnalyzeButton } from "./actions/analyze-button";
 import { LspConnectionIndicator } from "./controls/lsp-connection-indicator";
+import {AIOptimizeButton} from "@/features/problems/code/components/toolbar/actions/AIOptimizeButton";
 
 interface CodeToolbarProps {
   className?: string;
@@ -25,6 +26,7 @@ export const CodeToolbar = async ({ className }: CodeToolbarProps) => {
           <LspConnectionIndicator />
         </div>
         <div className="flex items-center gap-2">
+          <AIOptimizeButton />
           <AnalyzeButton />
           <ResetButton />
           <UndoButton />
diff --git a/src/stores/problem-editor.ts b/src/stores/problem-editor.ts
index 10863f8..e10166d 100644
--- a/src/stores/problem-editor.ts
+++ b/src/stores/problem-editor.ts
@@ -16,6 +16,8 @@ type ProblemEditorState = {
   editor: editor.IStandaloneCodeEditor | null;
   lspWebSocket: WebSocket | null;
   markers: editor.IMarker[];
+  useAIEditor: boolean;
+  loading: boolean;
 };
 
 type ProblemEditorAction = {
@@ -26,6 +28,8 @@ type ProblemEditorAction = {
   setEditor: (editor: editor.IStandaloneCodeEditor) => void;
   setLspWebSocket: (lspWebSocket: WebSocket) => void;
   setMarkers: (markers: editor.IMarker[]) => void;
+  setUseAIEditor: (flag: boolean) => void;
+  setLoading: (flag: boolean) => void;
 };
 
 type ProblemEditorStore = ProblemEditorState & ProblemEditorAction;
@@ -38,6 +42,10 @@ export const useProblemEditorStore = create<ProblemEditorStore>((set, get) => ({
   editor: null,
   lspWebSocket: null,
   markers: [],
+  useAIEditor: false,
+  loading: false,
+  setLoading: (loading) => set({loading}),
+  setUseAIEditor: (loading) => set({ useAIEditor: loading }),
   setProblem: (problemId, templates) => {
     const language = getLanguage(problemId);
     const value = getValue(problemId, language, templates);