fix: handler for parsing multiple generated code blocks

This commit is contained in:
Fahad
2025-10-18 00:28:17 +04:00
parent 95e69a7cb2
commit f4c20d2a20
2 changed files with 88 additions and 20 deletions

View File

@@ -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"

View File

@@ -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)