diff --git a/agent/deep_assistant.py b/agent/deep_assistant.py
index e973e67..49640e5 100644
--- a/agent/deep_assistant.py
+++ b/agent/deep_assistant.py
@@ -462,7 +462,7 @@ def create_custom_cli_agent(
         # Add memory middleware
         if enable_memory:
             agent_middleware.append(
-                AgentMemoryMiddleware(settings=settings, assistant_id=assistant_id)
+                CustomAgentMemoryMiddleware(settings=settings, assistant_id=assistant_id)
             )
 
         # Add skills middleware
diff --git a/agent/prompt_loader.py b/agent/prompt_loader.py
index e7077c9..4e88b0d 100644
--- a/agent/prompt_loader.py
+++ b/agent/prompt_loader.py
@@ -137,7 +137,7 @@ async def load_system_prompt_async(config) -> str:
 
 
 
-def replace_mcp_placeholders(mcp_settings: List[Dict], dataset_dir: str, bot_id: str) -> List[Dict]:
+def replace_mcp_placeholders(mcp_settings: List[Dict], dataset_dir: str, bot_id: str, dataset_ids: List[str]) -> List[Dict]:
     """
     替换 MCP 配置中的占位符
     """
@@ -150,15 +150,15 @@ def replace_mcp_placeholders(mcp_settings: List[Dict], dataset_dir: str, bot_id:
             for key, value in obj.items():
                 if key == 'args' and isinstance(value, list):
                     # 特别处理 args 列表
-                    obj[key] = [item.format(dataset_dir=dataset_dir, bot_id=bot_id) if isinstance(item, str) else item
+                    obj[key] = [item.format(dataset_dir=dataset_dir, bot_id=bot_id, dataset_ids=','.join(dataset_ids)) if isinstance(item, str) else item
                                 for item in value]
                 elif isinstance(value, (dict, list)):
                     obj[key] = replace_placeholders_in_obj(value)
                 elif isinstance(value, str):
-                    obj[key] = value.format(dataset_dir=dataset_dir, bot_id=bot_id)
+                    obj[key] = value.format(dataset_dir=dataset_dir, bot_id=bot_id, dataset_ids=','.join(dataset_ids))
         elif isinstance(obj, list):
             return [replace_placeholders_in_obj(item) if isinstance(item, (dict, list)) else 
-                    item.format(dataset_dir=dataset_dir, bot_id=bot_id) if isinstance(item, str) else item
+                    item.format(dataset_dir=dataset_dir, bot_id=bot_id, dataset_ids=','.join(dataset_ids)) if isinstance(item, str) else item
                     for item in obj]
         return obj
     
@@ -183,6 +183,7 @@ async def load_mcp_settings_async(config) -> List[Dict]:
     project_dir = getattr(config, 'project_dir', None)
     mcp_settings = getattr(config, 'mcp_settings', None)
     bot_id = getattr(config, 'bot_id', '')
+    dataset_ids = getattr(config, 'dataset_ids', [])
 
     # 1. ============ 首先合并skill目录下的plugin.json配置（不使用缓存，确保改动生效）============
     skill_mcp_settings = await merge_skill_mcp_configs(bot_id)
@@ -265,7 +266,7 @@ async def load_mcp_settings_async(config) -> List[Dict]:
     # 替换 MCP 配置中的 {dataset_dir} 占位符
     if dataset_dir is None:
         dataset_dir = ""
-    merged_settings = replace_mcp_placeholders(merged_settings, dataset_dir, bot_id)
+    merged_settings = replace_mcp_placeholders(merged_settings, dataset_dir, bot_id, dataset_ids)
     return merged_settings