diff --git a/frontend/app/api/merge/route.ts b/frontend/app/api/merge/route.ts
index c17eab9..c2a7600 100644
--- a/frontend/app/api/merge/route.ts
+++ b/frontend/app/api/merge/route.ts
@@ -19,7 +19,9 @@ export async function POST(request: Request) {
    - Explanations or comments
    - Markdown formatting
 
-The output should be the exact code that will replace the existing code, nothing more and nothing less.`
+The output should be the exact code that will replace the existing code, nothing more and nothing less.
+
+Important: When merging, preserve the original code structure as much as possible. Only make necessary changes to integrate the new code while maintaining the original code's organization and style.`
 
     const mergedCode = `Original file (${fileName}):\n${originalCode}\n\nNew code to merge:\n${newCode}`
 
diff --git a/frontend/components/editor/AIChat/ApplyButton.tsx b/frontend/components/editor/AIChat/ApplyButton.tsx
index f122339..8fb94ff 100644
--- a/frontend/components/editor/AIChat/ApplyButton.tsx
+++ b/frontend/components/editor/AIChat/ApplyButton.tsx
@@ -8,7 +8,7 @@ interface ApplyButtonProps {
   activeFileName: string
   activeFileContent: string
   editorRef: { current: any }
-  onApply: (mergedCode: string) => void
+  onApply: (mergedCode: string, originalCode: string) => void
 }
 
 export default function ApplyButton({
@@ -48,7 +48,7 @@ export default function ApplyButton({
           mergedCode += decoder.decode(value, { stream: true })
         }
       }
-      onApply(mergedCode.trim())
+      onApply(mergedCode.trim(), activeFileContent)
     } catch (error) {
       console.error("Error applying code:", error)
       toast.error(
diff --git a/frontend/components/editor/AIChat/lib/markdownComponents.tsx b/frontend/components/editor/AIChat/lib/markdownComponents.tsx
index 36130d5..770ee09 100644
--- a/frontend/components/editor/AIChat/lib/markdownComponents.tsx
+++ b/frontend/components/editor/AIChat/lib/markdownComponents.tsx
@@ -17,7 +17,7 @@ export const createMarkdownComponents = (
   activeFileName: string,
   activeFileContent: string,
   editorRef: any,
-  handleApplyCode: (mergedCode: string) => void,
+  handleApplyCode: (mergedCode: string, originalCode: string) => void,
   selectFile: (tab: TTab) => void,
   mergeDecorationsCollection?: monaco.editor.IEditorDecorationsCollection,
   setMergeDecorationsCollection?: (collection: undefined) => void
@@ -61,7 +61,7 @@ export const createMarkdownComponents = (
                       mergeDecorationsCollection &&
                       editorRef?.current
                     ) {
-                      mergeDecorationsCollection.clear()
+                      mergeDecorationsCollection?.clear()
                       setMergeDecorationsCollection(undefined)
                     }
                   }}
@@ -75,14 +75,15 @@ export const createMarkdownComponents = (
                 <div className="w-px bg-input"></div>
                 <Button
                   onClick={() => {
-                    if (
-                      setMergeDecorationsCollection &&
-                      mergeDecorationsCollection &&
-                      editorRef?.current
-                    ) {
-                      editorRef.current.getModel()?.setValue(activeFileContent)
-                      mergeDecorationsCollection.clear()
-                      setMergeDecorationsCollection(undefined)
+                    if (editorRef?.current && mergeDecorationsCollection) {
+                      const model = editorRef.current.getModel()
+                      if (model && (model as any).originalContent) {
+                        editorRef.current?.setValue(
+                          (model as any).originalContent
+                        )
+                        mergeDecorationsCollection.clear()
+                        setMergeDecorationsCollection?.(undefined)
+                      }
                     }
                   }}
                   size="sm"
diff --git a/frontend/components/editor/AIChat/types/index.ts b/frontend/components/editor/AIChat/types/index.ts
index 5eb7362..7fe24c4 100644
--- a/frontend/components/editor/AIChat/types/index.ts
+++ b/frontend/components/editor/AIChat/types/index.ts
@@ -57,7 +57,7 @@ export interface AIChatProps {
   templateType: string
   templateConfig?: TemplateConfig
   projectName: string
-  handleApplyCode: (mergedCode: string) => void
+  handleApplyCode: (mergedCode: string, originalCode: string) => void
   mergeDecorationsCollection?: monaco.editor.IEditorDecorationsCollection
   setMergeDecorationsCollection?: (collection: undefined) => void
   selectFile: (tab: TTab) => void
@@ -110,7 +110,7 @@ export interface MessageProps {
   ) => void
   setIsContextExpanded: (isExpanded: boolean) => void
   socket: Socket | null
-  handleApplyCode: (mergedCode: string) => void
+  handleApplyCode: (mergedCode: string, originalCode: string) => void
   activeFileName: string
   activeFileContent: string
   editorRef: any
diff --git a/frontend/components/editor/index.tsx b/frontend/components/editor/index.tsx
index c470f17..c992d61 100644
--- a/frontend/components/editor/index.tsx
+++ b/frontend/components/editor/index.tsx
@@ -384,10 +384,10 @@ export default function CodeEditor({
 
   // handle apply code
   const handleApplyCode = useCallback(
-    (mergedCode: string) => {
+    (mergedCode: string, originalCode: string) => {
       if (!editorRef) return
 
-      const originalLines = activeFileContent.split("\n")
+      const originalLines = originalCode.split("\n")
       const mergedLines = mergedCode.split("\n")
 
       const decorations: monaco.editor.IModelDeltaDecoration[] = []
@@ -397,32 +397,31 @@ export default function CodeEditor({
         i < Math.max(originalLines.length, mergedLines.length);
         i++
       ) {
-        if (originalLines[i] !== mergedLines[i]) {
+        // Only highlight new lines (green highlights)
+        if (i >= originalLines.length || originalLines[i] !== mergedLines[i]) {
           decorations.push({
             range: new monaco.Range(i + 1, 1, i + 1, 1),
             options: {
               isWholeLine: true,
-              className:
-                i < originalLines.length
-                  ? "removed-line-decoration"
-                  : "added-line-decoration",
-              glyphMarginClassName:
-                i < originalLines.length
-                  ? "removed-line-glyph"
-                  : "added-line-glyph",
+              className: "added-line-decoration",
+              glyphMarginClassName: "added-line-glyph",
             },
           })
         }
       }
 
-      // Update editor content
+      // Store original content in the model
+      const model = editorRef.getModel()
+      if (model) {
+        ;(model as any).originalContent = originalCode
+      }
       editorRef.setValue(mergedCode)
 
       // Apply decorations
       const newDecorations = editorRef.createDecorationsCollection(decorations)
       setMergeDecorationsCollection(newDecorations)
     },
-    [editorRef, activeFileContent]
+    [editorRef]
   )
 
   // Generate widget effect