Merge pull request #923 from inclusionAI/tool_match

Handle incomplete tool-call history during replay

Author: wuman001

aworld/agents/llm_agent.py

@@ -865,17 +865,45 @@ async def async_messages_transform(self,
             tool_calls_map = {}
             last_tool_calls = []
             matched_tool_call_ids = set()
+
+            def _is_tool_history(history) -> bool:
+                if isinstance(history, MemoryMessage):
+                    return isinstance(history, MemoryToolMessage)
+                return history.metadata.get('role') == 'tool'
+
+            def _drop_incomplete_tool_call_turn(reason: str):
+                nonlocal tool_calls_map, last_tool_calls
+                if not last_tool_calls:
+                    return
+                dropped_message = None
+                if messages and messages[-1].get("role") == "assistant" and messages[-1].get("tool_calls"):
+                    dropped_message = messages.pop()
+                logger.warning(
+                    "Skip incomplete tool-call turn in memory replay: "
+                    f"reason={reason}, missing_tool_call_ids={last_tool_calls}, "
+                    f"matched_tool_result_ids={list(tool_calls_map.keys())}, "
+                    f"dropped_assistant_message={bool(dropped_message)}, agent={self.id()}"
+                )
+                tool_calls_map = {}
+                last_tool_calls = []
+
+            def _append_complete_tool_results():
+                nonlocal tool_calls_map, last_tool_calls
+                for tool_call_id in last_tool_calls:
+                    if tool_call_id not in tool_calls_map:
+                        _drop_incomplete_tool_call_turn(f"missing tool result for {tool_call_id}")
+                        return
+                    messages.append(tool_calls_map.get(tool_call_id))
+                    matched_tool_call_ids.add(tool_call_id)
+                tool_calls_map = {}
+                last_tool_calls = []
+
             for history in histories:
                 if len(last_tool_calls) > 0 and len(tool_calls_map) == len(last_tool_calls):
                     # Maintain the order of tool calls
-                    for tool_call_id in last_tool_calls:
-                        if tool_call_id not in tool_calls_map:
-                            raise AWorldRuntimeException(
-                                f"tool_calls mismatch! {tool_call_id} not found in {tool_calls_map}, last_tool_calls: {last_tool_calls}, messages: {messages}, histories: {histories}")
-                        messages.append(tool_calls_map.get(tool_call_id))
-                        matched_tool_call_ids.add(tool_call_id)
-                    tool_calls_map = {}
-                    last_tool_calls = []
+                    _append_complete_tool_results()
+                elif last_tool_calls and not _is_tool_history(history):
+                    _drop_incomplete_tool_call_turn("next non-tool message encountered")
 
                 if isinstance(history, MemoryMessage):
                     if isinstance(history, MemoryToolMessage):
@@ -887,9 +915,10 @@ async def async_messages_transform(self,
                                 f"tool_call_id={history.tool_call_id}, agent={self.id()}"
                             )
                         else:
-                            raise AWorldRuntimeException(
-                                f"tool_calls mismatch! {history.tool_call_id} not found in {last_tool_calls}, "
-                                f"messages: {messages}, histories: {histories}")
+                            logger.warning(
+                                "Skip orphan tool result in memory replay: "
+                                f"tool_call_id={history.tool_call_id}, agent={self.id()}"
+                            )
                     else:
                         messages.append(history.to_openai_message())
                         if isinstance(history, MemoryAIMessage) and history.tool_calls:
@@ -908,9 +937,10 @@ async def async_messages_transform(self,
                                 f"tool_call_id={tool_call_id}, agent={self.id()}"
                             )
                         else:
-                            raise AWorldRuntimeException(
-                                f"tool_calls mismatch! {tool_call_id} not found in {last_tool_calls}, "
-                                f"messages: {messages}, histories: {histories}")
+                            logger.warning(
+                                "Skip orphan tool result in memory replay: "
+                                f"tool_call_id={tool_call_id}, agent={self.id()}"
+                            )
                     else:
                         if not self.use_tools_in_prompt and history.metadata.get('tool_calls'):
                             messages.append({'role': history.metadata['role'], 'content': history.content,
@@ -922,31 +952,13 @@ async def async_messages_transform(self,
                                              "tool_call_id": history.metadata.get("tool_call_id")})
                 if len(last_tool_calls) > 0 and len(tool_calls_map) == len(last_tool_calls):
                     # Maintain the order of tool calls
-                    for tool_call_id in last_tool_calls:
-                        if tool_call_id not in tool_calls_map:
-                            raise AWorldRuntimeException(
-                                f"tool_calls mismatch! {tool_call_id} not found in {tool_calls_map}, last_tool_calls: {last_tool_calls}, messages: {messages}, histories: {histories}")
-                        messages.append(tool_calls_map.get(tool_call_id))
-                        matched_tool_call_ids.add(tool_call_id)
-                    tool_calls_map = {}
-                    last_tool_calls = []
+                    _append_complete_tool_results()
                 elif len(tool_calls_map) > len(last_tool_calls):
-                    raise AWorldRuntimeException(
-                        f"tool_calls mismatch! {len(tool_calls_map)} tool messages > {len(last_tool_calls)} tool calls: "
-                        f"tool_calls_map={tool_calls_map}, last_tool_calls={last_tool_calls}, messages={messages}, histories={histories}")
-            if len(tool_calls_map) == len(last_tool_calls):
-                for tool_call_id in last_tool_calls:
-                    if tool_call_id not in tool_calls_map:
-                        raise AWorldRuntimeException(
-                            f"tool_calls mismatch! {tool_call_id} not found in {tool_calls_map}, last_tool_calls: {last_tool_calls}, messages: {messages}, histories: {histories}")
-                    messages.append(tool_calls_map.get(tool_call_id))
-                    matched_tool_call_ids.add(tool_call_id)
-                tool_calls_map = {}
-                last_tool_calls = []
+                    _drop_incomplete_tool_call_turn("more tool results than tool calls")
+            if last_tool_calls and len(tool_calls_map) == len(last_tool_calls):
+                _append_complete_tool_results()
             else:
-                raise AWorldRuntimeException(
-                    f"tool_calls mismatch! {len(tool_calls_map)} tool messages != {len(last_tool_calls)} tool calls: "
-                    f"tool_calls_map={tool_calls_map}, last_tool_calls={last_tool_calls}, messages={messages}, histories={histories}")
+                _drop_incomplete_tool_call_turn("end of history reached")
 
         return messages
 

tests/runners/test_memory_tool_result_compaction.py

@@ -304,6 +304,96 @@ async def test_llm_message_replay_skips_duplicate_tool_result(monkeypatch):
     assert tool_messages[0]["content"] == [{"type": "text", "text": "first cron result"}]
 
 
+@pytest.mark.asyncio
+async def test_llm_message_replay_drops_incomplete_tool_call_turn(monkeypatch):
+    meta = MessageMetadata(
+        session_id="session-1",
+        user_id="user-1",
+        task_id="task-1",
+        agent_id="agent-1",
+        agent_name="Aworld",
+    )
+    ai_message = MemoryAIMessage(
+        content="",
+        tool_calls=[
+            ToolCall.from_dict({
+                "id": "cron__cron_tool:missing",
+                "function": {"name": "cron__cron_tool", "arguments": "{}"},
+            })
+        ],
+        metadata=meta,
+    )
+    fake_memory = _FakeMemory()
+    fake_memory.items = [(ai_message, None)]
+    monkeypatch.setattr(
+        "aworld.agents.llm_agent.MemoryFactory",
+        type("MemoryFactory", (), {"instance": staticmethod(lambda: fake_memory)}),
+    )
+
+    context = _build_context()
+    agent = LLMAgent(
+        name="Aworld",
+        agent_id="agent-1",
+        conf=AgentConfig(
+            llm_model_name="test-model",
+            llm_api_key="test-key",
+            memory_config=AgentMemoryConfig(history_rounds=10),
+        ),
+    )
+    message = Message(headers={"context": context})
+
+    messages = await agent.async_messages_transform(
+        image_urls=[],
+        observation=Observation(action_result=[ActionResult(content="continue current turn")]),
+        message=message,
+    )
+
+    assert not any(message.get("tool_calls") for message in messages)
+    assert not any(message.get("role") == "tool" for message in messages)
+
+
+@pytest.mark.asyncio
+async def test_llm_message_replay_skips_orphan_tool_result(monkeypatch):
+    meta = MessageMetadata(
+        session_id="session-1",
+        user_id="user-1",
+        task_id="task-1",
+        agent_id="agent-1",
+        agent_name="Aworld",
+    )
+    orphan_tool = MemoryToolMessage(
+        content="orphan result",
+        tool_call_id="missing-call",
+        metadata=meta,
+    )
+    fake_memory = _FakeMemory()
+    fake_memory.items = [(orphan_tool, None)]
+    monkeypatch.setattr(
+        "aworld.agents.llm_agent.MemoryFactory",
+        type("MemoryFactory", (), {"instance": staticmethod(lambda: fake_memory)}),
+    )
+
+    context = _build_context()
+    agent = LLMAgent(
+        name="Aworld",
+        agent_id="agent-1",
+        conf=AgentConfig(
+            llm_model_name="test-model",
+            llm_api_key="test-key",
+            memory_config=AgentMemoryConfig(history_rounds=10),
+        ),
+    )
+    message = Message(headers={"context": context})
+
+    messages = await agent.async_messages_transform(
+        image_urls=[],
+        observation=Observation(action_result=[ActionResult(content="continue current turn")]),
+        message=message,
+    )
+
+    assert not any(message.get("role") == "tool" for message in messages)
+
+
 @pytest.mark.asyncio
 async def test_default_memory_handler_compacts_large_tool_results_by_char_length(monkeypatch):
     fake_memory = _FakeMemory()