Fix reasoning manager response validation

mobilerun/agent/droid/droid_agent.py

@@ -42,6 +42,7 @@
 from mobilerun.agent.fast_agent import FastAgent
 from mobilerun.agent.fast_agent.events import FastAgentOutputEvent
 from mobilerun.agent.manager import ManagerAgent, StatelessManagerAgent
+from mobilerun.agent.manager.prompts import ManagerResponseValidationError
 from mobilerun.agent.oneflows.structured_output_agent import StructuredOutputAgent
 from mobilerun.agent.trajectory import TrajectoryWriter
 from mobilerun.agent.utils.llm_loader import (
@@ -875,6 +876,9 @@ async def run_manager(
         except DeviceDisconnectedError as e:
             logger.error(f"Device disconnected: {e}")
             return FinalizeEvent(success=False, reason=f"Device disconnected: {e}")
+        except ManagerResponseValidationError as e:
+            logger.error(f"Manager response validation failed: {e}")
+            return FinalizeEvent(success=False, reason=str(e))
 
         event = ManagerPlanEvent(
             plan=result["plan"],
@@ -900,7 +904,7 @@ async def handle_manager_plan(
                     extra={"color": "cyan"},
                 )
                 return ManagerInputEvent()
-            success = ev.success if ev.success is not None else True
+            success = ev.success if ev.success is not None else False
             self.shared_state.progress_summary = f"Answer: {ev.answer}"
             return FinalizeEvent(success=success, reason=ev.answer)
 

mobilerun/agent/manager/manager_agent.py

@@ -35,7 +35,13 @@
     ManagerPlanDetailsEvent,
     ManagerResponseEvent,
 )
-from mobilerun.agent.manager.prompts import parse_manager_response
+from mobilerun.agent.manager.prompts import (
+    ManagerResponseValidationError,
+    add_default_success_to_final_tag,
+    parse_manager_response,
+    strip_manager_final_tags,
+    validate_manager_response,
+)
 from mobilerun.agent.usage import get_usage_from_response
 from mobilerun.agent.utils.chat_utils import filter_empty_messages
 from mobilerun.agent.utils.inference import acall_with_retries
@@ -324,55 +330,54 @@ async def _validate_and_retry(
         retry_count = 0
 
         while retry_count < max_retries:
-            error_message = None
-
-            if parsed["answer"] and not parsed["plan"]:
-                if parsed["success"] is None:
-                    error_message = (
-                        'You must include success="true" or success="false" attribute '
-                        "in the <request_accomplished> tag.\n"
-                        'Example: <request_accomplished success="true">Task completed</request_accomplished>\n'
-                        "Retry again."
-                    )
-                else:
-                    break  # Valid
-            elif parsed["plan"] and parsed["answer"]:
-                error_message = (
-                    "You cannot use both request_accomplished tag while the plan is not finished. "
-                    "If you want to use request_accomplished tag, please make sure the plan is finished.\n"
-                    "Retry again."
-                )
-            elif not parsed["plan"]:
-                error_message = (
-                    "You must provide a plan to complete the task. "
-                    "Please provide a plan with the correct format."
-                )
-            else:
-                break  # Valid: plan without answer
-
-            if error_message:
-                retry_count += 1
-                logger.warning(
-                    f"Manager response invalid (retry {retry_count}/{max_retries}): {error_message}"
-                )
+            validation = validate_manager_response(parsed)
+            if validation.is_valid:
+                return output
+
+            error_message = f"{validation.error_message}\nRetry again."
+            retry_count += 1
+            logger.warning(
+                f"Manager response invalid (retry {retry_count}/{max_retries}): {error_message}"
+            )
 
-                # Build retry messages
-                retry_messages = messages + [
-                    ChatMessage(role="assistant", content=output),
-                    ChatMessage(role="user", content=error_message),
-                ]
+            retry_messages = messages + [
+                ChatMessage(role="assistant", content=output),
+                ChatMessage(role="user", content=error_message),
+            ]
 
-                try:
-                    response = await acall_with_retries(
-                        self.llm, retry_messages, stream=self.agent_config.streaming
-                    )
-                    output = response.message.content
-                    parsed = parse_manager_response(output)
-                except Exception as e:
-                    logger.error(f"LLM retry failed: {e}")
-                    break
+            try:
+                response = await acall_with_retries(
+                    self.llm, retry_messages, stream=self.agent_config.streaming
+                )
+                output = response.message.content
+                parsed = parse_manager_response(output)
+            except Exception as e:
+                logger.error(f"LLM retry failed: {e}")
+                if validation.can_continue_with_plan:
+                    return strip_manager_final_tags(output)
+                if validation.can_accept_final_without_success:
+                    return add_default_success_to_final_tag(output)
+                raise ManagerResponseValidationError(error_message) from e
+
+        validation = validate_manager_response(parsed)
+        if validation.is_valid:
+            return output
+        if validation.can_continue_with_plan:
+            logger.warning(
+                "Manager response stayed invalid after retries; continuing with the plan "
+                "and ignoring the conflicting final answer."
+            )
+            return strip_manager_final_tags(output)
+        if validation.can_accept_final_without_success:
+            logger.warning(
+                "Manager response omitted final success after retries; accepting the "
+                "terminal answer as success for backward compatibility."
+            )
+            return add_default_success_to_final_tag(output)
 
-        return output
+        raise ManagerResponseValidationError(
+            validation.error_message or "Invalid manager response."
+        )
 
     # ========================================================================
     # Workflow Steps

mobilerun/agent/manager/prompts.py

@@ -3,6 +3,73 @@
 """
 
 import re
+from dataclasses import dataclass
+
+FINAL_RESPONSE_TAGS = ("request_accomplished", "answer")
+
+_SUCCESS_ATTR_RE = re.compile(
+    r"""\bsuccess\s*=\s*(?P<quote>["'])(?P<value>true|false)(?P=quote)""",
+    re.IGNORECASE | re.DOTALL,
+)
+_SUCCESS_ATTR_PRESENT_RE = re.compile(r"\bsuccess\s*=", re.IGNORECASE | re.DOTALL)
+_FINAL_TAG_RE = re.compile(
+    r"<(?P<tag>request_accomplished|answer)\b(?P<attrs>[^>]*)>"
+    r"(?P<body>.*?)"
+    r"</(?P=tag)>",
+    re.IGNORECASE | re.DOTALL,
+)
+_FINAL_OPEN_TAG_RE = re.compile(
+    r"<(?P<tag>request_accomplished|answer)\b(?P<attrs>[^>]*)>",
+    re.IGNORECASE | re.DOTALL,
+)
+
+
+@dataclass(frozen=True)
+class ManagerResponseValidation:
+    is_valid: bool
+    error_message: str | None = None
+    can_continue_with_plan: bool = False
+    can_accept_final_without_success: bool = False
+
+
+class ManagerResponseValidationError(RuntimeError):
+    """Raised when the manager cannot produce a valid response after retries."""
+
+
+def _find_tag_matches(response: str, tag: str) -> list[re.Match[str]]:
+    pattern = re.compile(
+        rf"<{tag}\b(?P<attrs>[^>]*)>(?P<body>.*?)</{tag}>",
+        re.IGNORECASE | re.DOTALL,
+    )
+    return list(pattern.finditer(response))
+
+
+def _tag_content(match: re.Match[str]) -> str:
+    return match.group("body").strip()
+
+
+def _success_from_attrs(attrs: str) -> bool | None:
+    match = _SUCCESS_ATTR_RE.search(attrs)
+    if not match:
+        return None
+    return match.group("value").lower() == "true"
+
+
+def strip_manager_final_tags(response: str) -> str:
+    """Remove final answer tags while leaving reasoning and plan tags intact."""
+    return _FINAL_TAG_RE.sub("", response).strip()
+
+
+def add_default_success_to_final_tag(response: str) -> str:
+    """Add success=true to a final tag that omitted the success attribute."""
+
+    def replace(match: re.Match[str]) -> str:
+        attrs = match.group("attrs")
+        if _SUCCESS_ATTR_PRESENT_RE.search(attrs):
+            return match.group(0)
+        return f'<{match.group("tag")}{attrs} success="true">'
+
+    return _FINAL_OPEN_TAG_RE.sub(replace, response, count=1)
 
 
 def parse_manager_response(response: str) -> dict:
@@ -35,43 +102,48 @@ def parse_manager_response(response: str) -> dict:
 
     def extract(tag: str) -> str:
         """Extract content between XML-style tags (handles attributes)."""
-        # Use regex to handle tags with attributes like <tag attr="value">
-        pattern = rf"<{tag}(?:\s+[^>]*)?>(.+?)</{tag}>"
-        match = re.search(pattern, response, re.DOTALL)
-        if match:
-            return match.group(1).strip()
+        matches = _find_tag_matches(response, tag)
+        if matches:
+            return _tag_content(matches[0])
         return ""
 
     def extract_all(tag: str) -> str:
         """Extract and combine content from all occurrences of a tag."""
-        pattern = rf"<{tag}(?:\s+[^>]*)?>(.+?)</{tag}>"
-        matches = re.findall(pattern, response, re.DOTALL)
+        matches = [_tag_content(match) for match in _find_tag_matches(response, tag)]
         if not matches:
             return ""
-        return "\n".join(m.strip() for m in matches if m.strip())
+        return "\n".join(match for match in matches if match)
 
     thought = extract("thought")
     memory_section = extract_all("add_memory")
     plan = extract("plan")
     progress_summary = extract("progress_summary")
-    answer = extract("request_accomplished")
+    plan_matches = _find_tag_matches(response, "plan")
+    final_matches = list(_FINAL_TAG_RE.finditer(response))
+    final_matches.sort(key=lambda match: match.start())
 
-    # Also support <answer> tag as alternative to <request_accomplished>
-    if not answer:
-        answer = extract("answer")
+    final_match = None
+    for match in final_matches:
+        if _tag_content(match):
+            final_match = match
+            break
+    if final_match is None and final_matches:
+        final_match = final_matches[0]
 
-    # Parse success attribute from request_accomplished or answer tag
+    answer = _tag_content(final_match) if final_match else ""
     success = None
-    if answer:
-        # Try to extract success attribute from tag
-        success_match = re.search(
-            r'<(?:request_accomplished|answer)\s+success="(true|false)">', response
-        )
-        if success_match:
-            success = success_match.group(1) == "true"
-        else:
-            # Default to True for backward compatibility if no attribute present
-            success = True
+    final_tag = None
+    success_attr_present = False
+    if final_match:
+        final_tag = final_match.group("tag").lower()
+        final_attrs = final_match.group("attrs")
+        success_attr_present = bool(_SUCCESS_ATTR_PRESENT_RE.search(final_attrs))
+        success = _success_from_attrs(final_attrs)
+
+    final_counts = {
+        tag: sum(1 for match in final_matches if match.group("tag").lower() == tag)
+        for tag in FINAL_RESPONSE_TAGS
+    }
 
     # Parse current subgoal from first line of plan
     current_goal_text = plan
@@ -110,5 +182,79 @@ def extract_all(tag: str) -> str:
         "current_subgoal": current_subgoal,
         "answer": answer,
         "success": success,
+        "success_attr_present": success_attr_present,
         "progress_summary": progress_summary,
+        "final_tag": final_tag,
+        "tag_counts": {
+            "plan": len(plan_matches),
+            "request_accomplished": final_counts["request_accomplished"],
+            "answer": final_counts["answer"],
+            "final": len(final_matches),
+        },
     }
+
+
+def validate_manager_response(parsed: dict) -> ManagerResponseValidation:
+    """Validate the manager output contract parsed by parse_manager_response."""
+    plan = (parsed.get("plan") or "").strip()
+    answer = (parsed.get("answer") or "").strip()
+    success = parsed.get("success")
+    success_attr_present = bool(parsed.get("success_attr_present"))
+    tag_counts = parsed.get("tag_counts") or {}
+    plan_count = tag_counts.get("plan", 1 if plan else 0)
+    final_count = tag_counts.get("final", 1 if answer else 0)
+
+    if plan_count > 1:
+        return ManagerResponseValidation(
+            is_valid=False,
+            error_message="Manager response contains multiple <plan> tags. "
+            "Provide exactly one <plan>, or exactly one final answer tag.",
+        )
+
+    if final_count > 1:
+        return ManagerResponseValidation(
+            is_valid=False,
+            error_message="Manager response contains multiple final answer tags. "
+            "Provide exactly one <request_accomplished> or <answer> tag.",
+            can_continue_with_plan=plan_count == 1 and bool(plan),
+        )
+
+    if plan_count == 1 and not plan:
+        return ManagerResponseValidation(
+            is_valid=False,
+            error_message="Manager <plan> tag must not be empty.",
+        )
+
+    if final_count == 1 and not answer:
+        return ManagerResponseValidation(
+            is_valid=False,
+            error_message="Manager final answer tag must not be empty.",
+            can_continue_with_plan=plan_count == 1 and bool(plan),
+        )
+
+    if plan and answer:
+        return ManagerResponseValidation(
+            is_valid=False,
+            error_message="Manager response must provide exactly one of <plan> "
+            "or a final answer tag, not both.",
+            can_continue_with_plan=plan_count == 1 and final_count == 1,
+        )
+
+    if answer:
+        if success is None:
+            return ManagerResponseValidation(
+                is_valid=False,
+                error_message='Final answer tag must include success="true" '
+                'or success="false".',
+                can_accept_final_without_success=not success_attr_present,
+            )
+        return ManagerResponseValidation(is_valid=True)
+
+    if plan:
+        return ManagerResponseValidation(is_valid=True)
+
+    return ManagerResponseValidation(
+        is_valid=False,
+        error_message="Manager response must provide exactly one of <plan> "
+        "or a final answer tag.",
+    )