Filter Jupyter terminal wrapper output
This commit is contained in:
+26
-2
@@ -535,6 +535,7 @@ PY
|
|||||||
if not command.strip():
|
if not command.strip():
|
||||||
raise JupyterRuntimeError('Remote command is empty.')
|
raise JupyterRuntimeError('Remote command is empty.')
|
||||||
marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__'
|
marker = f'__ZK_AGENT_EXIT_{uuid.uuid4().hex}__'
|
||||||
|
start_marker = f'__ZK_AGENT_START_{uuid.uuid4().hex}__'
|
||||||
workspace = shlex.quote(self.binding.workspace_cwd)
|
workspace = shlex.quote(self.binding.workspace_cwd)
|
||||||
run_dir = f'{self.binding.workspace_cwd}/scratchpad/.run_pids'
|
run_dir = f'{self.binding.workspace_cwd}/scratchpad/.run_pids'
|
||||||
quoted_run_dir = shlex.quote(run_dir)
|
quoted_run_dir = shlex.quote(run_dir)
|
||||||
@@ -544,6 +545,7 @@ PY
|
|||||||
wrapped = (
|
wrapped = (
|
||||||
f'mkdir -p {workspace} {quoted_run_dir} && '
|
f'mkdir -p {workspace} {quoted_run_dir} && '
|
||||||
f'cd {workspace} && '
|
f'cd {workspace} && '
|
||||||
|
f'printf "\\n{start_marker}\\n"; '
|
||||||
f'( setsid bash -lc {shlex.quote(inner_command)} ) & '
|
f'( setsid bash -lc {shlex.quote(inner_command)} ) & '
|
||||||
'__zk_pid=$!; '
|
'__zk_pid=$!; '
|
||||||
f'printf "%s\\n" "$__zk_pid" > {pid_file}; '
|
f'printf "%s\\n" "$__zk_pid" > {pid_file}; '
|
||||||
@@ -581,6 +583,7 @@ PY
|
|||||||
return self._collect_until_marker(
|
return self._collect_until_marker(
|
||||||
ws,
|
ws,
|
||||||
marker=marker,
|
marker=marker,
|
||||||
|
start_marker=start_marker,
|
||||||
timeout_seconds=timeout_seconds,
|
timeout_seconds=timeout_seconds,
|
||||||
max_output_chars=max_output_chars,
|
max_output_chars=max_output_chars,
|
||||||
cancel_event=cancel_event,
|
cancel_event=cancel_event,
|
||||||
@@ -1079,6 +1082,7 @@ else:
|
|||||||
ws: Any,
|
ws: Any,
|
||||||
*,
|
*,
|
||||||
marker: str,
|
marker: str,
|
||||||
|
start_marker: str,
|
||||||
timeout_seconds: float,
|
timeout_seconds: float,
|
||||||
max_output_chars: int,
|
max_output_chars: int,
|
||||||
cancel_event: Any | None,
|
cancel_event: Any | None,
|
||||||
@@ -1117,12 +1121,22 @@ else:
|
|||||||
pass
|
pass
|
||||||
return RemoteCommandResult(
|
return RemoteCommandResult(
|
||||||
exit_code=124,
|
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,
|
timed_out=True,
|
||||||
)
|
)
|
||||||
|
stdout = extract_terminal_command_output(
|
||||||
|
clean_terminal_output(''.join(transcript)),
|
||||||
|
start_marker=start_marker,
|
||||||
|
)
|
||||||
return RemoteCommandResult(
|
return RemoteCommandResult(
|
||||||
exit_code=exit_code,
|
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:
|
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')
|
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:
|
def truncate_text(text: str, limit: int) -> str:
|
||||||
if limit <= 0 or len(text) <= limit:
|
if limit <= 0 or len(text) <= limit:
|
||||||
return text
|
return text
|
||||||
|
|||||||
@@ -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()
|
||||||
Reference in New Issue
Block a user