import asyncio import threading import time from unittest.mock import MagicMock import pytest from opik.api_objects.prompt import mask_context as prompt_mask_context_module from opik.rest_api.core.api_error import ApiError from opik.rest_api.types.local_runner_heartbeat_response import ( LocalRunnerHeartbeatResponse, ) from opik.rest_api.types.local_runner_job import LocalRunnerJob from opik.runner import in_process_loop, registry @pytest.fixture def mock_api(): api = MagicMock() api.runners.heartbeat = MagicMock( return_value=LocalRunnerHeartbeatResponse(cancelled_job_ids=[]) ) api.runners.next_job = MagicMock(return_value=None) api.runners.report_job_result = MagicMock() api.runners.register_agents = MagicMock() return api @pytest.fixture def shutdown_event(): return threading.Event() @pytest.fixture def loop(mock_api, shutdown_event): return in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, heartbeat_interval_seconds=100, poll_idle_interval_seconds=0.01, ) class TestHeartbeatLoop: def test_heartbeat__410__triggers_shutdown(self, mock_api, shutdown_event, loop): loop._heartbeat_interval_seconds = 0.05 mock_api.runners.heartbeat.side_effect = ApiError(status_code=410) t = threading.Thread(target=loop._heartbeat_loop) t.start() t.join(timeout=5) assert shutdown_event.is_set() def test_heartbeat__cancelled_ids__tracked(self, mock_api, shutdown_event, loop): loop._heartbeat_interval_seconds = 0.05 call_count = 0 def heartbeat_side_effect(runner_id, **kwargs): nonlocal call_count call_count += 1 if call_count == 1: return LocalRunnerHeartbeatResponse(cancelled_job_ids=["j-1"]) shutdown_event.set() return LocalRunnerHeartbeatResponse(cancelled_job_ids=[]) mock_api.runners.heartbeat.side_effect = heartbeat_side_effect t = threading.Thread(target=loop._heartbeat_loop) t.start() t.join(timeout=5) assert "j-1" in loop._cancelled_jobs class TestPollLoop: def test_poll__no_jobs__polls_repeatedly(self, mock_api, shutdown_event, loop): call_count = 0 def side_effect(runner_id): nonlocal call_count call_count += 1 if call_count >= 3: shutdown_event.set() return None mock_api.runners.next_job.side_effect = side_effect loop._loop = asyncio.new_event_loop() t = threading.Thread(target=loop._poll_loop) t.start() t.join(timeout=5) assert call_count >= 3 def test_poll__error__backs_off(self, mock_api, shutdown_event): lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, heartbeat_interval_seconds=100, backoff_cap_seconds=0.01, ) lp._loop = asyncio.new_event_loop() call_count = 0 def side_effect(runner_id): nonlocal call_count call_count += 1 if call_count >= 2: shutdown_event.set() return None raise ConnectionError("fail") mock_api.runners.next_job.side_effect = side_effect t = threading.Thread(target=lp._poll_loop) t.start() t.join(timeout=5) assert call_count >= 2 def test_poll__job_available__queued(self, mock_api, shutdown_event, loop): job = LocalRunnerJob(id="j-1", agent_name="test", inputs={"q": "hi"}) call_count = 0 def side_effect(runner_id): nonlocal call_count call_count += 1 if call_count == 1: return job shutdown_event.set() return None mock_api.runners.next_job.side_effect = side_effect aio_loop = asyncio.new_event_loop() loop._loop = aio_loop calls = [] def mock_call_soon(fn, *args): calls.append(args) fn(*args) loop._loop.call_soon_threadsafe = mock_call_soon t = threading.Thread(target=loop._poll_loop) t.start() t.join(timeout=5) assert len(calls) >= 1 assert calls[0][0].id == "j-1" class TestPollFailureLogging: @staticmethod def _make_loop(mock_api, shutdown_event): lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, heartbeat_interval_seconds=100, poll_idle_interval_seconds=0.01, initial_backoff_seconds=0.001, backoff_cap_seconds=0.001, ) lp._loop = asyncio.new_event_loop() return lp def test_poll__sustained_429__warns_with_firewall_and_interval_hint( self, mock_api, shutdown_event, capfd ): lp = self._make_loop(mock_api, shutdown_event) call_count = 0 def side_effect(runner_id): nonlocal call_count call_count += 1 if call_count >= in_process_loop._POLL_FAILURE_HINT_THRESHOLD: shutdown_event.set() raise ApiError(status_code=429, body="too many requests") mock_api.runners.next_job.side_effect = side_effect t = threading.Thread(target=lp._poll_loop) t.start() t.join(timeout=5) stderr = capfd.readouterr().err assert "OPIK_RUNNER_POLL_INTERVAL" in stderr assert "firewall or proxy" in stderr assert "429" in stderr # Full error detail is surfaced, not just a friendly message. assert "status_code: 429" in stderr assert "too many requests" in stderr def test_poll__sustained_connection_error__hint_without_rate_limit_note( self, mock_api, shutdown_event, capfd ): lp = self._make_loop(mock_api, shutdown_event) call_count = 0 def side_effect(runner_id): nonlocal call_count call_count += 1 if call_count >= in_process_loop._POLL_FAILURE_HINT_THRESHOLD: shutdown_event.set() raise ConnectionError("name resolution failed") mock_api.runners.next_job.side_effect = side_effect t = threading.Thread(target=lp._poll_loop) t.start() t.join(timeout=5) stderr = capfd.readouterr().err assert "firewall or proxy" in stderr assert "429" not in stderr # Full error detail (message + type via traceback) is surfaced. assert "name resolution failed" in stderr assert "ConnectionError" in stderr def test_poll__recovers_after_failures__logs_reconnected( self, mock_api, shutdown_event, capfd ): lp = self._make_loop(mock_api, shutdown_event) call_count = 0 def side_effect(runner_id): nonlocal call_count call_count += 1 if call_count == 1: raise ConnectionError("transient") shutdown_event.set() return None mock_api.runners.next_job.side_effect = side_effect t = threading.Thread(target=lp._poll_loop) t.start() t.join(timeout=5) stderr = capfd.readouterr().err assert "Reconnected to Opik server" in stderr def test_poll__success_resets_counter__intermittent_failures_stay_quiet( self, mock_api, shutdown_event, capfd ): lp = self._make_loop(mock_api, shutdown_event) # Two failures, a success (resets), two more failures: never 3 in a row. outcomes = ["fail", "fail", "ok", "fail", "fail", "stop"] call_count = 0 def side_effect(runner_id): nonlocal call_count outcome = outcomes[min(call_count, len(outcomes) - 1)] call_count += 1 if outcome == "stop": shutdown_event.set() return None if outcome == "ok": return None raise ConnectionError("blip") mock_api.runners.next_job.side_effect = side_effect t = threading.Thread(target=lp._poll_loop) t.start() t.join(timeout=5) stderr = capfd.readouterr().err assert "firewall or proxy" not in stderr assert "Reconnected to Opik server" in stderr class TestJobExecution: def test_execute_job__sync_entrypoint__calls_function( self, mock_api, shutdown_event ): def my_agent(q, **kwargs): return f"answer: {q}" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="my_agent", inputs={"q": "hello"}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "completed" def test_execute_job__async_entrypoint__calls_function( self, mock_api, shutdown_event ): async def my_agent(q, **kwargs): return f"answer: {q}" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="my_agent", inputs={"q": "hello"}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "completed" def test_execute_job__unknown_agent__reports_failed(self, mock_api, shutdown_event): lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="missing", inputs={}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "failed" assert "Unknown agent" in call_kwargs["error"] def test_execute_job__exception__reports_failed(self, mock_api, shutdown_event): def bad_agent(**kwargs): raise ValueError("boom") registry.register("bad", bad_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="bad", inputs={}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "failed" assert "ValueError" in call_kwargs["error"] def test_execute_job__timeout__reports_failed(self, mock_api, shutdown_event): def slow_agent(**kwargs): time.sleep(5) registry.register("slow", slow_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="slow", inputs={}, timeout=1) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "failed" assert "timed out" in call_kwargs["error"].lower() def test_execute_job__cancelled__skipped(self, mock_api, shutdown_event): def my_agent(**kwargs): return "ok" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) lp._cancelled_jobs["j-cancel"] = 0.0 job = LocalRunnerJob(id="j-cancel", agent_name="my_agent", inputs={}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() mock_api.runners.report_job_result.assert_not_called() def test_execute_job__trace_id__merges_opik_args(self, mock_api, shutdown_event): captured_kwargs = {} def my_agent(**kwargs): captured_kwargs.update(kwargs) return "ok" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob( id="j-1", agent_name="my_agent", inputs={ "opik_args": { "trace": {"tags": ["existing"]}, "span": {"metadata": {"k": "v"}}, } }, ) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() assert len(captured_kwargs["opik_args"]["trace"]["id"]) > 0 # injected trace ID assert captured_kwargs["opik_args"]["trace"]["tags"] == ["existing"] assert captured_kwargs["opik_args"]["span"]["metadata"] == {"k": "v"} def test_execute_job__reports_running_before_function__running_precedes_func_and_completed( self, mock_api, shutdown_event ): call_order = [] def my_agent(**kwargs): call_order.append("func") return "ok" def track_report(job_id, *, status, **kwargs): call_order.append(status) mock_api.runners.report_job_result.side_effect = track_report registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) job = LocalRunnerJob(id="j-1", agent_name="my_agent", inputs={}) aio_loop = asyncio.new_event_loop() aio_loop.run_until_complete(lp._execute_job(job)) aio_loop.close() assert call_order == ["running", "func", "completed"] def test_execute_job__running_report__uses_same_generated_trace_id_as_completed( self, mock_api, shutdown_event ): def my_agent(**kwargs): return "ok" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) job = LocalRunnerJob(id="j-1", agent_name="my_agent", inputs={}) aio_loop = asyncio.new_event_loop() aio_loop.run_until_complete(lp._execute_job(job)) aio_loop.close() calls = mock_api.runners.report_job_result.call_args_list assert len(calls) == 2 running_kwargs = calls[0][1] completed_kwargs = calls[1][1] assert running_kwargs["status"] == "running" assert completed_kwargs["status"] == "completed" # Both calls share the same generated trace_id assert running_kwargs["trace_id"] == completed_kwargs["trace_id"] assert len(running_kwargs["trace_id"]) > 0 def test_execute_job__typed_params__string_inputs_cast_before_call( self, mock_api, shutdown_event ): captured = {} def my_agent(query: str, count: int, score: float, active: bool, **kwargs): captured["query"] = query captured["count"] = count captured["score"] = score captured["active"] = active params = registry.extract_params(my_agent) registry.register("typed_agent", my_agent, "proj", params, "") lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) job = LocalRunnerJob( id="j-1", agent_name="typed_agent", inputs={"query": "hello", "count": "5", "score": "3.14", "active": "true"}, ) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() assert captured == {"query": "hello", "count": 5, "score": 3.14, "active": True} assert type(captured["count"]) is int assert type(captured["score"]) is float assert type(captured["active"]) is bool def test_execute_job__invalid_input_type__reports_failed( self, mock_api, shutdown_event ): def my_agent(count: int, **kwargs): pass params = registry.extract_params(my_agent) registry.register("typed_fail", my_agent, "proj", params, "") lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) job = LocalRunnerJob( id="j-1", agent_name="typed_fail", inputs={"count": "3.9"}, ) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "failed" assert "TypeError" in call_kwargs["error"] def test_execute_job__report_failure__does_not_raise( self, mock_api, shutdown_event ): def my_agent(**kwargs): return "ok" registry.register("my_agent", my_agent, "proj", [], "") mock_api.runners.report_job_result.side_effect = RuntimeError("network down") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="my_agent", inputs={}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() # All report calls go through _safe_report_job_result — failures are swallowed. assert mock_api.runners.report_job_result.call_count == 2 def test_execute_job__prompt_masks__activates_mask_context_during_execution( self, mock_api, shutdown_event ): captured = {} def my_agent(**kwargs): captured["p1"] = prompt_mask_context_module.get_mask_for_prompt("prompt-1") captured["p2"] = prompt_mask_context_module.get_mask_for_prompt("prompt-2") captured["unknown"] = prompt_mask_context_module.get_mask_for_prompt( "prompt-unknown" ) return "ok" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob( id="j-1", agent_name="my_agent", inputs={}, prompt_masks={"prompt-1": "mask-a", "prompt-2": "mask-b"}, ) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() assert captured["p1"] == "mask-a" assert captured["p2"] == "mask-b" assert captured["unknown"] is None assert prompt_mask_context_module.get_active_prompt_masks() is None def test_execute_job__prompt_masks_absent__mask_context_inactive( self, mock_api, shutdown_event ): captured = {} def my_agent(**kwargs): captured["p1"] = prompt_mask_context_module.get_mask_for_prompt("prompt-1") captured["all"] = prompt_mask_context_module.get_active_prompt_masks() return "ok" registry.register("my_agent", my_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop( mock_api, "r-1", shutdown_event, ) job = LocalRunnerJob(id="j-1", agent_name="my_agent", inputs={}) loop = asyncio.new_event_loop() loop.run_until_complete(lp._execute_job(job)) loop.close() assert captured["p1"] is None assert captured["all"] is None class TestInjectTraceId: def test_absent_key__injects_fresh_opik_args(self): inputs: dict = {} in_process_loop._inject_trace_id(inputs, "tid-1") assert inputs["opik_args"]["trace"]["id"] == "tid-1" def test_dict_value__merges_trace_id(self): inputs = {"opik_args": {"trace": {"tags": ["t"]}, "span": {"k": "v"}}} in_process_loop._inject_trace_id(inputs, "tid-2") assert inputs["opik_args"]["trace"]["id"] == "tid-2" assert inputs["opik_args"]["trace"]["tags"] == ["t"] assert inputs["opik_args"]["span"] == {"k": "v"} def test_explicit_none__leaves_inputs_unchanged(self): inputs: dict = {"opik_args": None} in_process_loop._inject_trace_id(inputs, "tid-3") assert inputs["opik_args"] is None def test_trace_none__treats_as_empty(self): inputs = {"opik_args": {"trace": None}} in_process_loop._inject_trace_id(inputs, "tid-4") assert inputs["opik_args"]["trace"]["id"] == "tid-4" def test_non_dict_opik_args__replaces_with_fresh(self): inputs: dict = {"opik_args": "unexpected"} in_process_loop._inject_trace_id(inputs, "tid-5") assert inputs["opik_args"]["trace"]["id"] == "tid-5" def test_does_not_mutate_original_dict(self): original_trace = {"tags": ["x"]} original_opik = {"trace": original_trace} inputs = {"opik_args": original_opik} in_process_loop._inject_trace_id(inputs, "tid-6") assert original_opik.get("trace", {}).get("id") is None assert original_trace.get("id") is None class TestJobLogs: @staticmethod def _run_with_streamer(lp, coro): loop = asyncio.new_event_loop() lp._loop = loop from opik.runner.log_streamer import LogStreamer lp._log_streamer = LogStreamer(lp._api, loop) async def _wrapper(): lp._log_streamer.start() await coro await asyncio.sleep(0.1) await lp._log_streamer.stop() loop.run_until_complete(_wrapper()) loop.close() @staticmethod def _enqueue(lp, job_id, text): from opik.rest_api.types.local_runner_log_entry import LocalRunnerLogEntry lp._log_streamer._queue.put_nowait( (job_id, LocalRunnerLogEntry(stream="stdout", text=text)) ) def test_execute_job__logs_sent_when_present(self, mock_api, shutdown_event): lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) async def my_agent(**kwargs): self._enqueue(lp, "j-log-1", "hello\n") return "ok" registry.register("my_agent", my_agent, "proj", [], "") job = LocalRunnerJob(id="j-log-1", agent_name="my_agent", inputs={}) self._run_with_streamer(lp, lp._execute_job(job)) mock_api.runners.append_job_logs.assert_called_once() call_kwargs = mock_api.runners.append_job_logs.call_args[1] assert call_kwargs["job_id"] == "j-log-1" assert any("hello" in e.text for e in call_kwargs["request"]) def test_execute_job__no_logs__append_not_called(self, mock_api, shutdown_event): async def silent_agent(**kwargs): return "ok" registry.register("silent_agent", silent_agent, "proj", [], "") lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) job = LocalRunnerJob(id="j-log-2", agent_name="silent_agent", inputs={}) self._run_with_streamer(lp, lp._execute_job(job)) mock_api.runners.append_job_logs.assert_not_called() def test_execute_job__exception__logs_sent_then_result_failed( self, mock_api, shutdown_event ): lp = in_process_loop.InProcessRunnerLoop(mock_api, "r-1", shutdown_event) async def failing_agent(**kwargs): self._enqueue(lp, "j-log-3", "about to fail\n") raise ValueError("boom") registry.register("failing_agent", failing_agent, "proj", [], "") job = LocalRunnerJob(id="j-log-3", agent_name="failing_agent", inputs={}) self._run_with_streamer(lp, lp._execute_job(job)) mock_api.runners.append_job_logs.assert_called_once() call_kwargs = mock_api.runners.report_job_result.call_args[1] assert call_kwargs["status"] == "failed"