diff --git a/runtime/task.py b/runtime/task.py index 916298f..380d55f 100644 --- a/runtime/task.py +++ b/runtime/task.py @@ -24,7 +24,7 @@ from storage.timeline import Timeline from tools.describe_screen import describe_screen from tools.screenshot import take_screenshot from world.config import WorldConfig, load_config as load_world_config -from world.model import WorldModel +from world.model import TaskWorldView, WorldModel logger = logging.getLogger(__name__) @@ -87,8 +87,9 @@ class TaskRunner: def run(self, task: Task) -> Task: context = TaskContext(task_id=task.id, goal=task.goal) - if self.world_model is not None: - context.world = self.world_model.start_task(task.id) + world_handle = self._start_world_view(task.id) + if world_handle is not None: + context.world = world_handle.state self._update_task(task, status="running") for _ in range(self.config.max_steps): @@ -109,8 +110,7 @@ class TaskRunner: context=context, ) context.add_step_result(result) - self._update_world(context, scene, step, result) - self._append_timeline(task, scene, step, result) + self._record_step_result(world_handle, context, task, scene, step, result) if not result.success: self._update_task( task, @@ -128,6 +128,34 @@ class TaskRunner: ) return task + def _start_world_view(self, task_id: str) -> TaskWorldView | None: + """Start (or resume) this task's isolated `WorldModel` handle, if enabled. + + Shared by `run()` and by any other driver (e.g. `CollaborativeTaskRunner`) + that composes this `TaskRunner` for its world-model bookkeeping. + """ + if self.world_model is None: + return None + return self.world_model.start_task(task_id) + + def _record_step_result( + self, + world_handle: TaskWorldView | None, + context: TaskContext, + task: Task, + scene: Scene, + step: PlannedStep, + result: object, + ) -> None: + """Record one executed step's bookkeeping: world-model observe + timeline append. + + Shared by `run()` and by any other driver (e.g. `CollaborativeTaskRunner`) + that executes steps outside of this class's own loop, so the two paths + cannot drift apart on what gets recorded for a step. + """ + self._update_world(world_handle, context, scene, step, result) + self._append_timeline(task, scene, step, result) + def _complete_task(self, task: Task) -> Task: self._update_task(task, status="completed", completed=True) self._notify_task_succeeded(task) @@ -194,20 +222,21 @@ class TaskRunner: def _update_world( self, + world_handle: TaskWorldView | None, context: TaskContext, scene: Scene, step: PlannedStep, result: object, ) -> None: - if self.world_model is None: + if world_handle is None: return - self.world_model.observe( + world_handle.observe( scene, self._semantic_scene_from_result(result), step, result, # type: ignore[arg-type] ) - context.world = self.world_model.state + context.world = world_handle.state def _semantic_scene_from_result(self, result: object) -> SemanticScene | None: value = getattr(result, "result", None) diff --git a/tests/test_world_model.py b/tests/test_world_model.py index 96df499..2855964 100644 --- a/tests/test_world_model.py +++ b/tests/test_world_model.py @@ -129,6 +129,52 @@ def test_observe_skips_malformed_fields_without_raising() -> None: assert len(model.state.history) == 2 +def test_concurrent_tasks_on_same_world_model_stay_isolated() -> None: + model = WorldModel(config=WorldConfig(history_size=5)) + launch_a = PlannedStep( + action="launch_app", + description="launch a's app", + args={"app_id": "com.example.a"}, + ) + launch_b = PlannedStep( + action="launch_app", + description="launch b's app", + args={"app_id": "com.example.b"}, + ) + + task_a = model.start_task("task-a") + task_a.observe(_scene(), _semantic_scene("A-Home"), launch_a, _result(launch_a)) + + # Task B starts on the SAME WorldModel instance *after* task A has + # already observed an event but *before* task A observes its next one -- + # this is the interleaving that corrupted state when WorldModel tracked + # "the current task" via a single shared mutable attribute. + task_b = model.start_task("task-b") + task_b.observe(_scene(), _semantic_scene("B-Home"), launch_b, _result(launch_b)) + + remember_a = PlannedStep( + action="tap", + description="a remembers", + args={"remember": {"contact": "Alice"}}, + ) + task_a.observe(_scene(), _semantic_scene("A-Chat"), remember_a, _result(remember_a)) + + assert task_a.state.current_app == "com.example.a" + assert task_a.state.current_page == "A-Chat" + assert task_a.state.variables == {"contact": "Alice"} + assert [event.action for event in task_a.state.history] == ["launch_app", "tap"] + + assert task_b.state.current_app == "com.example.b" + assert task_b.state.current_page == "B-Home" + assert task_b.state.variables == {} + assert [event.action for event in task_b.state.history] == ["launch_app"] + + # External readers keep working via state_for(). + assert model.state_for("task-a") is task_a.state + assert model.state_for("task-b") is task_b.state + assert model.state_for("task-a") is not model.state_for("task-b") + + def test_observe_appends_semantic_or_raw_scene_history_with_eviction() -> None: model = WorldModel(config=WorldConfig(history_size=2)) first = PlannedStep(action="first", description="first") diff --git a/world/model.py b/world/model.py index 02bf32a..05777a0 100644 --- a/world/model.py +++ b/world/model.py @@ -24,13 +24,23 @@ class WorldModel: @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) -> WorldState: + def start_task(self, task_id: str) -> "TaskWorldView": self._current_task_id = task_id - return self._states.setdefault(task_id, self._new_state()) + 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) @@ -41,26 +51,51 @@ class WorldModel: 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(semantic_scene) - self._update_app(step, result) - self._update_variables(step) - self._append_history(scene, semantic_scene, step, result) + 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, semantic_scene: SemanticScene | None) -> None: + def _update_page(self, state: WorldState, semantic_scene: SemanticScene | None) -> None: if semantic_scene is None: return page = semantic_scene.page.strip() if page: - self.state.current_page = page + state.current_page = page - def _update_app(self, step: "PlannedStep", result: "StepResult") -> None: + def _update_app( + self, state: WorldState, step: "PlannedStep", result: "StepResult" + ) -> None: if not getattr(result, "success", False): return action = getattr(step, "action", None) @@ -76,11 +111,11 @@ class WorldModel: logger.info("world model skipped app update: missing app identifier") return if action == "launch_app": - self.state.current_app = app_id + state.current_app = app_id return - self.state.current_app = None + state.current_app = None - def _update_variables(self, step: "PlannedStep") -> 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 @@ -88,16 +123,17 @@ class WorldModel: if not isinstance(remember, dict): logger.info("world model skipped remember update: value is not a mapping") return - self.state.variables.update(remember) + state.variables.update(remember) def _append_history( self, + state: WorldState, scene: Scene, semantic_scene: SemanticScene | None, step: "PlannedStep", result: "StepResult", ) -> None: - self.state.history.append( + state.history.append( WorldEvent( scene_summary=semantic_scene or scene, action=str(getattr(step, "action", "unknown")), @@ -106,6 +142,40 @@ class WorldModel: ) +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