fix: simplify snapshot cleaning to only update function result message

Address reviewer feedback:
- Remove unnecessary recreation of assistant message (already recorded
  separately by _record_assistant_tool_calls_from_requests)
- Fixes format mismatch issue (FunctionCallingMessage vs BaseMessage)
- Fixes multi-tool-call support (assistant message preserves all tool calls)
- Add test for snapshot cleaning functionality

Author: Chase-Xuu

camel/agents/chat_agent.py

@@ -169,8 +169,6 @@ class _ToolOutputHistoryEntry:
     result_text: str
     record_uuids: List[str]
     record_timestamps: List[float]
-    args: Optional[Dict[str, Any]] = None
-    extra_content: Optional[Dict[str, Any]] = None
     cached: bool = False
 
 
@@ -1447,8 +1445,6 @@ def _register_tool_output_for_cache(
         tool_call_id: str,
         result_text: str,
         records: List[MemoryRecord],
-        args: Optional[Dict[str, Any]] = None,
-        extra_content: Optional[Dict[str, Any]] = None,
     ) -> None:
         if not records:
             return
@@ -1459,8 +1455,6 @@ def _register_tool_output_for_cache(
             result_text=result_text,
             record_uuids=[str(record.uuid) for record in records],
             record_timestamps=[record.timestamp for record in records],
-            args=args,
-            extra_content=extra_content,
         )
         self._tool_output_history.append(entry)
         self._process_tool_output_cache()
@@ -1502,66 +1496,31 @@ def _clean_snapshot_in_memory(
                 if record["uuid"] not in entry.record_uuids
             ]
 
