fix: handler for parsing multiple generated code blocks

tests/test_chat_simple.py

@@ -147,11 +147,86 @@ class TestChatTool:
         saved_path = tmp_path / "zen_generated.code"
         saved_content = saved_path.read_text(encoding="utf-8")
 
-        assert "print('hello')" in saved_content
         assert "print('world')" in saved_content
-        assert saved_content.count("<GENERATED-CODE>") == 2
+        assert "print('hello')" not in saved_content
+        assert saved_content.count("<GENERATED-CODE>") == 1
+        assert "<GENERATED-CODE>print('hello')" in formatted
         assert str(saved_path) in formatted
 
+    def test_format_response_single_generated_code_block(self, tmp_path):
+        """Single <GENERATED-CODE> block should be saved and removed from narrative."""
+        tool = ChatTool()
+        tool._model_context = SimpleNamespace(capabilities=SimpleNamespace(allow_code_generation=True))
+
+        response = (
+            "Intro text before code.\n"
+            "<GENERATED-CODE>print('only-once')</GENERATED-CODE>\n"
+            "Closing thoughts after code."
+        )
+
+        request = ChatRequest(prompt="Test", working_directory=str(tmp_path))
+
+        formatted = tool.format_response(response, request)
+
+        saved_path = tmp_path / "zen_generated.code"
+        saved_content = saved_path.read_text(encoding="utf-8")
+
+        assert "print('only-once')" in saved_content
+        assert "<GENERATED-CODE>" in saved_content
+        assert "print('only-once')" not in formatted
+        assert "Closing thoughts after code." in formatted
+
+    def test_format_response_ignores_unclosed_generated_code(self, tmp_path):
+        """Unclosed generated-code tags should be ignored to avoid accidental clipping."""
+        tool = ChatTool()
+        tool._model_context = SimpleNamespace(capabilities=SimpleNamespace(allow_code_generation=True))
+
+        response = "Intro text\n<GENERATED-CODE>print('oops')\nStill ongoing"
+
+        request = ChatRequest(prompt="Test", working_directory=str(tmp_path))
+
+        formatted = tool.format_response(response, request)
+
+        saved_path = tmp_path / "zen_generated.code"
+        assert not saved_path.exists()
+        assert "print('oops')" in formatted
+
+    def test_format_response_ignores_orphaned_closing_tag(self, tmp_path):
+        """Stray closing tags should not trigger extraction."""
+        tool = ChatTool()
+        tool._model_context = SimpleNamespace(capabilities=SimpleNamespace(allow_code_generation=True))
+
+        response = "Intro text\n</GENERATED-CODE> just text"
+
+        request = ChatRequest(prompt="Test", working_directory=str(tmp_path))
+
+        formatted = tool.format_response(response, request)
+
+        saved_path = tmp_path / "zen_generated.code"
+        assert not saved_path.exists()
+        assert "</GENERATED-CODE> just text" in formatted
+
+    def test_format_response_preserves_narrative_after_generated_code(self, tmp_path):
+        """Narrative content after generated code must remain intact in the formatted output."""
+        tool = ChatTool()
+        tool._model_context = SimpleNamespace(capabilities=SimpleNamespace(allow_code_generation=True))
+
+        response = (
+            "Summary before code.\n"
+            "<GENERATED-CODE>print('demo')</GENERATED-CODE>\n"
+            "### Follow-up\n"
+            "Further analysis and guidance after the generated snippet.\n"
+        )
+
+        request = ChatRequest(prompt="Test", working_directory=str(tmp_path))
+
+        formatted = tool.format_response(response, request)
+
+        assert "Summary before code." in formatted
+        assert "### Follow-up" in formatted
+        assert "Further analysis and guidance after the generated snippet." in formatted
+        assert "print('demo')" not in formatted
+
     def test_tool_name(self):
         """Test tool name is correct"""
         assert self.tool.get_name() == "chat"

tools/chat.py

@@ -28,9 +28,11 @@ from .simple.base import SimpleTool
 CHAT_FIELD_DESCRIPTIONS = {
     "prompt": (
         "Your question or idea for collaborative thinking. Provide detailed context, including your goal, what you've tried, and any specific challenges. "
-        "CRITICAL: To discuss code, use 'files' parameter instead of pasting code blocks here."
+        "WARNING: Large inline code must NOT be shared in prompt. Provide full-path to files on disk as separate parameter."
+    ),
+    "files": (
+        "Absolute file or folder paths for code context. Required whenever you reference source code—supply the FULL absolute path (do not shorten)."
     ),
-    "files": "Absolute file or folder paths for code context.",
     "images": "Image paths (absolute) or base64 strings for optional visual context.",
     "working_directory": (
         "Absolute directory path where generated code artifacts are stored. The directory must already exist."
@@ -310,24 +312,15 @@ class ChatTool(SimpleTool):
         if not matches:
             return None, text, 0
 
-        blocks = [match.group(0).strip() for match in matches]
+        last_match = matches[-1]
-        combined_block = "\n\n".join(blocks)
+        block = last_match.group(0).strip()
 
-        remainder_parts: list[str] = []
+        # Merge the text before and after the final block while trimming excess whitespace
-        last_end = 0
+        before = text[: last_match.start()]
-        for match in matches:
+        after = text[last_match.end() :]
-            start, end = match.span()
+        remainder = self._join_sections(before, after)
-            segment = text[last_end:start]
-            if segment:
-                remainder_parts.append(segment)
-            last_end = end
-        tail = text[last_end:]
-        if tail:
-            remainder_parts.append(tail)
 
-        remainder = self._join_sections(*remainder_parts)
+        return block, remainder, len(matches)
-
-        return combined_block, remainder, len(blocks)
 
     def _persist_generated_code_block(self, block: str, working_directory: str) -> Path:
         expanded = os.path.expanduser(working_directory)