From 88abb29de99fc9a32e8fdf61a88b5c327bbc1b0f Mon Sep 17 00:00:00 2001
From: Kartik <kartik.labhshetwar@mem0.ai>
Date: Thu, 19 Mar 2026 18:18:11 +0530
Subject: [PATCH] fix: handle truncated code blocks in removeCodeBlocks
 function (#4421)

---
 mem0-ts/src/oss/src/prompts/index.ts          |  8 ++--
 .../src/oss/tests/remove-code-blocks.test.ts  | 38 +++++++++++++++++++
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/mem0-ts/src/oss/src/prompts/index.ts b/mem0-ts/src/oss/src/prompts/index.ts
index 03125cff1..c39ac25db 100644
--- a/mem0-ts/src/oss/src/prompts/index.ts
+++ b/mem0-ts/src/oss/src/prompts/index.ts
@@ -278,9 +278,7 @@ export function parseMessages(messages: string[]): string {
 }
 
 export function removeCodeBlocks(text: string): string {
-  // Extract content inside code fences instead of deleting it.
-  // The old regex /```[^`]*```/g replaced the entire block (including
-  // its content) with an empty string, so when an LLM returned JSON
-  // wrapped in ```json ... ``` the actual payload was discarded.
-  return text.replace(/```(?:\w+)?\n?([\s\S]*?)```/g, "$1").trim();
+  // Extract content inside code fences, handling both complete and
+  // truncated blocks (where the closing ``` never arrives).
+  return text.replace(/```(?:\w+)?\n?([\s\S]*?)(?:```|$)/g, "$1").trim();
 }
diff --git a/mem0-ts/src/oss/tests/remove-code-blocks.test.ts b/mem0-ts/src/oss/tests/remove-code-blocks.test.ts
index 439d6b468..4453bfa14 100644
--- a/mem0-ts/src/oss/tests/remove-code-blocks.test.ts
+++ b/mem0-ts/src/oss/tests/remove-code-blocks.test.ts
@@ -27,4 +27,42 @@ describe("removeCodeBlocks", () => {
     expect(removeCodeBlocks(input)).toContain('"facts"');
     expect(removeCodeBlocks(input)).not.toContain("```");
   });
+
+  // Truncated LLM response cases (issue #4401)
+  it("handles truncated code block missing closing fence", () => {
+    const input = '```json\n{"facts": ["hello"]}';
+    expect(removeCodeBlocks(input)).toBe('{"facts": ["hello"]}');
+  });
+
+  it("handles truncated code block with incomplete JSON", () => {
+    const input = '```json\n{"key": "value"';
+    expect(removeCodeBlocks(input)).toBe('{"key": "value"');
+  });
+
+  it("handles orphan trailing fence", () => {
+    const input = '{"result": true}\n```';
+    expect(removeCodeBlocks(input)).toBe('{"result": true}');
+  });
+
+  it("handles truncated block with bare fence (no language tag)", () => {
+    const input = '```\n{"facts": ["test"]}';
+    expect(removeCodeBlocks(input)).toBe('{"facts": ["test"]}');
+  });
+
+  it("handles complete block followed by truncated block", () => {
+    const input = '```json\n{"a":1}\n```\nsome text\n```python\nprint("hi")';
+    const result = removeCodeBlocks(input);
+    expect(result).toContain('{"a":1}');
+    expect(result).toContain('print("hi")');
+    expect(result).not.toMatch(/^```/);
+  });
+
+  it("returns empty string for empty input", () => {
+    expect(removeCodeBlocks("")).toBe("");
+  });
+
+  it("handles CRLF line endings from LLM proxies", () => {
+    const input = '```json\r\n{"facts": ["hello"]}\r\n```';
+    expect(removeCodeBlocks(input)).toBe('{"facts": ["hello"]}');
+  });
 });