chore: import upstream snapshot with attribution

sdks/python/tests/unit/runner/test_supervisor.py

@@ -0,0 +1,579 @@
+import os
+import signal
+import subprocess
+import sys
+import threading
+import time
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+
+from opik.cli.local_runner.pairing import RunnerType
+from opik.rest_api.core.api_error import ApiError
+from opik.rest_api.types.local_runner_heartbeat_response import (
+    LocalRunnerHeartbeatResponse,
+)
+from opik.runner.supervisor import Supervisor
+
+
+_SENTINEL = object()
+
+
+def _make_supervisor(
+    command=_SENTINEL,
+    env=None,
+    repo_root=None,
+    runner_id="runner-1",
+    api=None,
+    watch=None,
+    runner_type=RunnerType.ENDPOINT,
+) -> Supervisor:
+    if command is _SENTINEL:
+        command = [sys.executable, "-c", "import time; time.sleep(60)"]
+    if env is None:
+        env = dict(os.environ)
+    if repo_root is None:
+        repo_root = Path.cwd()
+    if api is None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+        api.runners.next_bridge_commands.return_value = MagicMock(commands=[])
+    return Supervisor(
+        command=command,
+        env=env,
+        repo_root=repo_root,
+        runner_id=runner_id,
+        api=api,
+        watch=watch,
+        runner_type=runner_type,
+        project_name="test-project",
+        workspace=None,
+        heartbeat_interval_seconds=0.05,
+        graceful_timeout_seconds=0.1,
+        main_loop_tick_seconds=0.05,
+    )
+
+
+class TestStartChild:
+    def test_launches_process(self, tmp_path: Path) -> None:
+        marker = tmp_path / "started"
+        sup = _make_supervisor(
+            command=[sys.executable, "-c", f"open('{marker}', 'w').write('ok')"],
+            repo_root=tmp_path,
+        )
+        child = sup._start_child()
+        child.wait(timeout=5)
+        assert marker.read_text() == "ok"
+
+    def test_captures_output_via_pipe(self) -> None:
+        sup = _make_supervisor()
+        with patch("opik.runner.supervisor.subprocess.Popen") as mock_popen:
+            mock_proc = MagicMock()
+            mock_proc.stdout = MagicMock()
+            mock_proc.stdout.readline = MagicMock(return_value=b"")
+            mock_proc.stderr = MagicMock()
+            mock_proc.stderr.readline = MagicMock(return_value=b"")
+            mock_popen.return_value = mock_proc
+            sup._start_child()
+            call_kwargs = mock_popen.call_args.kwargs
+            assert call_kwargs.get("stdout") == subprocess.PIPE
+            assert call_kwargs.get("stderr") == subprocess.PIPE
+
+
+class TestStopChild:
+    def test_sigterm_then_wait(self) -> None:
+        sup = _make_supervisor()
+        with sup._child_lock:
+            sup._child = sup._start_child()
+        assert sup._child.poll() is None
+        sup._stop_child()
+        assert sup._child is None
+
+    def test_sigkill_after_timeout(self) -> None:
+        sup = _make_supervisor(
+            command=[
+                sys.executable,
+                "-c",
+                "import signal,time; signal.signal(signal.SIGTERM, signal.SIG_IGN); time.sleep(60)",
+            ],
+        )
+        with sup._child_lock:
+            sup._child = sup._start_child()
+        time.sleep(0.2)
+        sup._stop_child(graceful_timeout=0.1)
+        assert sup._child is None
+
+    def test_already_dead__no_error(self) -> None:
+        sup = _make_supervisor(
+            command=[sys.executable, "-c", "pass"],
+        )
+        with sup._child_lock:
+            sup._child = sup._start_child()
+        sup._child.wait(timeout=5)
+        sup._stop_child()
+        assert sup._child is None
+
+
+class TestRestart:
+    def test_stops_and_starts(self) -> None:
+        sup = _make_supervisor()
+        with sup._child_lock:
+            sup._child = sup._start_child()
+        old_pid = sup._child.pid
+
+        sup._restart_child("test reason")
+
+        assert sup._child is not None
+        assert sup._child.pid != old_pid
+
+        sup._stop_child()
+
+    def test_debounce(self) -> None:
+        sup = _make_supervisor()
+        with sup._child_lock:
+            sup._child = sup._start_child()
+
+        start_count = 0
+        original_start = sup._start_child
+
+        def counting_start():
+            nonlocal start_count
+            start_count += 1
+            return original_start()
+
+        sup._start_child = counting_start
+
+        # Three rapid calls — only the first should actually restart
+        for i in range(3):
+            sup._restart_child(f"trigger {i}")
+
+        assert start_count == 1
+
+        sup._stop_child()
+
+
+class TestChildExit:
+    def test_restarts_on_nonzero_exit_if_stable(self) -> None:
+        sup = _make_supervisor()
+        sup._on_child_restart = MagicMock()
+        sup._on_error = MagicMock()
+
+        # Guard stable, call _handle_child_exit with exit code 1
+        # (Guard tracks crashes in this method)
+        should_restart = sup._handle_child_exit(exit_code=1)
+
+        assert should_restart is True
+        sup._on_child_restart.assert_called_once_with("agent process has failed")
+        sup._on_error.assert_not_called()
+
+    def test_triggers_error_on_crash_loop(self) -> None:
+        sup = _make_supervisor()
+        sup._on_child_restart = MagicMock()
+        sup._on_error = MagicMock()
+        sup._guard._max_crashes = 2
+        sup._guard._window_seconds = 60.0
+
+        # Pre-record 2 crashes to make guard unstable
+        sup._guard.record_crash()
+        sup._guard.record_crash()
+        assert not sup._guard.is_stable()
+
+        # Call _handle_child_exit — should trigger error, not restart
+        should_restart = sup._handle_child_exit(exit_code=1)
+
+        assert should_restart is False
+        sup._on_error.assert_called_once()
+        sup._on_child_restart.assert_not_called()
+
+
+class TestErrorCallback:
+    def test_on_error_called_on_crash_loop(self) -> None:
+        error_messages = []
+        sup = _make_supervisor(
+            command=[sys.executable, "-c", "import sys; sys.exit(1)"],
+        )
+        sup._on_error = lambda msg: error_messages.append(msg)
+        sup._guard._max_crashes = 1
+        sup._guard._window_seconds = 60.0
+
+        # Exhaust stability guard so next crash triggers crash-loop
+        sup._guard.record_crash()
+
+        # Record another crash — guard should now be unstable
+        sup._guard.record_crash()
+        assert not sup._guard.is_stable()
+
+        # Simulate the on_error path directly
+        if not sup._guard.is_stable():
+            if sup._on_error:
+                sup._on_error("Crash loop detected — waiting for file change to retry")
+
+        assert len(error_messages) == 1
+        assert "Crash loop" in error_messages[0]
+
+    def test_on_child_restart_called_on_crash_restart(self) -> None:
+        restart_reasons = []
+        sup = _make_supervisor(
+            command=[sys.executable, "-c", "import sys; sys.exit(1)"],
+        )
+        sup._on_child_restart = lambda reason: restart_reasons.append(reason)
+
+        # Guard is stable — should trigger on_child_restart
+        assert sup._guard.is_stable()
+
+        if sup._on_child_restart:
+            sup._on_child_restart("agent process has failed")
+
+        assert len(restart_reasons) == 1
+        assert "agent process has failed" in restart_reasons[0]
+
+    def test_on_error_not_called_when_stable(self) -> None:
+        error_messages = []
+        sup = _make_supervisor(
+            command=[sys.executable, "-c", "import sys; sys.exit(1)"],
+        )
+        sup._on_error = lambda msg: error_messages.append(msg)
+
+        # Guard is stable — on_error should not fire
+        sup._guard.record_crash()
+        assert sup._guard.is_stable()
+
+        if not sup._guard.is_stable():
+            if sup._on_error:
+                sup._on_error("should not happen")
+
+        assert len(error_messages) == 0
+
+
+class TestShutdown:
+    def test_stops_all(self) -> None:
+        sup = _make_supervisor(watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.2)
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+        assert sup._child is None
+
+    def test_waits_for_child(self) -> None:
+        sup = _make_supervisor(watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.2)
+        sup._shutdown_event.set()
+        t.join(timeout=15)
+
+        assert not t.is_alive()
+
+
+class TestHeartbeat:
+    def test_sends_heartbeat__endpoint__reports_jobs(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+
+        sup = _make_supervisor(api=api, runner_type=RunnerType.ENDPOINT)
+
+        t = threading.Thread(target=sup._heartbeat_loop, daemon=True)
+        t.start()
+
+        time.sleep(0.1)
+        sup._shutdown_event.set()
+        t.join(timeout=5)
+
+        api.runners.heartbeat.assert_called()
+        call_kwargs = api.runners.heartbeat.call_args.kwargs
+        assert call_kwargs["capabilities"] == ["jobs"]
+
+    def test_sends_heartbeat__connect__reports_bridge(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+
+        sup = _make_supervisor(api=api, runner_type=RunnerType.CONNECT)
+
+        t = threading.Thread(target=sup._heartbeat_loop, daemon=True)
+        t.start()
+
+        time.sleep(0.1)
+        sup._shutdown_event.set()
+        t.join(timeout=5)
+
+        api.runners.heartbeat.assert_called()
+        call_kwargs = api.runners.heartbeat.call_args.kwargs
+        assert call_kwargs["capabilities"] == ["bridge"]
+
+    def test_410__shuts_down(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.side_effect = ApiError(status_code=410, body=None)
+
+        sup = _make_supervisor(api=api)
+
+        t = threading.Thread(target=sup._heartbeat_loop, daemon=True)
+        t.start()
+        t.join(timeout=5)
+
+        assert sup._shutdown_event.is_set()
+
+
+class TestDisconnectNotification:
+    def test_disconnect_called_on_shutdown(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+
+        sup = _make_supervisor(api=api, watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.2)
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+        api.runners.disconnect_runner.assert_called_once()
+        args, _ = api.runners.disconnect_runner.call_args
+        assert args[0] == sup._runner_id
+
+    def test_disconnect_error_is_non_fatal(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+        api.runners.disconnect_runner.side_effect = ApiError(status_code=503, body=None)
+
+        sup = _make_supervisor(api=api, watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.2)
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+        assert not t.is_alive()
+        api.runners.disconnect_runner.assert_called_once()
+
+    def test_disconnect_called_in_standalone_mode(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+
+        sup = _make_supervisor(command=None, api=api, watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.1)
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+        api.runners.disconnect_runner.assert_called_once()
+
+    def test_signal_handler_sets_shutdown_event(self) -> None:
+        # Without this, "if event set, disconnect fires" tests don't prove that SIGINT/SIGTERM
+        # actually drive the event. The CLI runs on the main thread in production, so
+        # signal.signal() installs cleanly there; this exercises that path directly.
+        sup = _make_supervisor()
+
+        original_sigint = signal.getsignal(signal.SIGINT)
+        original_sigterm = signal.getsignal(signal.SIGTERM)
+        try:
+            sup._install_signal_handlers()
+
+            installed_sigint = signal.getsignal(signal.SIGINT)
+            installed_sigterm = signal.getsignal(signal.SIGTERM)
+            assert callable(installed_sigint)
+            assert callable(installed_sigterm)
+            assert installed_sigint is installed_sigterm
+
+            assert not sup._shutdown_event.is_set()
+            installed_sigint(signal.SIGINT, None)
+            assert sup._shutdown_event.is_set()
+        finally:
+            signal.signal(signal.SIGINT, original_sigint)
+            signal.signal(signal.SIGTERM, original_sigterm)
+
+    def test_disconnect_called_on_heartbeat_410_eviction(self) -> None:
+        # A 410 from the heartbeat loop sets _shutdown_event from a background thread.
+        # The main loop must still tear down via the finally block, including disconnect.
+        api = MagicMock()
+        api.runners.heartbeat.side_effect = ApiError(status_code=410, body=None)
+
+        sup = _make_supervisor(api=api, watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+        t.join(timeout=10)
+
+        assert not t.is_alive()
+        api.runners.disconnect_runner.assert_called_once()
+
+
+# Subprocess driver used by TestSignalDeliveryEndToEnd. Lives at module scope so its
+# source can be passed via `python -c`. Runs Supervisor in standalone mode (no child
+# process) so the test isolates the shutdown path, and writes a marker file from a
+# fake disconnect_runner so the parent can verify the call happened.
+_SIGNAL_INTEGRATION_SCRIPT = """
+import json
+import sys
+from pathlib import Path
+
+ready_path = Path(sys.argv[1])
+marker_path = Path(sys.argv[2])
+
+from opik.cli.local_runner.pairing import RunnerType
+from opik.rest_api.types.local_runner_heartbeat_response import LocalRunnerHeartbeatResponse
+from opik.runner.supervisor import Supervisor
+
+
+class _FakeRunners:
+    def heartbeat(self, *_a, **_kw):
+        return LocalRunnerHeartbeatResponse()
+
+    def patch_checklist(self, *_a, **_kw):
+        return None
+
+    def disconnect_runner(self, runner_id, *, request_options=None):
+        marker_path.write_text(json.dumps({"runner_id": runner_id}))
+
+
+class _FakeApi:
+    runners = _FakeRunners()
+
+
+sup = Supervisor(
+    command=None,
+    env={},
+    repo_root=Path.cwd(),
+    runner_id="runner-xyz",
+    api=_FakeApi(),
+    runner_type=RunnerType.ENDPOINT,
+    project_name="signal-integration",
+    workspace=None,
+    heartbeat_interval_seconds=60,
+    graceful_timeout_seconds=0.1,
+    main_loop_tick_seconds=0.05,
+)
+
+# Signal readiness only after the signal handlers are in place. We piggyback on
+# the file-watcher branch: when watch=False (standalone mode), run() reaches the
+# wait() after installing handlers. Touch the ready file just before that by
+# wrapping _install_signal_handlers.
+_original_install = sup._install_signal_handlers
+
+
+def _install_and_signal_ready():
+    _original_install()
+    ready_path.write_text("ready")
+
+
+sup._install_signal_handlers = _install_and_signal_ready
+sup.run()
+"""
+
+
+class TestSignalDeliveryEndToEnd:
+    # Mock-driven tests above cover the wiring; this exercises the real OS-signal path
+    # by spawning a subprocess (so the supervisor runs on a real main thread) and
+    # checking that `disconnect_runner` actually fires before exit.
+
+    def _run_subprocess_with_signal(self, sig: signal.Signals, tmp_path: Path) -> Path:
+        ready = tmp_path / "ready"
+        marker = tmp_path / "disconnect-called.json"
+
+        proc = subprocess.Popen(
+            [sys.executable, "-c", _SIGNAL_INTEGRATION_SCRIPT, str(ready), str(marker)],
+            env={**os.environ},
+        )
+        try:
+            deadline = time.monotonic() + 10
+            while not ready.exists():
+                if time.monotonic() > deadline:
+                    proc.kill()
+                    raise AssertionError("Supervisor subprocess never became ready")
+                if proc.poll() is not None:
+                    raise AssertionError(
+                        f"Supervisor exited early (code={proc.returncode}) before ready"
+                    )
+                time.sleep(0.05)
+
+            proc.send_signal(sig)
+            exit_code = proc.wait(timeout=10)
+        finally:
+            if proc.poll() is None:
+                proc.kill()
+                proc.wait(timeout=5)
+
+        assert exit_code == 0, f"Subprocess exited with {exit_code}"
+        assert marker.exists(), "disconnect_runner was not called"
+        return marker
+
+    def test_sigint_triggers_disconnect(self, tmp_path: Path) -> None:
+        self._run_subprocess_with_signal(signal.SIGINT, tmp_path)
+
+    def test_sigterm_triggers_disconnect(self, tmp_path: Path) -> None:
+        marker = self._run_subprocess_with_signal(signal.SIGTERM, tmp_path)
+        # Sanity: the marker captures the runner_id we passed in, confirming we hit
+        # the real disconnect path and not some other exit code path.
+        import json
+
+        payload = json.loads(marker.read_text())
+        assert payload["runner_id"] == "runner-xyz"
+
+
+class TestStandaloneMode:
+    def test_no_command__runs_without_child(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+
+        sup = _make_supervisor(command=None, api=api, watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.1)
+
+        assert sup._child is None
+
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+    def test_no_command__sends_checklist_with_null_command(self) -> None:
+        api = MagicMock()
+        api.runners.heartbeat.return_value = LocalRunnerHeartbeatResponse()
+
+        sup = _make_supervisor(command=None, api=api, watch=False)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.1)
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+        api.runners.patch_checklist.assert_called()
+        checklist = api.runners.patch_checklist.call_args.kwargs.get(
+            "request"
+        ) or api.runners.patch_checklist.call_args[1].get("request")
+        assert checklist["command"] is None
+
+
+class TestBridgeIntegration:
+    def test_bridge_loop_runs(self) -> None:
+        sup = _make_supervisor(watch=False, runner_type=RunnerType.CONNECT)
+
+        t = threading.Thread(target=sup.run, daemon=True)
+        t.start()
+
+        time.sleep(0.1)
+
+        bridge_alive = False
+        for thread in threading.enumerate():
+            if thread.name == "bridge-poll":
+                bridge_alive = True
+                break
+
+        sup._shutdown_event.set()
+        t.join(timeout=10)
+
+        assert bridge_alive