feat: preserve planner context across task steps

runtime/ai_planner.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from inspect import Parameter, signature
 from typing import TYPE_CHECKING, Any
 
 from core.errors import TaskFailedError
@@ -39,11 +40,11 @@ class AIPlanner(Planner):
         world: "WorldState | None" = None,
         screenshot: bytes | None = None,
     ) -> list[PlannedStep]:
+        _sync_tool_results(context)
         scene_json = _without_ocr(scene.to_dict()) if self.config.multimodal else scene.to_dict()
         user_prompt = planner_user_prompt(
             goal=goal,
             scene_json=scene_json,
-            history_summary=_history_summary(world),
             device_platform=context.device_platform,
         )
         if context.step_results and not context.step_results[-1].success:
@@ -56,13 +57,18 @@ class AIPlanner(Planner):
                    "system_prompt": PLANNER_SYSTEM_PROMPT, "user_prompt": user_prompt,
                    "has_screenshot": screenshot is not None})
         try:
-            decision = self.client.decide(
-                system_prompt=PLANNER_SYSTEM_PROMPT,
-                user_prompt=user_prompt,
-                screenshot=screenshot,
-                tools=ALL_TOOL_SPECS,
-                timeout=self.config.timeout,
-            )
+            kwargs = {
+                "system_prompt": PLANNER_SYSTEM_PROMPT,
+                "user_prompt": user_prompt,
+                "screenshot": screenshot,
+                "tools": ALL_TOOL_SPECS,
+                "timeout": self.config.timeout,
+            }
+            if _accepts_history(self.client.decide):
+                kwargs["history"] = context.planner_history[
+                    -self.config.history_max_turns :
+                ]
+            decision = self.client.decide(**kwargs)
         except Exception as exc:
             self._log({"type": "agent_error", "task_id": context.task_id, "error": str(exc)})
             raise
@@ -76,6 +82,16 @@ class AIPlanner(Planner):
                 return []
             raise TaskFailedError(decision.arguments.get("reason") or "task failed")
 
+        context.planner_history.append(
+            {
+                "user_prompt": user_prompt,
+                "tool_name": decision.tool_name,
+                "arguments": _conversation_arguments(decision),
+                "rationale": decision.text_output,
+                "tool_result": None,
+            }
+        )
+
         return [
             PlannedStep(
                 action=decision.tool_name,
@@ -107,23 +123,6 @@ class AIPlanner(Planner):
                 pass
 
 
-def _history_summary(world: "WorldState | None") -> list[dict[str, Any]]:
-    if world is None:
-        return []
-    return [
-        {
-            "page": event.page,
-            "action": event.action,
-            "arguments": dict(event.arguments),
-            "rationale": event.rationale,
-            "purpose": event.purpose,
-            "expected_outcome": event.expected_outcome,
-            "success": event.success,
-        }
-        for event in world.history
-    ]
-
-
 def _without_ocr(scene_json: dict[str, Any]) -> dict[str, Any]:
     cleaned = dict(scene_json)
     elements = cleaned.get("elements")
@@ -135,3 +134,26 @@ def _without_ocr(scene_json: dict[str, Any]) -> dict[str, Any]:
         ]
     cleaned.pop("ocr_elements", None)
     return cleaned
+
+
+def _sync_tool_results(context: TaskContext) -> None:
+    for turn, result in zip(context.planner_history, context.step_results, strict=False):
+        if turn.get("tool_result") is None:
+            turn["tool_result"] = result.to_dict()
+
+
+def _conversation_arguments(decision: Any) -> dict[str, Any]:
+    arguments = dict(decision.arguments)
+    if decision.purpose is not None:
+        arguments["purpose"] = decision.purpose
+    if decision.expected_outcome is not None:
+        arguments["expected_outcome"] = decision.expected_outcome
+    return arguments
+
+
+def _accepts_history(method: Any) -> bool:
+    parameters = signature(method).parameters.values()
+    return any(
+        parameter.name == "history" or parameter.kind == Parameter.VAR_KEYWORD
+        for parameter in parameters
+    )