Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions mobilerun/agent/fast_agent/events.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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):
Expand Down
75 changes: 72 additions & 3 deletions mobilerun/agent/fast_agent/fast_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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 (
Expand Down Expand Up @@ -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)
Expand All @@ -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 = (
Expand Down
80 changes: 79 additions & 1 deletion mobilerun/agent/fast_agent/xml_parser.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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."""
Expand All @@ -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()

Expand All @@ -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:
Expand Down Expand Up @@ -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
Comment on lines +238 to +239

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Preserve DSML alongside parameter-like payload text

When a legitimate parameter value contains both literal DSML and parameter-like text—for example, type_text with Example: <parameter name="x"> then <|DSML|parameter>—this nested-opener check excludes the entire canonical payload span, so the subsequent DSML scan classifies the valid call as malformed and the agent eventually aborts without typing it. Fresh evidence in the updated code is that literal DSML alone is now preserved, but combining it with the raw <parameter...> content that _sanitize_param_content otherwise supports still triggers the same failure; avoid treating every nested parameter opener as proof that the regex crossed a corrupted boundary.

Useful? React with 👍 / 👎.

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]]:
Expand Down
Loading
Loading