diff --git a/tests/test_chat_simple.py b/tests/test_chat_simple.py index af45f2f..3ff1564 100644 --- a/tests/test_chat_simple.py +++ b/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("") == 2 + assert "print('hello')" not in saved_content + assert saved_content.count("") == 1 + assert "print('hello')" in formatted assert str(saved_path) in formatted + def test_format_response_single_generated_code_block(self, tmp_path): + """Single 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" + "print('only-once')\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 "" 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\nprint('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 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 " 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" + "print('demo')\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" diff --git a/tools/chat.py b/tools/chat.py index 9d42a3d..50f5d9a 100644 --- a/tools/chat.py +++ b/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] - combined_block = "\n\n".join(blocks) + last_match = matches[-1] + block = last_match.group(0).strip() - remainder_parts: list[str] = [] - last_end = 0 - for match in matches: - start, end = match.span() - segment = text[last_end:start] - if segment: - remainder_parts.append(segment) - last_end = end - tail = text[last_end:] - if tail: - remainder_parts.append(tail) + # Merge the text before and after the final block while trimming excess whitespace + before = text[: last_match.start()] + after = text[last_match.end() :] + remainder = self._join_sections(before, after) - remainder = self._join_sections(*remainder_parts) - - return combined_block, remainder, len(blocks) + return block, remainder, len(matches) def _persist_generated_code_block(self, block: str, working_directory: str) -> Path: expanded = os.path.expanduser(working_directory)