From 79b819d5b6b7836bcfcb2f58895efd763c22bda7 Mon Sep 17 00:00:00 2001 From: "rasul.osmanbayli" Date: Mon, 3 Aug 2026 19:20:09 +0400 Subject: [PATCH 1/2] Stop malformed tool-call retry loops --- mobilerun/agent/fast_agent/events.py | 2 + mobilerun/agent/fast_agent/fast_agent.py | 75 ++++- mobilerun/agent/fast_agent/xml_parser.py | 58 +++- tests/test_fast_agent_malformed_tool_guard.py | 266 ++++++++++++++++++ tests/test_fast_agent_xml_parser.py | 215 ++++++++++++++ 5 files changed, 612 insertions(+), 4 deletions(-) create mode 100644 tests/test_fast_agent_malformed_tool_guard.py diff --git a/mobilerun/agent/fast_agent/events.py b/mobilerun/agent/fast_agent/events.py index c1b66283..2b287f70 100644 --- a/mobilerun/agent/fast_agent/events.py +++ b/mobilerun/agent/fast_agent/events.py @@ -8,6 +8,7 @@ from llama_index.core.workflow import Event +from mobilerun.agent.fast_agent.xml_parser import ToolCallParseStatus from mobilerun.agent.usage import UsageResult @@ -23,6 +24,7 @@ class FastAgentResponseEvent(Event): thought: str code: Optional[str] = None usage: Optional[UsageResult] = None + tool_call_status: ToolCallParseStatus = ToolCallParseStatus.NO_MARKUP class FastAgentToolCallEvent(Event): diff --git a/mobilerun/agent/fast_agent/fast_agent.py b/mobilerun/agent/fast_agent/fast_agent.py index 091816e5..cc56b40d 100644 --- a/mobilerun/agent/fast_agent/fast_agent.py +++ b/mobilerun/agent/fast_agent/fast_agent.py @@ -33,11 +33,12 @@ FastAgentToolCallEvent, ) from mobilerun.agent.fast_agent.xml_parser import ( + ToolCallParseStatus, ToolResult, extract_add_memory, format_tool_calls, format_tool_results, - parse_tool_calls, + parse_tool_calls_detailed, ) from mobilerun.agent.usage import get_usage_from_response from mobilerun.agent.utils.chat_utils import limit_history @@ -59,6 +60,24 @@ logger = logging.getLogger("mobilerun") +_MALFORMED_TOOL_CALL_LIMIT = 3 + + +def _malformed_tool_call_correction(attempt: int) -> str: + """Build a focused retry instruction without echoing malformed model output.""" + return ( + "Your previous response contained tool-call markup that could not be parsed " + f"(attempt {attempt}/{_MALFORMED_TOOL_CALL_LIMIT}). No tool was executed.\n\n" + "Repeat the intended call using only the ASCII XML tags shown below. Replace " + "the placeholder names and value with the intended tool and arguments. Do not " + "use provider-specific markers, full-width punctuation, or alternate tag names.\n\n" + "\n" + '\n' + 'value\n' + "\n" + "" + ) + class FastAgent(Workflow): """Agent that uses XML tool-calling instead of code generation. @@ -108,6 +127,7 @@ def __init__( self.system_prompt: ChatMessage | None = None self.tool_call_counter = 0 + self._consecutive_malformed_tool_calls = 0 # Build tool descriptions and param types from registry self.tool_descriptions = self.registry.get_tool_descriptions_xml() @@ -180,6 +200,7 @@ async def _build_user_prompt(self, goal: str) -> ChatMessage: async def prepare_chat(self, ctx: Context, ev: StartEvent) -> FastAgentInputEvent: """Initialize message history with goal.""" logger.debug("Preparing chat for task execution...") + self._consecutive_malformed_tool_calls = 0 # Get available secrets (only if type_secret is actually in the registry) if ( @@ -373,7 +394,9 @@ async def handle_llm_input( response_text = response.message.content # Parse tool calls from response - thought, tool_calls = parse_tool_calls(response_text, self.param_types) + parse_result = parse_tool_calls_detailed(response_text, self.param_types) + thought = parse_result.thought + tool_calls = parse_result.calls # Extract from thought text and append to unified memory memory_update = extract_add_memory(thought) @@ -396,6 +419,7 @@ async def handle_llm_input( thought=thought, code=tool_calls_xml, usage=usage, + tool_call_status=parse_result.status, ) ctx.write_event_to_stream(event) return event @@ -403,10 +427,55 @@ async def handle_llm_input( @step async def handle_llm_output( self, ctx: Context, ev: FastAgentResponseEvent - ) -> FastAgentToolCallEvent | FastAgentInputEvent: + ) -> FastAgentToolCallEvent | FastAgentInputEvent | FastAgentEndEvent: """Route to execution or request tool call if missing.""" has_tool_calls = ev.code is not None + if ev.tool_call_status == ToolCallParseStatus.MALFORMED and not has_tool_calls: + self._consecutive_malformed_tool_calls += 1 + attempt = self._consecutive_malformed_tool_calls + logger.warning( + "Malformed tool-call markup detected (%d/%d)", + attempt, + _MALFORMED_TOOL_CALL_LIMIT, + ) + + if attempt >= _MALFORMED_TOOL_CALL_LIMIT: + pending = self.shared_state.drain_user_messages() + if pending: + logger.warning( + "⚠️ Dropping %d external user message(s) at malformed tool-call limit", + len(pending), + ) + ctx.write_event_to_stream( + ExternalUserMessageDroppedEvent( + message_ids=[message.id for message in pending], + reason="malformed_tool_call_limit_reached", + step_number=self.shared_state.step_number, + ) + ) + event = FastAgentEndEvent( + success=False, + reason=( + "Model produced malformed tool-call markup " + f"{_MALFORMED_TOOL_CALL_LIMIT} consecutive times; stopped to prevent " + "a retry loop. Switch models or verify tool-call protocol compatibility." + ), + tool_call_count=self.tool_call_counter, + ) + ctx.write_event_to_stream(event) + return event + + self.shared_state.message_history.append( + ChatMessage( + role="user", + content=_malformed_tool_call_correction(attempt), + ) + ) + return FastAgentInputEvent() + + self._consecutive_malformed_tool_calls = 0 + if not ev.thought: logger.warning("LLM provided tool calls without reasoning.") no_thoughts_text = ( diff --git a/mobilerun/agent/fast_agent/xml_parser.py b/mobilerun/agent/fast_agent/xml_parser.py index faa6483a..e311b888 100644 --- a/mobilerun/agent/fast_agent/xml_parser.py +++ b/mobilerun/agent/fast_agent/xml_parser.py @@ -10,6 +10,7 @@ import re import xml.etree.ElementTree as ET from dataclasses import dataclass, field +from enum import Enum from html import escape from typing import Any, Dict, List, Optional, Tuple @@ -23,6 +24,14 @@ re.DOTALL, ) +_DSML_MARKUP_PATTERN = r"(?:<|＜)\s*/?\s*｜+DSML｜+" +_DSML_MARKUP_RE = re.compile(_DSML_MARKUP_PATTERN, re.IGNORECASE) +_TOOL_CALL_MARKUP_RE = re.compile( + r"(?:<|＜)\s*/?\s*(?:function_calls|invoke|parameter)\b" + r"|" + _DSML_MARKUP_PATTERN, + re.IGNORECASE, +) + @dataclass class ToolCall: @@ -33,6 +42,23 @@ class ToolCall: error: Optional[str] = None +class ToolCallParseStatus(str, Enum): + """Classification of tool-call markup in an LLM response.""" + + VALID = "valid" + NO_MARKUP = "no_markup" + MALFORMED = "malformed" + + +@dataclass(frozen=True) +class ToolCallParseResult: + """Detailed result used internally to distinguish missing and invalid markup.""" + + thought: str + calls: List[ToolCall] + status: ToolCallParseStatus + + @dataclass class ToolResult: """Result from executing a single tool.""" @@ -59,6 +85,28 @@ def parse_tool_calls( if OPEN_TAG not in text: return text.strip(), [] + result = parse_tool_calls_detailed(text, param_types) + return result.thought, result.calls + + +def parse_tool_calls_detailed( + text: str, param_types: Optional[Dict[str, str]] = None +) -> ToolCallParseResult: + """Parse tool calls and classify whether markup was absent, valid, or malformed.""" + if OPEN_TAG not in text: + markup_match = _TOOL_CALL_MARKUP_RE.search(text) + if markup_match: + return ToolCallParseResult( + thought=text[: markup_match.start()].strip(), + calls=[], + status=ToolCallParseStatus.MALFORMED, + ) + return ToolCallParseResult( + thought=text.strip(), + calls=[], + status=ToolCallParseStatus.NO_MARKUP, + ) + parts = text.split(OPEN_TAG) text_before = parts[0].strip() @@ -72,12 +120,20 @@ def parse_tool_calls( if not block: continue + # Parameter-value sanitization intentionally preserves raw text such as + # code, but it must never turn provider-specific tool markup into an + # executable XML call. A separate valid wrapper can still win below. + if _DSML_MARKUP_RE.search(block): + continue + calls = _parse_tool_call_block(block, param_types) if calls: call_blocks.append(calls) deduped_blocks = _drop_adjacent_duplicate_blocks(call_blocks) - return text_before, [call for block in deduped_blocks for call in block] + calls = [call for block in deduped_blocks for call in block] + status = ToolCallParseStatus.VALID if calls else ToolCallParseStatus.MALFORMED + return ToolCallParseResult(thought=text_before, calls=calls, status=status) def format_tool_results(results: List[ToolResult]) -> str: diff --git a/tests/test_fast_agent_malformed_tool_guard.py b/tests/test_fast_agent_malformed_tool_guard.py new file mode 100644 index 00000000..99311be6 --- /dev/null +++ b/tests/test_fast_agent_malformed_tool_guard.py @@ -0,0 +1,266 @@ +import asyncio +from types import SimpleNamespace +from unittest.mock import patch + +from llama_index.core.base.llms.types import ChatMessage, ChatResponse + +from mobilerun.agent.action_result import ActionResult +from mobilerun.agent.droid.state import MobileAgentState +from mobilerun.agent.fast_agent.fast_agent import FastAgent + +MALFORMED_RESPONSE = """I'm still on the main feed. I need to tap the Profile tab. + + + +<｜DSML｜ name="index">189 + +""" + +DSML_ONLY_RESPONSE = """I will tap the Profile tab. +<｜DSML｜tool_calls> +<｜DSML｜invoke name="click"> +<｜DSML｜parameter name="index">189 + +""" + +CORRUPTED_PARAMETER_RESPONSE = """Settings home is visible. Marking the task complete. + + +true +Completed 20 verified cycles. + +""" + +COMPLETE_RESPONSE = """The task is complete. + + +true +Done + +""" + +CLICK_RESPONSE = """I will tap the target. + + +12 + +""" + + +class _SequencedResponder: + def __init__(self, responses): + self.responses = list(responses) + self.requests = [] + + async def __call__(self, _llm, messages, *, stream): + self.requests.append(list(messages)) + if not self.responses: + raise AssertionError("FastAgent made an unexpected extra LLM request") + content = self.responses.pop(0) + return ChatResponse(message=ChatMessage(role="assistant", content=content)) + + +class _PromptResolver: + def get_prompt(self, name): + if name == "fast_agent_system": + return "Use the documented XML tools: {{ tool_descriptions }}" + if name == "fast_agent_user": + return "Goal: {{ goal }}" + return None + + +class _StateProvider: + requires_coordinate_tools = False + + async def get_state(self): + return SimpleNamespace( + formatted_text="Settings home with Network & internet visible", + focused_text="", + elements=[], + phone_state={ + "packageName": "com.android.settings", + "currentApp": "com.android.settings.Settings", + }, + ) + + +class _Registry: + def __init__(self, shared_state): + self.shared_state = shared_state + self.executed = [] + self.tools = {"click": object(), "complete": object()} + + def get_tool_descriptions_xml(self): + return '' + + def get_param_types(self): + return {"index": "number", "success": "boolean"} + + async def execute(self, name, parameters, _action_ctx, workflow_ctx=None): + self.executed.append((name, parameters)) + if name == "complete": + await self.shared_state.complete( + parameters.get("success", False), + message=parameters.get("message", ""), + ) + return ActionResult(success=True, summary=f"{name} executed") + + +def _run_sequence(responses, pending_message=None): + async def _run(): + shared_state = MobileAgentState() + if pending_message: + shared_state.queue_user_message(pending_message) + registry = _Registry(shared_state) + responder = _SequencedResponder(responses) + config = SimpleNamespace( + fast_agent=SimpleNamespace(vision=False, parallel_tools=False), + max_steps=20, + streaming=False, + after_sleep_action=0, + ) + action_ctx = SimpleNamespace( + driver=SimpleNamespace(), + credential_manager=None, + ui=None, + ) + agent = FastAgent( + llm=SimpleNamespace(), + agent_config=config, + registry=registry, + action_ctx=action_ctx, + state_provider=_StateProvider(), + shared_state=shared_state, + prompt_resolver=_PromptResolver(), + ) + + with ( + patch( + "mobilerun.agent.fast_agent.fast_agent.acall_with_retries", + new=responder, + ), + patch( + "mobilerun.agent.fast_agent.fast_agent.get_usage_from_response", + return_value=None, + ), + ): + result = await agent.run(input="Exercise the malformed-call guard") + + return result, responder, registry, agent, shared_state + + return asyncio.run(_run()) + + +def _correction_messages(shared_state): + return [ + message.content + for message in shared_state.message_history + if message.role == "user" + and message.content + and "tool-call markup that could not be parsed" in message.content + ] + + +def _request_text(request): + return "\n".join(message.content or "" for message in request) + + +def test_three_malformed_responses_stop_without_a_fourth_request(): + result, responder, registry, _agent, shared_state = _run_sequence( + [MALFORMED_RESPONSE] * 3 + ) + + assert result["success"] is False + assert "3 consecutive times" in result["reason"] + assert len(responder.requests) == 3 + assert registry.executed == [] + + corrections = _correction_messages(shared_state) + assert len(corrections) == 2 + assert "attempt 1/3" in corrections[0] + assert "attempt 2/3" in corrections[1] + assert "" in corrections[0] + assert " + +<｜DSML｜ name="index">189 + +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertIn("Profile tab", result.thought) + self.assertEqual(result.calls, []) + + def test_classifies_captured_hybrid_dsml_as_malformed(self): + text = """I will return to Settings home. + + +<｜｜DSML｜｜rameter name="button">back + +""" + + result = parse_tool_calls_detailed(text) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_rejects_dsml_corruption_inside_parameter_value(self): + text = """Settings home is visible. Marking the task complete. + + +true +Completed 20 verified cycles. + +""" + + result = parse_tool_calls_detailed(text, {"success": "boolean"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_allows_ordinary_dsml_text_inside_parameter_value(self): + text = """I will type the literal markup. + + +literal payload + +""" + + result = parse_tool_calls_detailed(text) + + self.assertEqual(result.status, ToolCallParseStatus.VALID) + self.assertEqual(len(result.calls), 1) + self.assertEqual( + result.calls[0].parameters, + {"text": 'literal payload'}, + ) + + def test_allows_unmatched_tool_like_text_inside_parameter_value(self): + text = """I will type the literal snippet. + + +literal token + +""" + + result = parse_tool_calls_detailed(text) + + self.assertEqual(result.status, ToolCallParseStatus.VALID) + self.assertEqual(len(result.calls), 1) + self.assertEqual( + result.calls[0].parameters, + {"text": 'literal token'}, + ) + + def test_classifies_dsml_without_xml_wrapper_as_malformed(self): + text = """I will tap the target. +<｜DSML｜tool_calls> +<｜DSML｜invoke name="click"> +<｜DSML｜parameter name="index">12 + +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.thought, "I will tap the target.") + self.assertEqual(result.calls, []) + + def test_classifies_standalone_invoke_as_malformed(self): + text = """I will tap the target. +12""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_classifies_noncanonical_wrapper_as_malformed(self): + text = """I will tap the target. + +12 +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_classifies_missing_close_tag_as_malformed(self): + text = """I will tap the target. + +12""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_classifies_empty_wrapper_as_malformed(self): + result = parse_tool_calls_detailed("\n") + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_classifies_wrapper_without_named_invoke_as_malformed(self): + text = """ +12 +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + + def test_classifies_valid_xml_as_valid(self): + text = """I will tap the target. + +12 +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.VALID) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].parameters, {"index": 12}) + + def test_valid_sibling_block_wins_over_malformed_block(self): + text = """I will retry and then complete. + +<｜｜DSML｜｜ name="click"> +12 + + + + +true +Done + +""" + + result = parse_tool_calls_detailed(text, {"success": "boolean"}) + + self.assertEqual(result.status, ToolCallParseStatus.VALID) + self.assertEqual([call.name for call in result.calls], ["complete"]) + self.assertEqual( + result.calls[0].parameters, + {"success": True, "message": "Done"}, + ) + + def test_argument_error_is_valid_markup(self): + text = """ +not-a-number +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.VALID) + self.assertEqual(len(result.calls), 1) + self.assertIsNotNone(result.calls[0].error) + + def test_public_tuple_parser_remains_compatible(self): + text = """I will tap the target. + +12 +""" + + thought, calls = parse_tool_calls(text, {"index": "number"}) + detailed = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual((thought, calls), (detailed.thought, detailed.calls)) + + def test_public_tuple_parser_preserves_marker_only_text(self): + cases = [ + "I will tap.\n<｜DSML｜tool_calls>", + 'I will tap.\n', + "I will tap.\n", + ] + + for text in cases: + with self.subTest(text=text): + thought, calls = parse_tool_calls(text) + self.assertEqual(thought, text.strip()) + self.assertEqual(calls, []) + def test_drops_adjacent_exact_duplicate_tool_calls(self): text = """ I will tap the target. From 2c0271dd883e6a05c086f27cd9cf00c35a814070 Mon Sep 17 00:00:00 2001 From: "rasul.osmanbayli" Date: Mon, 3 Aug 2026 19:43:08 +0400 Subject: [PATCH 2/2] Preserve literal DSML parameter payloads --- mobilerun/agent/fast_agent/xml_parser.py | 28 +++++++++++-- tests/test_fast_agent_xml_parser.py | 50 ++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 3 deletions(-) diff --git a/mobilerun/agent/fast_agent/xml_parser.py b/mobilerun/agent/fast_agent/xml_parser.py index e311b888..26930dd0 100644 --- a/mobilerun/agent/fast_agent/xml_parser.py +++ b/mobilerun/agent/fast_agent/xml_parser.py @@ -23,6 +23,7 @@ r'()(.*?)()', re.DOTALL, ) +_PARAM_OPEN_RE = re.compile(r"/])") _DSML_MARKUP_PATTERN = r"(?:<|＜)\s*/?\s*｜+DSML｜+" _DSML_MARKUP_RE = re.compile(_DSML_MARKUP_PATTERN, re.IGNORECASE) @@ -121,9 +122,9 @@ def parse_tool_calls_detailed( continue # Parameter-value sanitization intentionally preserves raw text such as - # code, but it must never turn provider-specific tool markup into an - # executable XML call. A separate valid wrapper can still win below. - if _DSML_MARKUP_RE.search(block): + # code and literal DSML. Provider-specific markup is malformed only + # when it occurs outside a trustworthy canonical parameter payload. + if _has_structural_dsml_markup(block): continue calls = _parse_tool_call_block(block, param_types) @@ -226,6 +227,27 @@ def _parse_tool_call_block( return calls +def _has_structural_dsml_markup(block: str) -> bool: + """Return whether DSML occurs outside a trustworthy parameter payload.""" + payload_spans: List[Tuple[int, int]] = [] + for parameter in _PARAM_RE.finditer(block): + payload = parameter.group(2) + # A nested parameter opener means the regex may have crossed + # from a missing close tag into a later parameter. Do not hide DSML in + # that span from structural validation. + if _PARAM_OPEN_RE.search(payload): + continue + payload_spans.append(parameter.span(2)) + + for marker in _DSML_MARKUP_RE.finditer(block): + if not any( + start <= marker.start() and marker.end() <= end + for start, end in payload_spans + ): + return True + return False + + def _drop_adjacent_duplicate_blocks( blocks: List[List[ToolCall]], ) -> List[List[ToolCall]]: diff --git a/tests/test_fast_agent_xml_parser.py b/tests/test_fast_agent_xml_parser.py index 08cab4fc..0382458f 100644 --- a/tests/test_fast_agent_xml_parser.py +++ b/tests/test_fast_agent_xml_parser.py @@ -59,6 +59,20 @@ def test_rejects_dsml_corruption_inside_parameter_value(self): self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) self.assertEqual(result.calls, []) + def test_rejects_dsml_corruption_crossing_xml_parameter_variant(self): + text = """Settings home is visible. Marking the task complete. + + +true +Completed 20 verified cycles. + +""" + + result = parse_tool_calls_detailed(text, {"success": "boolean"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + def test_allows_ordinary_dsml_text_inside_parameter_value(self): text = """I will type the literal markup. @@ -76,6 +90,42 @@ def test_allows_ordinary_dsml_text_inside_parameter_value(self): {"text": 'literal payload'}, ) + def test_allows_literal_dsml_syntax_inside_parameter_value(self): + text = """I will type the literal DSML syntax. + + +before <｜DSML｜tool_calls>literal after + +""" + + result = parse_tool_calls_detailed(text) + + self.assertEqual(result.status, ToolCallParseStatus.VALID) + self.assertEqual(len(result.calls), 1) + self.assertEqual( + result.calls[0].parameters, + {"text": "before <｜DSML｜tool_calls>literal after"}, + ) + thought, calls = parse_tool_calls(text) + self.assertEqual(thought, result.thought) + self.assertEqual(calls, result.calls) + + def test_literal_dsml_payload_does_not_hide_structural_dsml(self): + text = """I will type text and then tap. + + +literal <｜DSML｜payload> + +<｜DSML｜ name="click"> +12 + +""" + + result = parse_tool_calls_detailed(text, {"index": "number"}) + + self.assertEqual(result.status, ToolCallParseStatus.MALFORMED) + self.assertEqual(result.calls, []) + def test_allows_unmatched_tool_like_text_inside_parameter_value(self): text = """I will type the literal snippet.