agentic-mobile-control/world/model.py

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

from core.models import Scene
from semantic.models import SemanticScene
from world.config import WorldConfig, load_config
from world.models import WorldEvent, WorldState

if TYPE_CHECKING:
    from runtime.executor import StepResult
    from runtime.planner import PlannedStep

logger = logging.getLogger(__name__)


class WorldModel:
    def __init__(self, *, config: WorldConfig | None = None) -> None:
        self.config = config or load_config()
        self._states: dict[str, WorldState] = {}
        self._current_task_id: str | None = None
        self._fallback_state = self._new_state()

    @property
    def state(self) -> WorldState:
        """Convenience accessor for single-active-task usage only.

        Relies on the shared ``_current_task_id`` pointer, so it is not
        safe when multiple tasks run concurrently against the same
        ``WorldModel`` instance (a later ``start_task()`` call from another
        task can move this pointer out from under you). Concurrent callers
        must use the ``TaskWorldView`` returned by ``start_task()`` (or
        ``state_for(task_id)``) instead.
        """
        if self._current_task_id is None:
            return self._fallback_state
        return self._states[self._current_task_id]

    def start_task(self, task_id: str) -> "TaskWorldView":
        self._current_task_id = task_id
        state = self._states.setdefault(task_id, self._new_state())
        return TaskWorldView(self, task_id, state)

    def state_for(self, task_id: str) -> WorldState | None:
        return self._states.get(task_id)

    def observe(
        self,
        scene: Scene,
        semantic_scene: SemanticScene | None,
        step: "PlannedStep",
        result: "StepResult",
    ) -> None:
        """Update the shared "current task" state. See `state`'s docstring."""
        self._apply(self.state, scene, semantic_scene, step, result)

    def observe_task(
        self,
        task_id: str,
        scene: Scene,
        semantic_scene: SemanticScene | None,
        step: "PlannedStep",
        result: "StepResult",
    ) -> None:
        """Update `task_id`'s own state explicitly, independent of any other task."""
        state = self._states.setdefault(task_id, self._new_state())
        self._apply(state, scene, semantic_scene, step, result)

    def _apply(
        self,
        state: WorldState,
        scene: Scene,
        semantic_scene: SemanticScene | None,
        step: "PlannedStep",
        result: "StepResult",
    ) -> None:
        try:
            self._update_page(state, semantic_scene)
            self._update_app(state, step, result)
            self._update_variables(state, step)
            self._append_history(state, scene, semantic_scene, step, result)
        except Exception as exc:
            logger.info("world model update skipped after unexpected error: %s", exc)

    def _new_state(self) -> WorldState:
        return WorldState.with_history_bound(self.config.history_size)

    def _update_page(self, state: WorldState, semantic_scene: SemanticScene | None) -> None:
        if semantic_scene is None:
            return
        page = semantic_scene.page.strip()
        if page:
            state.current_page = page

    def _update_app(
        self, state: WorldState, step: "PlannedStep", result: "StepResult"
    ) -> None:
        if not getattr(result, "success", False):
            return
        action = getattr(step, "action", None)
        if action not in {"launch_app", "terminate_app"}:
            return

        args = getattr(step, "args", {})
        if not isinstance(args, dict):
            logger.info("world model skipped app update: step args are not a mapping")
            return
        app_id = _app_identifier(args)
        if not app_id:
            logger.info("world model skipped app update: missing app identifier")
            return
        if action == "launch_app":
            state.current_app = app_id
            return
        state.current_app = None

    def _update_variables(self, state: WorldState, step: "PlannedStep") -> None:
        args = getattr(step, "args", {})
        if not isinstance(args, dict) or "remember" not in args:
            return
        remember = args["remember"]
        if not isinstance(remember, dict):
            logger.info("world model skipped remember update: value is not a mapping")
            return
        state.variables.update(remember)

    def _append_history(
        self,
        state: WorldState,
        scene: Scene,
        semantic_scene: SemanticScene | None,
        step: "PlannedStep",
        result: "StepResult",
    ) -> None:
        state.history.append(
            WorldEvent(
                scene_summary=semantic_scene or scene,
                action=str(getattr(step, "action", "unknown")),
                success=bool(getattr(result, "success", False)),
            )
        )


class TaskWorldView:
    """Per-task handle into a `WorldModel`, scoped explicitly by `task_id`.

    Returned by `WorldModel.start_task()` so a caller (e.g. `TaskRunner`)
    can keep observing/reading its own task's `WorldState` for the rest of
    that task's lifetime by holding onto this handle, rather than relying on
    `WorldModel`'s shared "current task" pointer. This makes it safe for
    multiple tasks to run concurrently against the same `WorldModel`
    instance without corrupting each other's state.
    """

    def __init__(self, model: WorldModel, task_id: str, state: WorldState) -> None:
        self._model = model
        self._task_id = task_id
        self._state = state

    @property
    def task_id(self) -> str:
        return self._task_id

    @property
    def state(self) -> WorldState:
        return self._state

    def observe(
        self,
        scene: Scene,
        semantic_scene: SemanticScene | None,
        step: "PlannedStep",
        result: "StepResult",
    ) -> None:
        self._model._apply(self._state, scene, semantic_scene, step, result)


def _app_identifier(args: dict[str, Any]) -> str | None:
    value = args.get("bundle_id") or args.get("app_id")
    return value if isinstance(value, str) and value.strip() else None