Coverage for merco/tools/middleware.py: 97%

1"""ToolContext + ToolMiddleware + ToolMiddlewareChain""" 

2from __future__ import annotations 

3 

4from abc import ABC, abstractmethod 

5from dataclasses import dataclass, field 

6from pathlib import Path 

7 

8from merco.sandbox.guard import GuardAction 

9from merco.sandbox.confirm import confirm_edit 

10import merco.sandbox.snapshot as snapshot 

11 

12 

13@dataclass 

14class ToolContext: 

15 """工具执行上下文，贯穿 before/during/after""" 

16 tool_name: str 

17 arguments: dict 

18 tool: object | None = None 

19 result: dict | None = None 

20 error: BaseException | None = None 

21 metadata: dict = field(default_factory=dict) 

22 

23 

24class ToolMiddleware(ABC): 

25 """工具中间件基类""" 

26 name: str = "" 

27 

28 @abstractmethod 

29 async def before(self, ctx: ToolContext): 

30 """执行前。返回 dict 短路；返回 ctx/None 继续""" 

31 ... 

32 

33 @abstractmethod 

34 async def after(self, ctx: ToolContext): 

35 """执行后。可修改 ctx.result；返回 dict 替换 result""" 

36 ... 

37 

38 @abstractmethod 

39 async def on_error(self, ctx: ToolContext): 

40 """异常处理。返回 dict → 错误结果；None → 抛""" 

41 ... 

42 

43 

44class ToolMiddlewareChain: 

45 """洋葱模型：before 正序，after/on_error 逆序""" 

46 

47 def __init__(self): 

48 self._middlewares: list[ToolMiddleware] = [] 

49 

50 def use(self, middleware: ToolMiddleware) -> "ToolMiddlewareChain": 

51 self._middlewares.append(middleware) 

52 return self 

53 

54 async def execute(self, ctx: ToolContext, call_tool) -> dict: 

55 for mw in self._middlewares: 

56 r = await mw.before(ctx) 

57 if isinstance(r, dict): 

58 return r 

59 if isinstance(r, ToolContext): 

60 ctx = r 

61 

62 try: 

63 result = call_tool() 

64 if hasattr(result, "__await__"): 

65 result = await result 

66 ctx.result = result 

67 except BaseException as e: 

68 ctx.error = e 

69 for mw in reversed(self._middlewares): 

70 r = await mw.on_error(ctx) 

71 if isinstance(r, dict): 

72 return r 

73 raise 

74 

75 for mw in reversed(self._middlewares): 

76 r = await mw.after(ctx) 

77 if isinstance(r, dict): 

78 ctx.result = r 

79 elif isinstance(r, ToolContext): 

80 ctx = r 

81 return ctx.result 

82 

83 

84class GuardMiddleware(ToolMiddleware): 

85 """包装 ToolGuard — DENY 返回错误，ASK 抛异常，ALLOW 继续""" 

86 name = "guard" 

87 

88 def __init__(self, guard): 

89 self.guard = guard 

90 

91 async def before(self, ctx: ToolContext): 

92 result = await self.guard.check(ctx.tool_name, ctx.arguments) 

93 if result.action == GuardAction.DENY: 

94 return {"error": f"操作被安全守卫拒绝: {result.reason}", "tool": ctx.tool_name} 

95 if result.action == GuardAction.ASK: 

96 from merco.sandbox.guard import GuardConfirmationRequired 

97 raise GuardConfirmationRequired(result) 

98 return None 

99 

100 async def after(self, ctx: ToolContext): 

101 return None 

102 

103 async def on_error(self, ctx: ToolContext): 

104 return None 

105 

106 

107class ErrorHandlingMiddleware(ToolMiddleware): 

108 """工具异常 → 结构化 tool_error 结果""" 

109 name = "error_handling" 

110 

111 async def before(self, ctx: ToolContext): 

112 return None 

113 

114 async def after(self, ctx: ToolContext): 

115 return None 

116 

117 async def on_error(self, ctx: ToolContext): 

118 from merco.tools.errors import tool_error 

119 return tool_error( 

120 ctx.error, 

121 ctx.tool_name, 

122 getattr(ctx.tool, 'parameters', None) if ctx.tool else None, 

123 ) 

124 

125 

126class EditApplyMiddleware(ToolMiddleware): 

127 """应用 EditFile planned_edit：确认、快照、写入""" 

128 name = "edit_apply" 

129 

130 def __init__(self, diff_view: str = "unified"): 

131 self.diff_view = diff_view 

132 

133 async def before(self, ctx: ToolContext): 

134 return None 

135 

136 async def after(self, ctx: ToolContext): 

137 result = ctx.result or {} 

138 if not isinstance(result, dict) or not result.get("planned_edit"): 

139 return None 

140 

141 path = result["path"] 

142 old_content = result["old_content"] 

143 new_content = result["new_content"] 

144 diff = result["diff"] 

145 

146 approved = await confirm_edit(diff, path, 1, old_content, new_content, self.diff_view) 

147 if not approved: 

148 return { 

149 "success": False, 

150 "path": path, 

151 "message": "用户已取消修改", 

152 "diff": diff, 

153 } 

154 

155 snapshot.track(path, old_content) 

156 Path(path).write_text(new_content, encoding="utf-8") 

157 return { 

158 "success": True, 

159 "path": path, 

160 "diff": diff, 

161 "message": f"已修改 `{path}`", 

162 } 

163 

164 async def on_error(self, ctx: ToolContext): 

165 return None