Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion codeflash/code_utils/code_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {"<", ">", ":", '"', "|", "?", "*"}

Expand Down Expand Up @@ -458,6 +458,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)
Expand Down
31 changes: 31 additions & 0 deletions tests/test_code_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Loading