From f7b537dc97f9a0cee5bfd27b38ba3f1034fd8faa Mon Sep 17 00:00:00 2001 From: Nicholas Vitsentzatos Date: Wed, 8 Jul 2026 22:30:15 -0400 Subject: [PATCH 1/3] feat(orchestration): Added TaskRunner loop + tests (initial commit of chunk B2) --- mobilerun/orchestration/runner.py | 60 +++++++ tests/test_task_runner.py | 257 ++++++++++++++++++++++++++++++ 2 files changed, 317 insertions(+) create mode 100644 mobilerun/orchestration/runner.py create mode 100644 tests/test_task_runner.py diff --git a/mobilerun/orchestration/runner.py b/mobilerun/orchestration/runner.py new file mode 100644 index 00000000..8d294c73 --- /dev/null +++ b/mobilerun/orchestration/runner.py @@ -0,0 +1,60 @@ +import logging +from typing import Any, Callable + +from mobilerun.orchestration.invoker import AgentInvoker +from mobilerun.orchestration.models import TaskRequest, TaskResult +from mobilerun.orchestration.queue import TaskQueue + +logger = logging.getLogger("mobilerun.orchestration") + +EventCallback = Callable[[Any], None] + + +class TaskRunner: + """Drives TaskQueue -> AgentInvoker sequentially, one task at a time.""" + + def __init__( + self, + queue: TaskQueue, + invoker: AgentInvoker, + event_callback: EventCallback | None = None, + ) -> None: + self._queue = queue + self._invoker = invoker + self._event_callback = event_callback + self._stop_requested = False + + async def run_once(self) -> TaskResult: + request: TaskRequest = await self._queue.dequeue() + self._queue.mark_running(request.id) + try: + result = await self._invoker.run_task(request, event_callback=self._event_callback) + except Exception as exc: + # Invoker exceptions must never kill the loop. Note this except only wraps _invoker.run_task + logger.exception("invoker.run_task raised for task %s", request.id) + result = TaskResult( + task_id=request.id, + success=False, + reason=f"invoker raised {type(exc).__name__}: {exc}", + steps=0, + error=str(exc) or type(exc).__name__, + ) + + + if result.success: + self._queue.mark_completed(request.id, result) + else: + self._queue.mark_failed(request.id, result) + return result + + async def run_forever(self) -> None: + """Runs run_once in a loop until stop() is called between tasks. stop() only takes effect between tasks + (graceful drain of whatever is currently running); it will not interrupt an idle dequeue() that is blocked + waiting for the next ready task. + """ + self._stop_requested = False + while not self._stop_requested: + await self.run_once() + + def stop(self) -> None: + self._stop_requested = True diff --git a/tests/test_task_runner.py b/tests/test_task_runner.py new file mode 100644 index 00000000..1cc82bb5 --- /dev/null +++ b/tests/test_task_runner.py @@ -0,0 +1,257 @@ +import asyncio + +from mobilerun.orchestration.models import TaskRequest, TaskResult +from mobilerun.orchestration.runner import TaskRunner + + +class FakeQueue: + """In-memory queue double. Logs mark_* calls; dequeue drains a preloaded list.""" + + def __init__(self, requests): + self._requests = list(requests) + self.calls = [] # ordered log of (op, task_id) tuples + self.results = {} # task_id -> TaskResult passed to mark_completed/mark_failed + + async def dequeue(self): + if not self._requests: + # Mirror the real queue idling on an empty/not-ready queue: block + # forever. Callers stop the loop via TaskRunner.stop() or by + # cancelling the asyncio.Task wrapping run_forever(). + await asyncio.Event().wait() + request = self._requests.pop(0) + self.calls.append(("dequeue", request.id)) + return request + + def mark_running(self, task_id): + self.calls.append(("mark_running", task_id)) + + def mark_completed(self, task_id, result): + self.calls.append(("mark_completed", task_id)) + self.results[task_id] = result + + def mark_failed(self, task_id, result): + self.calls.append(("mark_failed", task_id)) + self.results[task_id] = result + + +class FakeInvoker: + """Returns preloaded TaskResults in order, or raises a preloaded exception. + + results: list of TaskResult returned per successive run_task call. + raise_on: dict {call_index -> Exception} -- raise instead of returning. + on_call: optional callable(index) invoked at the start of each run_task, + used by tests to trigger stop() mid-task. + """ + + def __init__(self, results=None, raise_on=None, on_call=None): + self._results = list(results or []) + self._raise_on = raise_on or {} + self._on_call = on_call + self.received_requests = [] + self.received_event_callbacks = [] + + async def run_task(self, request, event_callback=None): + index = len(self.received_requests) + self.received_requests.append(request) + self.received_event_callbacks.append(event_callback) + if self._on_call is not None: + self._on_call(index) + if index in self._raise_on: + raise self._raise_on[index] + return self._results[index] + + +def _request(goal="do a thing"): + return TaskRequest(goal=goal) + + +def _result(task_id, success=True, reason="ok", steps=1): + return TaskResult(task_id=task_id, success=success, reason=reason, steps=steps) + + +def test_run_once_happy_path_transitions_and_call_order(): + request = _request() + invoker_result = _result(request.id, success=True) + queue = FakeQueue([request]) + invoker = FakeInvoker(results=[invoker_result]) + runner = TaskRunner(queue, invoker) + + returned = asyncio.run(runner.run_once()) + + assert queue.calls == [ + ("dequeue", request.id), + ("mark_running", request.id), + ("mark_completed", request.id), + ] + assert returned is invoker_result + assert invoker.received_requests == [request] + + +def test_run_once_forwards_event_callback(): + request = _request() + queue = FakeQueue([request]) + invoker = FakeInvoker(results=[_result(request.id)]) + sentinel = object() + runner = TaskRunner(queue, invoker, event_callback=sentinel) + + asyncio.run(runner.run_once()) + + assert invoker.received_event_callbacks == [sentinel] + + +def test_run_once_invoker_returns_failure_marks_failed(): + request = _request() + failed = _result(request.id, success=False, reason="agent gave up") + queue = FakeQueue([request]) + invoker = FakeInvoker(results=[failed]) + runner = TaskRunner(queue, invoker) + + returned = asyncio.run(runner.run_once()) + + assert ("mark_failed", request.id) in queue.calls + assert ("mark_completed", request.id) not in queue.calls + assert queue.results[request.id] is failed + assert returned is failed + + +def test_run_once_invoker_raises_converts_to_synthetic_failed_result(): + request = _request() + queue = FakeQueue([request]) + invoker = FakeInvoker(raise_on={0: RuntimeError("boom")}) + runner = TaskRunner(queue, invoker) + + returned = asyncio.run(runner.run_once()) + + # Did not raise; converted to a synthetic failed result. + assert returned.success is False + assert returned.task_id == request.id + assert "RuntimeError" in returned.reason + assert "boom" in returned.reason + assert returned.error is not None and "boom" in returned.error + # mark_running happened before the failure; mark_failed, never mark_completed. + assert queue.calls == [ + ("dequeue", request.id), + ("mark_running", request.id), + ("mark_failed", request.id), + ] + + +def test_run_forever_processes_multiple_tasks_sequentially(): + requests = [_request(f"goal {i}") for i in range(3)] + queue = FakeQueue(requests) + invoker = FakeInvoker(results=[_result(r.id) for r in requests]) + runner = TaskRunner(queue, invoker) + + async def scenario(): + task = asyncio.create_task(runner.run_forever()) + # Wait until all three completions land, then stop and cancel. + while sum(1 for c in queue.calls if c[0] == "mark_completed") < 3: + await asyncio.sleep(0) + runner.stop() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(scenario()) + + # Strict alternation: dequeue -> mark_running -> mark_completed, per task, + # never two mark_running without an intervening completion/failure. + expected = [] + for r in requests: + expected += [ + ("dequeue", r.id), + ("mark_running", r.id), + ("mark_completed", r.id), + ] + assert queue.calls == expected + + +def test_run_forever_continues_after_failing_invoker(): + req_fail, req_ok = _request("fails"), _request("ok") + queue = FakeQueue([req_fail, req_ok]) + invoker = FakeInvoker( + # index 0 raises (see raise_on), so its result slot is never read. + results=[None, _result(req_ok.id)], + raise_on={0: RuntimeError("first task explodes")}, + ) + runner = TaskRunner(queue, invoker) + + async def scenario(): + task = asyncio.create_task(runner.run_forever()) + while ("mark_completed", req_ok.id) not in queue.calls: + await asyncio.sleep(0) + runner.stop() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + + asyncio.run(scenario()) + + assert ("mark_failed", req_fail.id) in queue.calls + assert ("mark_completed", req_ok.id) in queue.calls + + +def test_stop_after_current_task_exits_loop_before_next_dequeue(): + req1, req2 = _request("first"), _request("second") + queue = FakeQueue([req1, req2]) + + def _stop_during_first(index): + if index == 0: + runner.stop() + + invoker = FakeInvoker( + results=[_result(req1.id), _result(req2.id)], + on_call=_stop_during_first, + ) + runner = TaskRunner(queue, invoker) + + asyncio.run(asyncio.wait_for(runner.run_forever(), timeout=1.0)) + + # First task fully processed; second never dequeued. + assert ("mark_completed", req1.id) in queue.calls + assert ("dequeue", req2.id) not in queue.calls + assert ("mark_running", req2.id) not in queue.calls + + +def test_stop_called_while_idle_does_not_hang_forever_run(): + # Empty queue -> dequeue() blocks forever. stop() alone does NOT unblock it; + # this documents the deliberate tradeoff (cancel the task for immediate stop). + queue = FakeQueue([]) + invoker = FakeInvoker() + runner = TaskRunner(queue, invoker) + + async def scenario(): + task = asyncio.create_task(runner.run_forever()) + runner.stop() # flag set, but the loop is parked inside dequeue() + try: + await asyncio.wait_for(asyncio.shield(task), timeout=0.05) + raise AssertionError("run_forever should still be blocked in dequeue()") + except asyncio.TimeoutError: + pass + assert not task.done() + # Cancellation is the intended escape hatch for an idle loop. + task.cancel() + await asyncio.gather(task, return_exceptions=True) + assert task.cancelled() + + asyncio.run(scenario()) + + +def test_run_forever_cancellation_propagates_cleanly(): + # A blocked run_forever() cancels cleanly -- CancelledError is not swallowed + # by run_once's `except Exception` (it subclasses BaseException). + queue = FakeQueue([]) + invoker = FakeInvoker() + runner = TaskRunner(queue, invoker) + + async def scenario(): + task = asyncio.create_task(runner.run_forever()) + await asyncio.sleep(0) # let it reach the blocked dequeue() + task.cancel() + with_error = None + try: + await task + except asyncio.CancelledError as exc: + with_error = exc + assert with_error is not None + assert task.cancelled() + + asyncio.run(scenario()) From 10d05da05f780c34ffe4b3a1f3791e6f5df5df3b Mon Sep 17 00:00:00 2001 From: Anh Tuan Dang Date: Wed, 8 Jul 2026 22:57:47 -0400 Subject: [PATCH 2/3] feat(orchestration): add MobileAgentInvoker and shared models (B1) Co-authored-by: Cursor --- mobilerun/orchestration/__init__.py | 20 +++ mobilerun/orchestration/invoker.py | 76 +++++++++++ mobilerun/orchestration/models.py | 56 ++++++++ tests/test_agent_invoker.py | 205 ++++++++++++++++++++++++++++ 4 files changed, 357 insertions(+) create mode 100644 mobilerun/orchestration/__init__.py create mode 100644 mobilerun/orchestration/invoker.py create mode 100644 mobilerun/orchestration/models.py create mode 100644 tests/test_agent_invoker.py diff --git a/mobilerun/orchestration/__init__.py b/mobilerun/orchestration/__init__.py new file mode 100644 index 00000000..f2d0959e --- /dev/null +++ b/mobilerun/orchestration/__init__.py @@ -0,0 +1,20 @@ +"""Multi-task orchestration layer above the MobileAgent runtime.""" + +from mobilerun.orchestration.invoker import AgentInvoker, MobileAgentInvoker +from mobilerun.orchestration.models import ( + InvalidTransitionError, + TaskRecord, + TaskRequest, + TaskResult, + TaskStatus, +) + +__all__ = [ + "AgentInvoker", + "InvalidTransitionError", + "MobileAgentInvoker", + "TaskRecord", + "TaskRequest", + "TaskResult", + "TaskStatus", +] diff --git a/mobilerun/orchestration/invoker.py b/mobilerun/orchestration/invoker.py new file mode 100644 index 00000000..6519fbcc --- /dev/null +++ b/mobilerun/orchestration/invoker.py @@ -0,0 +1,76 @@ +"""Agent invocation seam — the only orchestration module that touches MobileAgent.""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any, Callable, Protocol + +from mobilerun import MobileAgent, ResultEvent +from mobilerun.agent.utils.llm_picker import load_llm +from mobilerun.config_manager.config_manager import MobileConfig +from mobilerun.config_manager.loader import ConfigLoader + +from mobilerun.orchestration.models import TaskRequest, TaskResult + + +class AgentInvoker(Protocol): + async def run_task( + self, + request: TaskRequest, + event_callback: Callable[[Any], None] | None = None, + ) -> TaskResult: ... + + +class MobileAgentInvoker: + """Default AgentInvoker wrapping the existing MobileAgent runtime.""" + + def __init__(self, config: MobileConfig | None = None) -> None: + self._config = config + + async def run_task( + self, + request: TaskRequest, + event_callback: Callable[[Any], None] | None = None, + ) -> TaskResult: + started_at = datetime.now() + try: + config = self._config or ConfigLoader.load(request.config_path) + if request.device_serial: + config.device.serial = request.device_serial + + llms = None + if request.llm_provider: + llms = load_llm(request.llm_provider, model=request.llm_model) + + agent = MobileAgent( + goal=request.goal, + config=config, + llms=llms, + timeout=request.timeout, + ) + handler = agent.run() + + async for event in handler.stream_events(): + if event_callback: + event_callback(event) + + result: ResultEvent = await handler + return TaskResult( + task_id=request.id, + success=result.success, + reason=result.reason, + steps=result.steps, + structured_output=result.structured_output, + started_at=started_at, + finished_at=datetime.now(), + ) + except Exception as exc: + return TaskResult( + task_id=request.id, + success=False, + reason="agent raised", + steps=0, + error=str(exc), + started_at=started_at, + finished_at=datetime.now(), + ) diff --git a/mobilerun/orchestration/models.py b/mobilerun/orchestration/models.py new file mode 100644 index 00000000..c9cf3892 --- /dev/null +++ b/mobilerun/orchestration/models.py @@ -0,0 +1,56 @@ +"""Shared orchestration data model (contract §0.1).""" + +from __future__ import annotations + +import uuid +from dataclasses import dataclass, field +from datetime import datetime, timedelta +from enum import Enum +from typing import Any + + +class TaskStatus(str, Enum): + WAITING = "waiting" + RUNNING = "running" + COMPLETED = "completed" + FAILED = "failed" + CANCELLED = "cancelled" + + +class InvalidTransitionError(Exception): + """Raised when a task status transition is not allowed.""" + + +@dataclass +class TaskRequest: + goal: str + id: str = field(default_factory=lambda: uuid.uuid4().hex) + created_at: datetime = field(default_factory=datetime.now) + scheduled_at: datetime | None = None + repeat_every: timedelta | None = None + priority: int = 0 + device_serial: str | None = None + llm_provider: str | None = None + llm_model: str | None = None + config_path: str | None = None + timeout: int = 1000 + metadata: dict = field(default_factory=dict) + + +@dataclass +class TaskResult: + task_id: str + success: bool + reason: str + steps: int + structured_output: Any | None = None + error: str | None = None + started_at: datetime | None = None + finished_at: datetime | None = None + + +@dataclass +class TaskRecord: + request: TaskRequest + status: TaskStatus = TaskStatus.WAITING + result: TaskResult | None = None diff --git a/tests/test_agent_invoker.py b/tests/test_agent_invoker.py new file mode 100644 index 00000000..b90a57b0 --- /dev/null +++ b/tests/test_agent_invoker.py @@ -0,0 +1,205 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import patch + +from mobilerun.config_manager.config_manager import MobileConfig +from mobilerun.orchestration.invoker import MobileAgentInvoker +from mobilerun.orchestration.models import TaskRequest + + +class FakeHandler: + def __init__(self, result=None, *, raise_on_await: Exception | None = None): + self._result = result or SimpleNamespace( + success=True, + reason="done", + steps=3, + structured_output=None, + ) + self._raise_on_await = raise_on_await + + async def stream_events(self): + yield SimpleNamespace(type="tap") + yield SimpleNamespace(type="swipe") + + def __await__(self): + async def done(): + if self._raise_on_await is not None: + raise self._raise_on_await + return self._result + + return done().__await__() + + +class FakeAgent: + instances: list[dict] = [] + + def __init__(self, **kwargs): + FakeAgent.instances.append(kwargs) + + def run(self): + return FakeHandler() + + +def test_run_task_maps_result_event_to_task_result(): + FakeAgent.instances = [] + request = TaskRequest(goal="open settings", id="task-1") + config = MobileConfig() + + with ( + patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + return_value=config, + ), + patch("mobilerun.orchestration.invoker.MobileAgent", FakeAgent), + ): + result = asyncio.run(MobileAgentInvoker().run_task(request)) + + assert result.task_id == "task-1" + assert result.success is True + assert result.reason == "done" + assert result.steps == 3 + assert result.error is None + assert result.started_at is not None + assert result.finished_at is not None + + agent_kwargs = FakeAgent.instances[0] + assert agent_kwargs["goal"] == "open settings" + assert agent_kwargs["config"] is config + assert agent_kwargs["timeout"] == 1000 + assert agent_kwargs["llms"] is None + + +def test_run_task_forwards_streamed_events_to_callback(): + FakeAgent.instances = [] + request = TaskRequest(goal="tap home") + received: list[object] = [] + + with ( + patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + return_value=MobileConfig(), + ), + patch("mobilerun.orchestration.invoker.MobileAgent", FakeAgent), + ): + asyncio.run( + MobileAgentInvoker().run_task(request, event_callback=received.append) + ) + + assert len(received) == 2 + assert received[0].type == "tap" + assert received[1].type == "swipe" + + +def test_run_task_without_callback_does_not_crash(): + FakeAgent.instances = [] + request = TaskRequest(goal="tap home") + + with ( + patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + return_value=MobileConfig(), + ), + patch("mobilerun.orchestration.invoker.MobileAgent", FakeAgent), + ): + result = asyncio.run(MobileAgentInvoker().run_task(request, event_callback=None)) + + assert result.success is True + + +def test_run_task_applies_device_serial_override(): + FakeAgent.instances = [] + request = TaskRequest(goal="check wifi", device_serial="emulator-5554") + config = MobileConfig() + + with ( + patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + return_value=config, + ), + patch("mobilerun.orchestration.invoker.MobileAgent", FakeAgent), + ): + asyncio.run(MobileAgentInvoker().run_task(request)) + + assert config.device.serial == "emulator-5554" + + +def test_run_task_uses_injected_config_without_loading(): + FakeAgent.instances = [] + request = TaskRequest(goal="check wifi", config_path="/ignored/config.yaml") + config = MobileConfig() + config.device.serial = "preset-device" + + with ( + patch("mobilerun.orchestration.invoker.ConfigLoader.load") as load_config, + patch("mobilerun.orchestration.invoker.MobileAgent", FakeAgent), + ): + asyncio.run(MobileAgentInvoker(config=config).run_task(request)) + + load_config.assert_not_called() + assert FakeAgent.instances[0]["config"] is config + + +def test_run_task_loads_llm_when_provider_set(): + FakeAgent.instances = [] + request = TaskRequest( + goal="summarize screen", + llm_provider="OpenAI", + llm_model="gpt-5.1", + ) + fake_llm = object() + + with ( + patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + return_value=MobileConfig(), + ), + patch("mobilerun.orchestration.invoker.load_llm", return_value=fake_llm) as load_llm, + patch("mobilerun.orchestration.invoker.MobileAgent", FakeAgent), + ): + asyncio.run(MobileAgentInvoker().run_task(request)) + + load_llm.assert_called_once_with("OpenAI", model="gpt-5.1") + assert FakeAgent.instances[0]["llms"] is fake_llm + + +def test_run_task_returns_failed_result_when_agent_raises(): + FakeAgent.instances = [] + + class RaisingFakeAgent: + def __init__(self, **kwargs): + FakeAgent.instances.append(kwargs) + + def run(self): + return FakeHandler(raise_on_await=RuntimeError("boom")) + + request = TaskRequest(goal="fail task", id="task-fail") + + with ( + patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + return_value=MobileConfig(), + ), + patch("mobilerun.orchestration.invoker.MobileAgent", RaisingFakeAgent), + ): + result = asyncio.run(MobileAgentInvoker().run_task(request)) + + assert result.task_id == "task-fail" + assert result.success is False + assert result.reason == "agent raised" + assert result.steps == 0 + assert result.error == "boom" + assert result.started_at is not None + assert result.finished_at is not None + + +def test_run_task_returns_failed_result_when_config_load_raises(): + request = TaskRequest(goal="bad config", id="task-config") + + with patch( + "mobilerun.orchestration.invoker.ConfigLoader.load", + side_effect=FileNotFoundError("missing config"), + ): + result = asyncio.run(MobileAgentInvoker().run_task(request)) + + assert result.success is False + assert result.error == "missing config" From b494d1100e648f9fd66cfaa1c01b6878b6228401 Mon Sep 17 00:00:00 2001 From: We1Yuan Date: Thu, 9 Jul 2026 10:41:29 -0400 Subject: [PATCH 3/3] Add orchestration event bridge --- manual_test_events.py | 56 +++++ mobilerun/orchestration/events.py | 330 ++++++++++++++++++++++++++++++ 2 files changed, 386 insertions(+) create mode 100644 manual_test_events.py create mode 100644 mobilerun/orchestration/events.py diff --git a/manual_test_events.py b/manual_test_events.py new file mode 100644 index 00000000..71a0cc4c --- /dev/null +++ b/manual_test_events.py @@ -0,0 +1,56 @@ +import logging +from dataclasses import dataclass + +from mobilerun.orchestration.events import ( + make_agent_event_callback, + raw_agent_event_to_orchestration_event, +) + + +logging.basicConfig( + level=logging.DEBUG, + format="%(levelname)s:%(name)s:%(message)s", +) + + +@dataclass +class FakeTapEvent: + action: str = "tap" + x: int = 120 + y: int = 300 + description: str = "Tapped login button" + + +received_events = [] + + +def external_callback(event): + print("CALLBACK RECEIVED:", event) + received_events.append(event) + + +raw_event = FakeTapEvent() + +normalized = raw_agent_event_to_orchestration_event( + raw_event, + task_id="task-123", +) + +print("NORMALIZED EVENT:") +print(normalized) + +agent_callback = make_agent_event_callback( + task_id="task-123", + event_callback=external_callback, + use_cli_handler=False, +) + +agent_callback(raw_event) + +assert len(received_events) == 1 +assert received_events[0].task_id == "task-123" +assert received_events[0].payload["action"] == "tap" +assert received_events[0].payload["x"] == 120 +assert received_events[0].payload["y"] == 300 + +print("events.py manual test passed") \ No newline at end of file diff --git a/mobilerun/orchestration/events.py b/mobilerun/orchestration/events.py new file mode 100644 index 00000000..93be2b3c --- /dev/null +++ b/mobilerun/orchestration/events.py @@ -0,0 +1,330 @@ +"""Event types and logging bridge for the orchestration runner. + +This module intentionally does not import MobileAgent. The invoker is the +only orchestration module that should touch the real MobileAgent runtime. +""" + +from __future__ import annotations + +import logging +from dataclasses import asdict, dataclass, field, is_dataclass +from datetime import datetime, timezone +from enum import Enum +from typing import Any, Callable, Mapping + +from mobilerun.orchestration.models import TaskRequest, TaskResult + +logger = logging.getLogger("mobilerun.orchestration") + + +class OrchestrationEventType(str, Enum): + TASK_STARTED = "task_started" + TASK_COMPLETED = "task_completed" + TASK_FAILED = "task_failed" + + AGENT_PROGRESS = "agent_progress" + AGENT_ACTION = "agent_action" + AGENT_RESULT = "agent_result" + AGENT_ERROR = "agent_error" + AGENT_EVENT = "agent_event" + + +@dataclass(frozen=True) +class OrchestrationEvent: + """Normalized event emitted by the orchestration layer.""" + + type: OrchestrationEventType + task_id: str | None = None + message: str | None = None + raw_event_type: str | None = None + payload: dict[str, Any] = field(default_factory=dict) + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + +EventCallback = Callable[[OrchestrationEvent], None] +RawAgentEventCallback = Callable[[Any], None] + + +def task_started_event(request: TaskRequest) -> OrchestrationEvent: + return OrchestrationEvent( + type=OrchestrationEventType.TASK_STARTED, + task_id=request.id, + message=f"Task started: {request.goal}", + payload={ + "goal": request.goal, + "priority": request.priority, + "device_serial": request.device_serial, + "llm_provider": request.llm_provider, + "llm_model": request.llm_model, + "timeout": request.timeout, + "metadata": request.metadata, + }, + ) + + +def task_completed_event(result: TaskResult) -> OrchestrationEvent: + return OrchestrationEvent( + type=OrchestrationEventType.TASK_COMPLETED, + task_id=result.task_id, + message=result.reason, + payload={ + "success": result.success, + "reason": result.reason, + "steps": result.steps, + "structured_output": _safe_value(result.structured_output), + "started_at": result.started_at, + "finished_at": result.finished_at, + }, + ) + + +def task_failed_event(result: TaskResult) -> OrchestrationEvent: + return OrchestrationEvent( + type=OrchestrationEventType.TASK_FAILED, + task_id=result.task_id, + message=result.error or result.reason, + payload={ + "success": result.success, + "reason": result.reason, + "steps": result.steps, + "error": result.error, + "started_at": result.started_at, + "finished_at": result.finished_at, + }, + ) + + +def raw_agent_event_to_orchestration_event( + raw_event: Any, + task_id: str | None = None, +) -> OrchestrationEvent: + """Convert a MobileRun runtime event into a normalized orchestration event.""" + + raw_event_type = raw_event.__class__.__name__ + payload = _payload_from_raw_event(raw_event) + + event_type = _classify_raw_event(raw_event_type, payload) + message = _message_from_payload(raw_event_type, payload) + + return OrchestrationEvent( + type=event_type, + task_id=task_id, + message=message, + raw_event_type=raw_event_type, + payload=payload, + ) + + +def emit_event( + event_callback: EventCallback | None, + event: OrchestrationEvent, + *, + log: bool = True, +) -> None: + """Log and forward an orchestration event without crashing the runner.""" + + if log: + log_event(event) + + if event_callback is None: + return + + try: + event_callback(event) + except Exception: + logger.exception( + "orchestration event callback failed for task %s", + event.task_id, + ) + + +def log_event(event: OrchestrationEvent) -> None: + """Default logging bridge for orchestration-level events.""" + + extra = { + "task_id": event.task_id, + "event_type": event.type.value, + "raw_event_type": event.raw_event_type, + } + + message = event.message or event.type.value + + if event.type in { + OrchestrationEventType.TASK_FAILED, + OrchestrationEventType.AGENT_ERROR, + }: + logger.warning(message, extra=extra) + elif event.type == OrchestrationEventType.TASK_COMPLETED: + logger.info(message, extra=extra) + else: + logger.debug(message, extra=extra) + + +class LoggingEventBridge: + """Bridge raw MobileRun agent events to orchestration events and logging.""" + + def __init__( + self, + task_id: str | None = None, + event_callback: EventCallback | None = None, + *, + use_cli_handler: bool = True, + ) -> None: + self._task_id = task_id + self._event_callback = event_callback + self._cli_handler = _load_cli_event_handler() if use_cli_handler else None + + def handle_agent_event(self, raw_event: Any) -> None: + """Handle one raw event from MobileAgent.handler.stream_events().""" + + if self._cli_handler is not None: + try: + self._cli_handler.handle(raw_event) + except Exception: + logger.exception( + "CLI event handler failed for event %s", + raw_event.__class__.__name__, + ) + + event = raw_agent_event_to_orchestration_event( + raw_event, + task_id=self._task_id, + ) + emit_event(self._event_callback, event) + + +def make_agent_event_callback( + task_id: str | None = None, + event_callback: EventCallback | None = None, + *, + use_cli_handler: bool = True, +) -> RawAgentEventCallback: + """Create the callback that TaskRunner passes into AgentInvoker.run_task().""" + + bridge = LoggingEventBridge( + task_id=task_id, + event_callback=event_callback, + use_cli_handler=use_cli_handler, + ) + return bridge.handle_agent_event + + +def _load_cli_event_handler() -> Any | None: + """Load the existing CLI event handler lazily so this module stays optional.""" + + try: + from mobilerun.cli.event_handler import EventHandler + except Exception: + logger.debug("CLI EventHandler is unavailable", exc_info=True) + return None + + return EventHandler() + + +def _classify_raw_event( + raw_event_type: str, + payload: Mapping[str, Any], +) -> OrchestrationEventType: + success = payload.get("success") + error = payload.get("error") + + if error or success is False: + return OrchestrationEventType.AGENT_ERROR + + if any(token in raw_event_type for token in ("Action", "ToolCall", "Execute")): + return OrchestrationEventType.AGENT_ACTION + + if any(token in raw_event_type for token in ("Result", "Output", "End", "Finalize")): + return OrchestrationEventType.AGENT_RESULT + + if any(token in raw_event_type for token in ("Plan", "Response", "Context", "Input")): + return OrchestrationEventType.AGENT_PROGRESS + + return OrchestrationEventType.AGENT_EVENT + + +def _message_from_payload( + raw_event_type: str, + payload: Mapping[str, Any], +) -> str: + for key in ( + "summary", + "reason", + "description", + "subgoal", + "answer", + "plan", + "thought", + "error", + ): + value = payload.get(key) + if value: + return _truncate(str(value), 200) + + return raw_event_type + + +def _payload_from_raw_event(raw_event: Any) -> dict[str, Any]: + if raw_event is None: + return {} + + if isinstance(raw_event, Mapping): + return { + str(key): _safe_value(value) + for key, value in raw_event.items() + if not str(key).startswith("_") + } + + if is_dataclass(raw_event): + return { + str(key): _safe_value(value) + for key, value in asdict(raw_event).items() + if not str(key).startswith("_") + } + + if hasattr(raw_event, "model_dump"): + try: + dumped = raw_event.model_dump() + if isinstance(dumped, Mapping): + return { + str(key): _safe_value(value) + for key, value in dumped.items() + if not str(key).startswith("_") + } + except Exception: + logger.debug("model_dump failed for %s", raw_event.__class__.__name__) + + if hasattr(raw_event, "__dict__"): + return { + str(key): _safe_value(value) + for key, value in vars(raw_event).items() + if not str(key).startswith("_") + } + + return {"repr": _truncate(repr(raw_event), 500)} + + +def _safe_value(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return _truncate(value, 500) if isinstance(value, str) else value + + if isinstance(value, datetime): + return value.isoformat() + + if isinstance(value, Mapping): + return { + str(key): _safe_value(item) + for key, item in value.items() + if not str(key).startswith("_") + } + + if isinstance(value, (list, tuple, set)): + return [_safe_value(item) for item in value] + + return _truncate(repr(value), 500) + + +def _truncate(value: str, limit: int) -> str: + if len(value) <= limit: + return value + return value[: limit - 3] + "..." \ No newline at end of file