diff --git a/src/jupyter_runtime.py b/src/jupyter_runtime.py index ecefd9c..c5df8bb 100644 --- a/src/jupyter_runtime.py +++ b/src/jupyter_runtime.py @@ -535,6 +535,7 @@ PY if not command.strip(): raise JupyterRuntimeError('Remote command is empty.') marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__' + start_marker = f'__ZK_AGENT_START_{uuid.uuid4().hex}__' workspace = shlex.quote(self.binding.workspace_cwd) run_dir = f'{self.binding.workspace_cwd}/scratchpad/.run_pids' quoted_run_dir = shlex.quote(run_dir) @@ -544,6 +545,7 @@ PY wrapped = ( f'mkdir -p {workspace} {quoted_run_dir} && ' f'cd {workspace} && ' + f'printf "\\n{start_marker}\\n"; ' f'( setsid bash -lc {shlex.quote(inner_command)} ) & ' '__zk_pid=$!; ' f'printf "%s\\n" "$__zk_pid" > {pid_file}; ' @@ -581,6 +583,7 @@ PY return self._collect_until_marker( ws, marker=marker, + start_marker=start_marker, timeout_seconds=timeout_seconds, max_output_chars=max_output_chars, cancel_event=cancel_event, @@ -1079,6 +1082,7 @@ else: ws: Any, *, marker: str, + start_marker: str, timeout_seconds: float, max_output_chars: int, cancel_event: Any | None, @@ -1117,12 +1121,22 @@ else: pass return RemoteCommandResult( exit_code=124, - stdout=truncate_text(''.join(transcript), max_output_chars), + stdout=truncate_text( + extract_terminal_command_output( + clean_terminal_output(''.join(transcript)), + start_marker=start_marker, + ), + max_output_chars, + ), timed_out=True, ) + stdout = extract_terminal_command_output( + clean_terminal_output(''.join(transcript)), + start_marker=start_marker, + ) return RemoteCommandResult( exit_code=exit_code, - stdout=truncate_text(clean_terminal_output(''.join(transcript)), max_output_chars), + stdout=truncate_text(stdout, max_output_chars), ) def _remote_write_text_command(self, path: str, content: str) -> str: @@ -1279,6 +1293,16 @@ def clean_terminal_output(text: str) -> str: return '\n'.join(lines).strip('\r\n') +def extract_terminal_command_output(text: str, *, start_marker: str) -> str: + """Drop terminal echo/prompt text before the wrapper's real start marker.""" + if not start_marker: + return text + marker_index = text.rfind(start_marker) + if marker_index < 0: + return text + return text[marker_index + len(start_marker) :].lstrip('\r\n') + + def truncate_text(text: str, limit: int) -> str: if limit <= 0 or len(text) <= limit: return text diff --git a/tests/test_jupyter_runtime.py b/tests/test_jupyter_runtime.py new file mode 100644 index 0000000..b00d7c6 --- /dev/null +++ b/tests/test_jupyter_runtime.py @@ -0,0 +1,40 @@ +from __future__ import annotations + +import unittest + +from src.jupyter_runtime import ( + clean_terminal_output, + extract_terminal_command_output, +) + + +class TestJupyterRuntimeOutputFiltering(unittest.TestCase): + def test_extract_terminal_command_output_drops_wrapper_echo(self) -> None: + start_marker = '__ZK_AGENT_START_test__' + raw = ( + 'root@host:/workspace# mkdir -p /workspace && cd /workspace && ' + 'printf "\\n__ZK_AGENT_START_test__\\n"; ' + '( setsid bash -lc \'export A=1; echo real-output\' ) &\r\n' + f'{start_marker}\r\n' + 'real-output\r\n' + ) + + output = extract_terminal_command_output( + clean_terminal_output(raw), + start_marker=start_marker, + ) + + self.assertEqual(output, 'real-output') + + def test_extract_terminal_command_output_uses_last_marker(self) -> None: + start_marker = '__ZK_AGENT_START_test__' + output = extract_terminal_command_output( + f'echo text containing {start_marker}\n{start_marker}\nactual', + start_marker=start_marker, + ) + + self.assertEqual(output, 'actual') + + +if __name__ == '__main__': + unittest.main()