1
@self._logged_tool(name='context', description='【防遗忘统一入口】上下文管理')
2
async def context(
3
    action: str = Field(description='操作类型：get_planner_summary/get_reflector_summary/add_objective/clear'),
4
    objective: Optional[str] = Field(default=None, description='长期目标（add_objective 时必填）'),
5
    goal: Optional[str] = Field(default=None, description='长期目标别名')
6
) -> str:
7
    """防遗忘统一入口 - 管理上下文和防止遗忘"""
8
    try:
9
        if not self.context_manager:
10
            return "错误：context_manager 未初始化"
11

12
        if goal and not objective:
13
            objective = goal
14

15
        if action == 'get_planner_summary':
16
            # 处理子命令 1
17
            summary = self.context_manager.get_planner_summary()
18
            return format_summary(summary)
19

20
        elif action == 'get_reflector_summary':
21
            # 处理子命令 2
22
            summary = self.context_manager.get_reflector_summary()
23
            return format_summary(summary)
24

25
        elif action == 'add_objective':
26
            # 处理子命令 3（需要额外参数）
27
            if not objective:
28
                return "错误：add_objective 操作需要 objective 参数"
29
            self.context_manager.add_long_term_objective(objective)
30
            return f"✅ 长期目标已添加：{objective}"
31

32
        elif action == 'clear':
33
            # 处理子命令 4
34
            self.context_manager.clear_planner_context()
35
            self.context_manager.clear_reflector_context()
36
            return "✅ 所有上下文已清空"
37

38
        else:
39
            return f"错误：未知的操作类型 '{action}'"
40

41
    except Exception as e:
42
        return f"错误：{str(e)}"

1
@self._logged_tool(name='planner', description='【逻辑断层统一入口】规划器')
2
async def planner(
3
    action: str = Field(description='操作类型：initial_plan/dynamic_plan/get_history/get_rejected/add_rejected'),
4
    goal: Optional[str] = Field(default=None, description='目标（initial_plan/dynamic_plan 时必填）'),
5
    context: Optional[str] = Field(default=None, description='上下文（JSON 格式字符串）'),
6
    strategy_name: Optional[str] = Field(default=None, description='策略名称（add_rejected 时必填）'),
7
    reason: Optional[str] = Field(default=None, description='原因（add_rejected 时必填）')
8
) -> str:
9
    """逻辑断层统一入口 - 规划器操作"""
10
    try:
11
        if action == 'initial_plan':
12
            # 验证必填参数
13
            if not goal:
14
                return "错误：initial_plan 操作需要 goal 参数"
15
            ctx = json.loads(context) if context else None
16
            result = self.planner.initial_plan(goal, ctx)
17
            return format_result(result)
18

19
        elif action == 'add_rejected':
20
            # 验证多个必填参数
21
            if not strategy_name or not reason:
22
                return "错误：add_rejected 操作需要 strategy_name 和 reason 参数"
23
            self.planner.add_rejected_strategy(strategy_name, reason)
24
            return f"✅ 策略已标记为拒绝：{strategy_name}"

1
def format_result(result: Dict) -> str:
2
    """统一的结果格式化"""
3
    lines = ["=" * 70, "📋 标题", "=" * 70]
4

5
    # 添加关键指标
6
    lines.append(f"指标 1: {result.get('key_metric_1')}")
7
    lines.append(f"指标 2: {result.get('key_metric_2')}")
8

9
    # 添加列表内容
10
    if 'items' in result:
11
        lines.append("\n【列表内容】:")
12
        for item in result['items']:
13
            lines.append(f"  - {item}")
14

15
    lines.append("=" * 70)
16
    return "\n".join(lines)

1
def _logged_tool(self, name, description):
2
    """带日志记录的工具装饰器工厂"""
3
    def decorator(func):
4
        # 1. 创建带日志包装的工具函数
5
        logged_func = self._create_logged_tool(name, func)
6

7
        # 2. 注册到 FastMCP 服务器
8
        self.server.tool(name=name, description=description)(logged_func)
9

10
        # 3. 自动记录到 registered_tools 字典
11
        if not hasattr(self, 'registered_tools'):
12
            self.registered_tools = {}
13
        self.registered_tools[name] = logged_func
14

15
        return logged_func
16
    return decorator

1
def _map_tool_arguments(self, tool_name: str, args: Dict[str, Any]) -> Dict[str, Any]:
2
    """根据工具名映射参数"""
3
    mapped_args = args.copy()
4
    tool_name_lower = tool_name.lower()
5

6
    # 处理 target_url -> url（browser, http_request 等）
7
    if 'target_url' in mapped_args:
8
        if 'browser' in tool_name_lower or 'http_request' in tool_name_lower:
9
            mapped_args['url'] = mapped_args.pop('target_url')
10

11
    # 处理 target -> goal（ai_pentest）
12
    if 'target' in mapped_args:
13
        if 'ai_pentest' in tool_name_lower:
14
            mapped_args['goal'] = mapped_args.pop('target')
15

16
    return mapped_args

1
async def call_tool(
2
    self,
3
    tool_name: str,
4
    args: Optional[Dict[str, Any]] = None,
5
    max_retries: int = 3,
6
    timeout: int = 300
7
) -> Dict[str, Any]:
8
    """调用 MCP 工具（带重试机制）"""
9

10
    tool_func = self.registered_tools[tool_name]
11
    args = args or {}
12

13
    for attempt in range(1, max_retries + 1):
14
        try:
15
            if asyncio.iscoroutinefunction(tool_func):
16
                result = await asyncio.wait_for(
17
                    tool_func(**args),
18
                    timeout=timeout
19
                )
20
            else:
21
                result = await asyncio.wait_for(
22
                    asyncio.get_event_loop().run_in_executor(
23
                        None, lambda: tool_func(**args)
24
                    ),
25
                    timeout=timeout
26
                )
27

28
            return result
29

30
        except asyncio.TimeoutError:
31
            if attempt == max_retries:
32
                return {
33
                    "success": False,
34
                    "error": f"工具调用超时：{tool_name} (>{timeout}s)",
35
                    "attempts": attempt
36
                }
37
            await asyncio.sleep(1 * attempt)  # 指数退避
38

39
        except Exception as e:
40
            if attempt == max_retries:
41
                return {
42
                    "success": False,
43
                    "error": f"工具执行失败：{str(e)}",
44
                    "attempts": attempt
45
                }
46
            await asyncio.sleep(1 * attempt)

1
# ✅ 使用枚举式字符串
2
action: str = Field(
3
    description='操作类型：search/add/get_detail/export'
4
)
5

6
# ✅ 在 docstring 中明确说明
7
"""
8
知识库统一入口 - 整合 6 个知识库相关工具
9

10
操作类型:
11
  - search: 搜索知识
12
  - add: 添加知识
13
  - get_detail: 获取详情
14
  - export: 导出知识库
15
"""

1
# ✅ 统一的错误处理模式
2
try:
3
    # 业务逻辑
4
    if not self.component:
5
        return f"错误：{component_name} 未初始化"
6

7
    # 参数验证
8
    if action == 'xxx' and not required_param:
9
        return f"错误：{action}操作需要{param_name}参数"
10

11
    # 执行逻辑
12
    result = self.component.do_something()
13
    return format_result(result)
14

15
except Exception as e:
16
    return f"错误：{str(e)}"