diff --git a/open_notebook/graphs/chat.py b/open_notebook/graphs/chat.py
index 810d29f..442396a 100644
--- a/open_notebook/graphs/chat.py
+++ b/open_notebook/graphs/chat.py
@@ -72,7 +72,7 @@ def call_model_with_messages(state: ThreadState, config: RunnableConfig) -> dict
     # Clean thinking content from AI response (e.g., <think>...</think> tags)
     content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content)
     cleaned_content = clean_thinking_content(content)
-    cleaned_message = AIMessage(content=cleaned_content)
+    cleaned_message = ai_message.model_copy(update={"content": cleaned_content})
 
     return {"messages": cleaned_message}
 
diff --git a/open_notebook/graphs/source_chat.py b/open_notebook/graphs/source_chat.py
index a173a56..99217a8 100644
--- a/open_notebook/graphs/source_chat.py
+++ b/open_notebook/graphs/source_chat.py
@@ -159,7 +159,7 @@ def call_model_with_source_context(
     # Clean thinking content from AI response (e.g., <think>...</think> tags)
     content = ai_message.content if isinstance(ai_message.content, str) else str(ai_message.content)
     cleaned_content = clean_thinking_content(content)
-    cleaned_message = AIMessage(content=cleaned_content)
+    cleaned_message = ai_message.model_copy(update={"content": cleaned_content})
 
     # Update state with context information
     return {
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 6052724..65e12e1 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -94,6 +94,14 @@ class TestTextUtilities:
         assert thinking == ""
         assert cleaned == "Just regular content"
 
+    def test_parse_thinking_content_malformed_no_open_tag(self):
+        """Test parsing malformed output where opening <think> tag is missing."""
+        content = "Some thinking content</think>Here is my answer"
+        thinking, cleaned = parse_thinking_content(content)
+
+        assert thinking == "Some thinking content"
+        assert cleaned == "Here is my answer"
+
     def test_parse_thinking_content_invalid_input(self):
         """Test parsing with invalid input types."""
         # Non-string input