Merge pull request #411 from droidrun/fix/deepseek-malformed-tool-call-guard

Stop FastAgent loops on malformed tool calls

Author: RasulOs

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):

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"
+        "<function_calls>\n"
+        '<invoke name="tool_name">\n'
+        '<parameter name="parameter_name">value</parameter>\n'
+        "</invoke>\n"
+        "</function_calls>"
+    )
+
 
 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 <add_memory> from thought text and append to unified memory
         memory_update = extract_add_memory(thought)
@@ -396,17 +419,63 @@ 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
 
     @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 = (

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
 
@@ -22,6 +23,15 @@
     r'(<parameter\s+name="[^"]*">)(.*?)(</parameter>)',
     re.DOTALL,
 )
+_PARAM_OPEN_RE = re.compile(r"<parameter(?=[\s>/])")
+
+_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
@@ -33,6 +43,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 +86,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 +121,20 @@ def parse_tool_calls(
         if not block:
             continue
 
+        # Parameter-value sanitization intentionally preserves raw text such as
+        # 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)
         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:
@@ -170,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]]: