diff --git a/src/codegen/extensions/tools/relace_edit.py b/src/codegen/extensions/tools/relace_edit.py
index 52111a0ac..0f7637bb8 100644
--- a/src/codegen/extensions/tools/relace_edit.py
+++ b/src/codegen/extensions/tools/relace_edit.py
@@ -2,7 +2,7 @@
 
 import difflib
 import os
-from typing import TYPE_CHECKING, ClassVar, Optional
+from typing import TYPE_CHECKING, ClassVar
 
 import requests
 from langchain_core.messages import ToolMessage
@@ -23,15 +23,15 @@ class RelaceEditObservation(Observation):
     filepath: str = Field(
         description="Path to the edited file",
     )
-    diff: Optional[str] = Field(
+    diff: str | None = Field(
         default=None,
         description="Unified diff showing the changes made",
     )
-    new_content: Optional[str] = Field(
+    new_content: str | None = Field(
         default=None,
         description="New content with line numbers",
     )
-    line_count: Optional[int] = Field(
+    line_count: int | None = Field(
         default=None,
         description="Total number of lines in file",
     )
@@ -135,7 +135,7 @@ def apply_relace_edit(api_key: str, initial_code: str, edit_snippet: str, stream
         raise Exception(msg)
 
 
-def relace_edit(codebase: Codebase, filepath: str, edit_snippet: str, api_key: Optional[str] = None) -> RelaceEditObservation:
+def relace_edit(codebase: Codebase, filepath: str, edit_snippet: str, api_key: str | None = None) -> RelaceEditObservation:
     """Edit a file using the Relace Instant Apply API.
 
     Args:
@@ -176,6 +176,8 @@ def relace_edit(codebase: Codebase, filepath: str, edit_snippet: str, api_key: O
     # Apply the edit using Relace API
     try:
         merged_code = apply_relace_edit(api_key, original_content, edit_snippet)
+        if original_content.endswith("\n") and not merged_code.endswith("\n"):
+            merged_code += "\n"
     except Exception as e:
         return RelaceEditObservation(
             status="error",