diff --git a/src/ci_orchestrator/commands/template.py b/src/ci_orchestrator/commands/template.py index eeba90a..9bd6f5d 100644 --- a/src/ci_orchestrator/commands/template.py +++ b/src/ci_orchestrator/commands/template.py @@ -367,12 +367,26 @@ def _run_with_heartbeat( check: bool = True, heartbeat: int = 30, ) -> WinRmResult: - """Run *script* via *transport* and print a heartbeat line every *heartbeat* s. + """Run *script* and show its output live; fall back to a heartbeat line. - pypsrp's execute_ps is fully blocking — no streaming. The heartbeat - thread reassures the operator that the command is still running during - long-running steps (e.g. Windows Update, which can take 20-40 min). + Prefers :meth:`WinRmTransport.run_streaming` so long steps (toolchain + install, Windows Update — 20-40 min) print progress as it happens. The + ``CI_EXITCODE:`` marker line is suppressed from the live echo. If streaming + is unavailable/fails, fall back to the blocking run + periodic heartbeat. """ + def _echo_line(line: str) -> None: + if not line.strip().startswith("CI_EXITCODE:"): + click.echo(line) + + try: + return transport.run_streaming(script, _echo_line, check=check) + except TransportError: + # Connection/command errors must propagate — e.g. the VM rebooting + # mid-update is caught upstream and treated as exit 3010. + raise + except Exception as exc: # pypsrp API quirk — degrade to the blocking path + click.echo(f" [WARN] live streaming failed ({exc}); using heartbeat.", err=True) + result_holder: list[WinRmResult] = [] exc_holder: list[BaseException] = [] stop_event = threading.Event() diff --git a/src/ci_orchestrator/transport/winrm.py b/src/ci_orchestrator/transport/winrm.py index 09e361a..5ef4ff0 100644 --- a/src/ci_orchestrator/transport/winrm.py +++ b/src/ci_orchestrator/transport/winrm.py @@ -10,6 +10,7 @@ from __future__ import annotations import contextlib import time +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path from typing import TYPE_CHECKING, Any @@ -44,6 +45,19 @@ class WinRmResult: return self.returncode == 0 +def _info_text(record: Any) -> str: + """Extract the displayed text of a PSRP information (Write-Host) record.""" + data = getattr(record, "message_data", None) + return str(data) if data is not None else str(record) + + +def _drain(seq: list[Any], start: int, emit: Callable[[str], None]) -> int: + """Emit ``str()`` of every item in ``seq[start:]``; return the new length.""" + for item in seq[start:]: + emit(str(item)) + return len(seq) + + class WinRmTransport: """Thin wrapper around :class:`pypsrp.client.Client`. @@ -180,6 +194,60 @@ class WinRmTransport: raise TransportCommandError(result.returncode, result.stdout, result.stderr) return result + def run_streaming( + self, script: str, on_line: Callable[[str], None], *, check: bool = True + ) -> WinRmResult: + """Run *script* and stream its output/host/error lines live via *on_line*. + + Unlike :meth:`run` (which buffers and returns only on completion), this + uses pypsrp's incremental ``poll_invoke`` so long-running scripts (the + toolchain install, Windows Update) print progress as it happens. + """ + from pypsrp.complex_objects import PSInvocationState + from pypsrp.powershell import PowerShell, RunspacePool + + client = self._connect() + pool = RunspacePool(client.wsman) + pool.open() + out_lines: list[str] = [] + err_lines: list[str] = [] + + def emit_out(s: str) -> None: + out_lines.append(s) + on_line(s) + + def emit_err(s: str) -> None: + err_lines.append(s) + on_line(f"[stderr] {s}") + + try: + ps = PowerShell(pool) + ps.add_script(script) + ps.begin_invoke() + out_i = info_i = err_i = 0 + running = (PSInvocationState.NOT_STARTED, PSInvocationState.RUNNING) + while ps.state in running: + ps.poll_invoke() + out_i = _drain(ps.output, out_i, emit_out) + info_i = _drain( + [_info_text(r) for r in ps.streams.information], info_i, on_line + ) + err_i = _drain([str(e) for e in ps.streams.error], err_i, emit_err) + time.sleep(0.5) + ps.end_invoke() + _drain(ps.output, out_i, emit_out) + _drain([str(e) for e in ps.streams.error], err_i, emit_err) + returncode = 1 if ps.streams.error else 0 + finally: + with contextlib.suppress(Exception): + pool.close() + result = WinRmResult( + stdout="\n".join(out_lines), stderr="\n".join(err_lines), returncode=returncode + ) + if check and not result.ok: + raise TransportCommandError(result.returncode, result.stdout, result.stderr) + return result + def copy(self, local_path: str | Path, remote_path: str) -> None: """Upload a file to the guest via PSRP.""" client = self._connect() diff --git a/tests/python/test_commands_template.py b/tests/python/test_commands_template.py index 7e1495f..19f84d5 100644 --- a/tests/python/test_commands_template.py +++ b/tests/python/test_commands_template.py @@ -515,6 +515,15 @@ def test_prepare_win_explicit_toolchain_not_found(tmp_path: Path) -> None: def _patch_prepare(monkeypatch: pytest.MonkeyPatch, backend: MagicMock, transport: MagicMock) -> None: + # _run_with_heartbeat prefers transport.run_streaming; route it through the + # test's transport.run side_effect and feed stdout to the on_line callback. + if not hasattr(transport.run_streaming, "side_effect") or transport.run_streaming.side_effect is None: + def _stream(script: str = "", on_line: object = None, **kwargs: object) -> WinRmResult: + res = transport.run(script, **kwargs) + if callable(on_line) and isinstance(res, WinRmResult): + on_line(res.stdout) + return res + transport.run_streaming.side_effect = _stream monkeypatch.setattr(tmpl_mod, "load_config", lambda: _make_config()) monkeypatch.setattr(tmpl_mod, "WorkstationVmrunBackend", lambda **_: backend) monkeypatch.setattr(tmpl_mod, "WinRmTransport", lambda **_: transport) diff --git a/tests/python/test_winrm_transport.py b/tests/python/test_winrm_transport.py index 7c5a33e..5712510 100644 --- a/tests/python/test_winrm_transport.py +++ b/tests/python/test_winrm_transport.py @@ -193,3 +193,89 @@ def test_is_ready_true_when_run_succeeds(monkeypatch: pytest.MonkeyPatch) -> Non _install_fake_pypsrp(monkeypatch, fake) t = WinRmTransport("10.0.0.1", "user", "pwd") assert t.is_ready() is True + + +def _install_fake_streaming(monkeypatch: pytest.MonkeyPatch, *, with_error: bool = False) -> None: + """Inject fake pypsrp.powershell + complex_objects for run_streaming tests.""" + class _State: + NOT_STARTED, RUNNING, COMPLETED, FAILED, STOPPED, STOPPING, DISCONNECTED = range(7) + + class _Streams: + def __init__(self) -> None: + self.information: list[Any] = [] + self.error: list[Any] = [] + + class _FakePS: + def __init__(self, _pool: Any) -> None: + self.state = _State.NOT_STARTED + self.output: list[Any] = [] + self.streams = _Streams() + self._poll = 0 + + def add_script(self, script: str) -> None: + self.script = script + + def begin_invoke(self) -> None: + self.state = _State.RUNNING + + def poll_invoke(self, *_a: Any, **_k: Any) -> None: + self._poll += 1 + if self._poll == 1: + self.output.append("line1") + self.streams.information.append(types.SimpleNamespace(message_data="hostmsg")) + elif self._poll == 2: + self.output.append("line2") + if with_error: + self.streams.error.append("boom") + self.state = _State.COMPLETED + + def end_invoke(self) -> None: + pass + + class _FakePool: + def __init__(self, wsman: Any) -> None: + self.wsman = wsman + + def open(self) -> None: + pass + + def close(self) -> None: + pass + + co = types.ModuleType("pypsrp.complex_objects") + co.PSInvocationState = _State # type: ignore[attr-defined] + psmod = types.ModuleType("pypsrp.powershell") + psmod.PowerShell = _FakePS # type: ignore[attr-defined] + psmod.RunspacePool = _FakePool # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "pypsrp.complex_objects", co) + monkeypatch.setitem(sys.modules, "pypsrp.powershell", psmod) + monkeypatch.setattr("ci_orchestrator.transport.winrm.time.sleep", lambda _s: None) + + +def test_run_streaming_emits_lines(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeClient() + _install_fake_pypsrp(monkeypatch, fake) + _install_fake_streaming(monkeypatch) + t = WinRmTransport("10.0.0.1", "u", "p") + lines: list[str] = [] + res = t.run_streaming("Write-Output 1", lines.append) + assert lines == ["line1", "hostmsg", "line2"] + assert res.ok and res.stdout == "line1\nline2" + + +def test_run_streaming_raises_on_error(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeClient() + _install_fake_pypsrp(monkeypatch, fake) + _install_fake_streaming(monkeypatch, with_error=True) + t = WinRmTransport("10.0.0.1", "u", "p") + with pytest.raises(TransportCommandError): + t.run_streaming("x", lambda _s: None) + + +def test_run_streaming_check_false(monkeypatch: pytest.MonkeyPatch) -> None: + fake = _FakeClient() + _install_fake_pypsrp(monkeypatch, fake) + _install_fake_streaming(monkeypatch, with_error=True) + t = WinRmTransport("10.0.0.1", "u", "p") + res = t.run_streaming("x", lambda _s: None, check=False) + assert not res.ok and "boom" in res.stderr