diff --git a/src/components/ai-optimized-editor.tsx b/src/components/ai-optimized-editor.tsx
index 9e4b11e..9bcaba2 100644
--- a/src/components/ai-optimized-editor.tsx
+++ b/src/components/ai-optimized-editor.tsx
@@ -5,24 +5,27 @@ import { DiffEditor } from "@monaco-editor/react";
 import { useMonacoTheme } from "@/hooks/use-monaco-theme";
 import { optimizeCode } from "@/app/actions/ai-improve";
 import { useProblemEditorStore } from "@/stores/problem-editor";
-// import {LanguageServerConfig} from "@/generated/client";
-// import type {editor} from "monaco-editor";
 
-export const AIEditorWrapper = (
-) => {
+export const AIEditorWrapper = () => {
     const {
         language,
         value: originalCode,
-        // setUseAIEditor,
         setLoading,
+        AIgenerate,
+        LastOptimizedCode,
+        setLastOptimizedCode,
     } = useProblemEditorStore();
 
     const [optimizedCode, setOptimizedCode] = useState<string>("");
     const { theme } = useMonacoTheme();
 
     useEffect(() => {
-        handleOptimize();
-    }, []);
+        if (AIgenerate) {
+            handleOptimize();
+        } else if (LastOptimizedCode) {
+            setOptimizedCode(LastOptimizedCode);
+        }
+    }, [AIgenerate]);
 
     const handleOptimize = async () => {
         setLoading(true);
@@ -33,6 +36,7 @@ export const AIEditorWrapper = (
                 problemId: "",
             });
             setOptimizedCode(res.optimizedCode);
+            setLastOptimizedCode(res.optimizedCode);
         } catch (err) {
             console.error("优化失败", err);
             setOptimizedCode("// 优化失败，请稍后重试");
diff --git a/src/features/problems/code/components/toolbar/actions/AIDisplayButton.tsx b/src/features/problems/code/components/toolbar/actions/AIDisplayButton.tsx
new file mode 100644
index 0000000..2146912
--- /dev/null
+++ b/src/features/problems/code/components/toolbar/actions/AIDisplayButton.tsx
@@ -0,0 +1,45 @@
+"use client";
+
+import { TooltipButton } from "@/components/tooltip-button";
+import { ArrowLeftRight, LoaderCircleIcon, Undo2Icon } from "lucide-react";
+import { useProblemEditorStore } from "@/stores/problem-editor";
+
+export const AIDisplayButton = () => {
+    const { useAIEditor, setUseAIEditor,setAIgenerate ,loading } = useProblemEditorStore();
+
+    const handleClick = () => {
+        setAIgenerate(false);
+        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" />
+            )
+                : (
+                <ArrowLeftRight size={16} strokeWidth={2} aria-hidden="true" />
+            )
+            }
+        </TooltipButton>
+    );
+};
diff --git a/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx b/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx
index 08b88a4..8ad26dd 100644
--- a/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx
+++ b/src/features/problems/code/components/toolbar/actions/AIOptimizeButton.tsx
@@ -1,40 +1,47 @@
 "use client";
 
 import { TooltipButton } from "@/components/tooltip-button";
-import { Wand2Icon, LoaderCircleIcon, Undo2Icon } from "lucide-react";
+import { Wand2Icon } from "lucide-react";
+//import { LoaderCircleIcon, Undo2Icon } from "lucide-react";
 import { useProblemEditorStore } from "@/stores/problem-editor";
 
 export const AIOptimizeButton = () => {
-    const { useAIEditor, setUseAIEditor, loading } = useProblemEditorStore();
+    const { useAIEditor, setUseAIEditor, setAIgenerate, loading } = useProblemEditorStore();
 
     const handleClick = () => {
+        setAIgenerate(true);
         if (!loading) {
             setUseAIEditor(!useAIEditor);
         }
     };
 
-    const tooltipContent = loading
-        ? "AI 正在优化中…"
-        : useAIEditor
-            ? "返回原始编辑器"
-            : "使用 AI 优化代码";
+    //     ? "AI 正在优化中…"
+    //     : useAIEditor
+    //         ? "返回原始编辑器"
+    //         : "使用 AI 优化代码";
+
+    const tooltipContent = "使用 AI 优化代码"; // 仅保留默认提示内容
 
     return (
         <TooltipButton
             tooltipContent={tooltipContent}
             onClick={handleClick}
-            disabled={loading}
+            disabled={loading || useAIEditor}
         >
-            {loading ? (
-                <LoaderCircleIcon
-                    className="opacity-60 animate-spin"
-                    size={16}
-                    strokeWidth={2}
-                    aria-hidden="true"
-                />
+            {
+                loading ? (
+                // <LoaderCircleIcon
+                //     className="opacity-60 animate-spin"
+                //     size={16}
+                //     strokeWidth={2}
+                //     aria-hidden="true"
+                // />
+                    null
             ) : useAIEditor ? (
-                <Undo2Icon size={16} strokeWidth={2} aria-hidden="true" />
-            ) : (
+                // <Undo2Icon size={16} strokeWidth={2} aria-hidden="true" />
+                    null
+            ) :
+                (
                 <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 dc9e12a..4bf1655 100644
--- a/src/features/problems/code/components/toolbar/code-toolbar.tsx
+++ b/src/features/problems/code/components/toolbar/code-toolbar.tsx
@@ -10,6 +10,7 @@ import {
 import { AnalyzeButton } from "./actions/analyze-button";
 import { LspConnectionIndicator } from "./controls/lsp-connection-indicator";
 import {AIOptimizeButton} from "@/features/problems/code/components/toolbar/actions/AIOptimizeButton";
+import {AIDisplayButton} from "@/features/problems/code/components/toolbar/actions/AIDisplayButton";
 
 interface CodeToolbarProps {
   className?: string;
@@ -27,6 +28,7 @@ export const CodeToolbar = async ({ className }: CodeToolbarProps) => {
         </div>
         <div className="flex items-center gap-2">
           <AIOptimizeButton />
+          <AIDisplayButton />
           <AnalyzeButton />
           <ResetButton />
           <UndoButton />
diff --git a/src/stores/problem-editor.ts b/src/stores/problem-editor.ts
index e10166d..52f39cc 100644
--- a/src/stores/problem-editor.ts
+++ b/src/stores/problem-editor.ts
@@ -18,6 +18,8 @@ type ProblemEditorState = {
   markers: editor.IMarker[];
   useAIEditor: boolean;
   loading: boolean;
+  AIgenerate: boolean;
+  LastOptimizedCode: string;
 };
 
 type ProblemEditorAction = {
@@ -30,6 +32,8 @@ type ProblemEditorAction = {
   setMarkers: (markers: editor.IMarker[]) => void;
   setUseAIEditor: (flag: boolean) => void;
   setLoading: (flag: boolean) => void;
+  setAIgenerate: (flag: boolean) => void;
+  setLastOptimizedCode: (code: string) => void;
 };
 
 type ProblemEditorStore = ProblemEditorState & ProblemEditorAction;
@@ -44,6 +48,10 @@ export const useProblemEditorStore = create<ProblemEditorStore>((set, get) => ({
   markers: [],
   useAIEditor: false,
   loading: false,
+  AIgenerate: false,
+  LastOptimizedCode: "",
+  setLastOptimizedCode: (code) => set({ LastOptimizedCode: code }),
+  setAIgenerate: (flag) => set({ AIgenerate: flag }),
   setLoading: (loading) => set({loading}),
   setUseAIEditor: (loading) => set({ useAIEditor: loading }),
   setProblem: (problemId, templates) => {