From 5f5b6698dbe11e7fcc016be2889900fb7dcb9bab Mon Sep 17 00:00:00 2001
From: Dat <tandat198@gmail.com>
Date: Mon, 14 Sep 2026 14:37:48 +0700
Subject: [PATCH] fix: prefer full-file block and handle nested fences in
 extract_code_block

Co-authored-by: CommandCodeBot <noreply@commandcode.ai>
---
 ownbench/runner.py    | 27 ++++++++++++++++++---
 tests/test_extract.py | 56 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_extract.py

diff --git a/ownbench/runner.py b/ownbench/runner.py
index b577304..67a592e 100644
--- a/ownbench/runner.py
+++ b/ownbench/runner.py
@@ -41,10 +41,31 @@ class RunResult:
     error: str | None = None
 
 
+_OPEN_FENCE = re.compile(r"^\s*```[^`]*$")
+_CLOSE_FENCE = re.compile(r"^```\s*$")
+
+
 def extract_code_block(text: str) -> str | None:
-    """Last fenced code block in *text*, or None."""
-    blocks = re.findall(r"```[^\n]*\n(.*?)```", text, re.DOTALL)
-    return blocks[-1] if blocks else None
+    """Largest top-level fenced code block in *text*, or None.
+
+    Line-based scan: an opening fence may be indented and carry an info string,
+    while a closing fence must sit at column 0. Picking the longest block keeps
+    the full file over a trailing snippet, and a `.*?`-free scan means a fence
+    embedded in indented content (e.g. inside a docstring) cannot truncate it.
+    """
+    blocks: list[str] = []
+    current: list[str] | None = None
+    for line in text.split("\n"):
+        if current is None:
+            if _OPEN_FENCE.match(line):
+                current = []
+        elif _CLOSE_FENCE.match(line):
+            body = "\n".join(current)
+            blocks.append(body + "\n" if current else body)
+            current = None
+        else:
+            current.append(line)
+    return max(blocks, key=len) if blocks else None
 
 
 def chat_completion(
diff --git a/tests/test_extract.py b/tests/test_extract.py
new file mode 100644
index 0000000..724090c
--- /dev/null
+++ b/tests/test_extract.py
@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from ownbench.runner import extract_code_block
+
+
+def test_trailing_snippet_ignored():
+    full_file = "\n".join(
+        [
+            "def divide(a, b):",
+            "    if b == 0:",
+            '        raise ValueError("division by zero")',
+            "    return a / b",
+        ]
+    )
+    reply = (
+        "```py\n"
+        f"{full_file}\n"
+        "```\n"
+        "\n"
+        "Why this works:\n"
+        "```py\n"
+        "b == 0  # note\n"
+        "```\n"
+    )
+
+    out = extract_code_block(reply)
+    assert out is not None
+    assert "def divide" in out
+    assert "raise ValueError" in out
+    assert out != "b == 0  # note"
+
+
+def test_nested_fence_in_docstring():
+    body = "\n".join(
+        [
+            "def f():",
+            '    """',
+            "    ```",
+            "    example",
+            "    ```",
+            '    """',
+            "    return 1",
+        ]
+    )
+    reply = "```py\n" + body + "\n```\n"
+
+    out = extract_code_block(reply)
+    assert out is not None
+    assert "def f" in out
+    assert "example" in out
+
+
+def test_regression():
+    reply = "```py\ndef f():\n    return 1\n```\n"
+    assert extract_code_block(reply) == "def f():\n    return 1\n"
+    assert extract_code_block("no code here, just prose.") is None