Merge pull request #1856 from codeflash-ai/fix/structured-error-output-subagent-mode

fix: output structured XML errors in subagent mode

Author: aseembits93

codeflash/code_utils/code_utils.py

@@ -17,7 +17,7 @@
 
 from codeflash.cli_cmds.console import logger, paneled_text
 from codeflash.code_utils.config_parser import find_pyproject_toml, get_all_closest_config_files
-from codeflash.lsp.helpers import is_LSP_enabled
+from codeflash.lsp.helpers import is_LSP_enabled, is_subagent_mode
 
 _INVALID_CHARS_NT = {"<", ">", ":", '"', "|", "?", "*"}
 
@@ -471,6 +471,11 @@ def exit_with_message(message: str, *, error_on_exit: bool = False) -> None:
     if is_LSP_enabled():
         logger.error(message)
         return
+    if is_subagent_mode():
+        from xml.sax.saxutils import escape
+
+        sys.stdout.write(f"<codeflash-error>{escape(message)}</codeflash-error>\n")
+        sys.exit(1 if error_on_exit else 0)
     paneled_text(message, panel_args={"style": "red"})
 
     sys.exit(1 if error_on_exit else 0)

tests/test_code_utils.py

@@ -8,6 +8,7 @@
 
 from codeflash.code_utils.code_utils import (
     cleanup_paths,
+    exit_with_message,
     file_name_from_test_module_name,
     file_path_from_module_name,
     get_all_function_names,
@@ -751,3 +752,33 @@ def __init__(self):
 """
     result = validate_python_code(code)
     assert result == code
+
+
+class TestExitWithMessageSubagent:
+    @patch("codeflash.code_utils.code_utils.is_subagent_mode", return_value=True)
+    def test_outputs_structured_xml_in_subagent_mode(self, _mock_subagent: MagicMock, capsys: pytest.CaptureFixture[str]) -> None:
+        with pytest.raises(SystemExit) as exc_info:
+            exit_with_message("Something went wrong", error_on_exit=True)
+        assert exc_info.value.code == 1
+        captured = capsys.readouterr()
+        assert "<codeflash-error>" in captured.out
+        assert "Something went wrong" in captured.out
+        assert "</codeflash-error>" in captured.out
+
+    @patch("codeflash.code_utils.code_utils.is_subagent_mode", return_value=True)
+    def test_escapes_xml_special_chars(self, _mock_subagent: MagicMock, capsys: pytest.CaptureFixture[str]) -> None:
+        with pytest.raises(SystemExit):
+            exit_with_message('File <foo> & "bar" not found', error_on_exit=True)
+        captured = capsys.readouterr()
+        assert "&lt;foo&gt;" in captured.out
+        assert "&amp;" in captured.out
+
+    @patch("codeflash.code_utils.code_utils.is_subagent_mode", return_value=False)
+    @patch("codeflash.code_utils.code_utils.is_LSP_enabled", return_value=False)
+    def test_no_xml_when_not_subagent(
+        self, _mock_lsp: MagicMock, _mock_subagent: MagicMock, capsys: pytest.CaptureFixture[str]
+    ) -> None:
+        with pytest.raises(SystemExit):
+            exit_with_message("Normal error", error_on_exit=True)
+        captured = capsys.readouterr()
+        assert "<codeflash-error>" not in captured.out