-            # Recreate both assistant and function messages
-            # For Gemini API, the tool call request needs to exist along with
-            # the tool call result.
-            assist_args = entry.args if entry.args is not None else {}
-
-            # Use timestamps from entry, ensuring proper ordering
-            if len(entry.record_timestamps) >= 2:
-                assist_timestamp = entry.record_timestamps[0]
-                func_timestamp = entry.record_timestamps[1]
-            elif entry.record_timestamps:
-                # If only one timestamp, use it for function and ensure
-                # assist comes before
-                func_timestamp = entry.record_timestamps[0]
-                assist_timestamp = func_timestamp - 1e-6
-            else:
-                # No timestamps, use current time with proper ordering
-                assist_timestamp = time.time_ns() / 1_000_000_000
-                func_timestamp = assist_timestamp + 1e-6
-
-            # Recreate assistant message (tool call request)
-            cleaned_assist_message = FunctionCallingMessage(
-                role_name=self.role_name,
-                role_type=self.role_type,
-                meta_dict={},
-                content="",
-                func_name=entry.tool_name,
-                args=assist_args,
-                tool_call_id=entry.tool_call_id,
-                extra_content=entry.extra_content,
+            # Recreate only the function result message with cleaned content.
+            # The assistant message with tool calls is already recorded
+            # separately by _record_assistant_tool_calls_from_requests and
+            # should not be modified here.
+            timestamp = (
+                entry.record_timestamps[0]
+                if entry.record_timestamps
+                else time.time_ns() / 1_000_000_000
             )
-
-            # Recreate function message (tool result)
-            cleaned_func_message = FunctionCallingMessage(
+            cleaned_message = FunctionCallingMessage(
                 role_name=self.role_name,
                 role_type=self.role_type,
                 meta_dict={},
                 content="",
                 func_name=entry.tool_name,
                 result=cleaned_result,
                 tool_call_id=entry.tool_call_id,
-                extra_content=entry.extra_content,
-            )
-
-            # Create new records for both assistant and function messages
-            new_assist_record = MemoryRecord(
-                message=cleaned_assist_message,
-                role_at_backend=OpenAIBackendRole.ASSISTANT,
-                timestamp=assist_timestamp,
-                agent_id=self.agent_id,
             )
-            new_func_record = MemoryRecord(
-                message=cleaned_func_message,
+            new_record = MemoryRecord(
+                message=cleaned_message,
                 role_at_backend=OpenAIBackendRole.FUNCTION,
-                timestamp=func_timestamp,
+                timestamp=timestamp,
                 agent_id=self.agent_id,
             )
-
-            # Add both records back
-            updated_records.append(new_assist_record.to_dict())
-            updated_records.append(new_func_record.to_dict())
+            updated_records.append(new_record.to_dict())
             updated_records.sort(key=lambda record: record["timestamp"])
             storage.clear()
             storage.save(updated_records)
@@ -1573,11 +1532,8 @@ def _clean_snapshot_in_memory(
             )
 
             entry.cached = True
-            entry.record_uuids = [
-                str(new_assist_record.uuid),
-                str(new_func_record.uuid),
-            ]
-            entry.record_timestamps = [assist_timestamp, func_timestamp]
+            entry.record_uuids = [str(new_record.uuid)]
+            entry.record_timestamps = [timestamp]
 
     def add_external_tool(
         self, tool: Union[FunctionTool, Callable, Dict[str, Any]]
@@ -4128,17 +4084,13 @@ def _record_tool_calling(
         )
 
         # Register tool output for snapshot cleaning if enabled
-        # Include args and extra_content so both the assistant message
-        # (tool call) and function message (tool result) can be recreated
         if self._enable_snapshot_clean and not mask_output and func_records:
             serialized_result = self._serialize_tool_result(result_for_memory)
             self._register_tool_output_for_cache(
                 func_name,
                 tool_call_id,
                 serialized_result,
                 cast(List[MemoryRecord], func_records),
-                args=args,
-                extra_content=extra_content,
             )
 
         if isinstance(result, ToolResult) and result.images:

test/agents/test_chat_agent.py

@@ -1837,3 +1837,84 @@ class MathResult(BaseModel):
     assert len(responses) > 1, "Should receive multiple streaming chunks"
     assert responses[-1].msg.parsed.answer == 6
     assert responses[-1].msg.parsed.explanation
+
+
+def test_clean_snapshot_in_memory():
+    """Test that snapshot content is properly cleaned in memory.
+
+    This tests the _clean_snapshot_in_memory functionality which removes
+    stale snapshot markers and references from tool output messages stored
+    in memory. The cleaning preserves the assistant message (tool call request)
+    and only updates the function result message.
+    """
+    from unittest.mock import MagicMock, patch
+
+    from camel.agents.chat_agent import _ToolOutputHistoryEntry
+    from camel.messages import FunctionCallingMessage
+
+    # Create a mock model to avoid API calls
+    mock_model = MagicMock()
+    mock_model.model_type = ModelType.DEFAULT
+    mock_model.model_config_dict = {}
+    mock_model.token_counter = None
+    mock_model.model_platform_name = "openai"
+
+    with patch.object(ChatAgent, '_init_model', return_value=mock_model):
+        agent = ChatAgent(
+            system_message="Test agent",
+            model=mock_model,
+        )
+
+    # Manually enable snapshot cleaning
+    agent._enable_snapshot_clean = True
+    agent._tool_output_history = []
+
+    # Create a mock memory storage
+    mock_storage = MagicMock()
+    mock_chat_history_block = MagicMock()
+    mock_chat_history_block.storage = mock_storage
+
+    agent.memory._chat_history_block = mock_chat_history_block
+
+    # Create a test entry with snapshot markers
+    test_uuid = "test-uuid-123"
+    test_timestamp = 1234567890.0
+    entry = _ToolOutputHistoryEntry(
+        tool_name="test_tool",
+        tool_call_id="call_123",
+        result_text="- Item 1 [ref=abc]\n- Item 2 [ref=def]\n",
+        record_uuids=[test_uuid],
+        record_timestamps=[test_timestamp],
+        cached=False,
+    )
+    agent._tool_output_history.append(entry)
+
+    # Mock the storage to return the existing record
+    mock_storage.load.return_value = [
+        {
+            "uuid": test_uuid,
+            "timestamp": test_timestamp,
+            "message": {"content": "- Item 1 [ref=abc]\n- Item 2 [ref=def]\n"},
+        }
+    ]
+
+    # Call the clean function
+    agent._clean_snapshot_in_memory(entry)
+
+    # Verify storage was updated
+    assert mock_storage.clear.called
+    assert mock_storage.save.called
+
+    # Get the saved records
+    saved_records = mock_storage.save.call_args[0][0]
+
+    # Should have one record (the cleaned function result)
+    assert len(saved_records) == 1
+
+    # The record should be a function message with cleaned content
+    saved_record = saved_records[0]
+    assert saved_record["role_at_backend"] == OpenAIBackendRole.FUNCTION.value
+
+    # Verify entry was marked as cached
+    assert entry.cached is True
+    assert len(entry.record_uuids) == 1  # Single function result record