1"""

2 AgentOS v1.1.4 Agent Runtime Middleware Pipeline — 可组合的执行生命周期中间件。

3

4 在 Agent 执行的每个阶段（pre-LLM / post-LLM / pre-tool / post-tool）

5 插入策略检查、日志、脱敏、预算控制等拦截逻辑。

6

7 灵感来自 Microsoft Agent Framework 1.0 的 Middleware Pipeline 和 CrewAI Runtime Hooks。

8"""

9

10 from __future__ import annotations

11

12 from abc import ABC, abstractmethod

13 from dataclasses import dataclass, field

14 from enum import Enum

15 from typing import Any, Callable, Optional

16

17

18 class MiddlewarePhase(str, Enum):

19 """中间件触发阶段。"""

20

21 PRE_LLM = "pre_llm" # LLM 调用前

22 POST_LLM = "post_llm" # LLM 调用后、输出解析前

23 PRE_TOOL = "pre_tool" # 工具调用前

24 POST_TOOL = "post_tool" # 工具调用后

25 ON_ERROR = "on_error" # 执行出错时

26 ON_START = "on_start" # Agent 启动时

27 ON_COMPLETE = "on_complete" # Agent 执行完成时

28

29

30@dataclass

31 class MiddlewareContext:

32 """中间件执行上下文。"""

33

34 phase: MiddlewarePhase

35 agent_name: str = ""

36 run_id: str = ""

37 # LLM 阶段

38 prompt: Optional[str] = None

39 model_name: Optional[str] = None

40 llm_output: Optional[str] = None

41 # Tool 阶段

42 tool_name: Optional[str] = None

43 tool_args: Optional[dict] = None

44 tool_result: Any = None

45 # Error

46 error: Optional[Exception] = None

47 # 额外元数据

48 metadata: dict[str, Any] = field(default_factory=dict)

49

50

51@dataclass

52 class MiddlewareDecision:

53 """中间件决策结果。"""

54

55 allow: bool = True

56 """是否允许继续执行。"""

57

58 reason: str = ""

59 """决策理由。"""

60

61 modified_context: Optional[MiddlewareContext] = None

62 """修改后的上下文（如脱敏后的 prompt）。"""

63

64 action: str = "allow" # allow / warn / block / transform / escalate

65 """决策动作。"""

66

67 blocked_by: str = ""

68 """阻断方名称。"""

69

70

71 class AgentMiddleware(ABC):

72 """Agent 运行时中间件基类。

73

74 每个中间件声明自己监听的阶段，通过 process() 返回决策。

75 返回 MiddlewareDecision(allow=False) 阻断执行链。

76 """

77

78 name: str = "base_middleware"

79

80 @property

81 def phases(self) -> list[MiddlewarePhase]:

82 """返回此中间件监听的阶段列表。"""

83 return [MiddlewarePhase.PRE_LLM]

84

85 @abstractmethod

86 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision:

87 """处理中间件逻辑。返回决策。"""

88 ...

89

90

91# ── 内置中间件 ──────────────────────────────────────────────────────────────

92

93 class PIIMaskingMiddleware(AgentMiddleware):

94 """PII脱敏中间件：在 pre-LLM 阶段对 prompt 脱敏。"""

95

96 name = "pii_masking"

97

98 @property

99 def phases(self) -> list[MiddlewarePhase]:

100 return [MiddlewarePhase.PRE_LLM]

101

102 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision:

103 if not ctx.prompt:

104 return MiddlewareDecision(allow=True)

105 from agentos.security.guard import PIIDetector

106 detector = PIIDetector(auto_redact=True)

107 sanitized, items = detector.redact(ctx.prompt)

108 count = len(items)

109 if count > 0:

110 new_ctx = MiddlewareContext(**{**ctx.__dict__})

111 new_ctx.prompt = sanitized

112 new_ctx.metadata["pii_count"] = count

113 return MiddlewareDecision(

114 allow=True, action="transform",

115 reason=f"Masked {count} PII instances",

116 modified_context=new_ctx,

117 )

118 return MiddlewareDecision(allow=True)

119

120

121 class BudgetGuardMiddleware(AgentMiddleware):

122 """预算守护中间件：pre-LLM 阶段检查预算。"""

123

124 name = "budget_guard"

125

126 def __init__(self, tracker=None, budget_limit: float = 0.0, warn_ratio: float = 0.8):

127 self.tracker = tracker

128 self.budget_limit = budget_limit

129 self.warn_ratio = warn_ratio

130

131 @property

132 def phases(self) -> list[MiddlewarePhase]:

133 return [MiddlewarePhase.PRE_LLM, MiddlewarePhase.PRE_TOOL]

134

135 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision:

136 if not self.tracker or self.budget_limit <= 0:

137 return MiddlewareDecision(allow=True)

138 spent = self.tracker.total_cost

139 ratio = spent / self.budget_limit

140 if ratio >= 1.0:

141 return MiddlewareDecision(

142 allow=False, action="block",

143 reason=f"Budget exceeded: ${spent:.4f} / ${self.budget_limit:.2f}",

144 blocked_by=self.name,

145 )

146 if ratio >= self.warn_ratio:

147 return MiddlewareDecision(

148 allow=True, action="warn",

149 reason=f"Budget warning: {ratio:.0%} used (${spent:.4f} / ${self.budget_limit:.2f})",

150 )

151 return MiddlewareDecision(allow=True)

152

153

154 class ToolRiskGuardMiddleware(AgentMiddleware):

155 """工具风险守护中间件：pre-tool 阶段根据风险等级决定是否阻断。"""

156

157 name = "tool_risk_guard"

158

159 def __init__(self, max_auto_level: str = "medium"):

160 from agentos.tools.risk import ToolRiskLevel

161 self.max_auto_level = ToolRiskLevel(max_auto_level)

162

163 @property

164 def phases(self) -> list[MiddlewarePhase]:

165 return [MiddlewarePhase.PRE_TOOL]

166

167 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision:

168 if not ctx.tool_name:

169 return MiddlewareDecision(allow=True)

170

171 from agentos.tools.risk import infer_risk_level

172 risk = infer_risk_level(ctx.tool_name, tool_args=ctx.tool_args)

173

174 if risk.requires_user_confirm():

175 return MiddlewareDecision(

176 allow=False, action="escalate",

177 reason=f"Tool '{ctx.tool_name}' requires user approval: {risk.description}",

178 blocked_by=self.name,

179 )

180

181 levels = ["low", "medium", "high", "critical"]

182 if levels.index(risk.level.value) > levels.index(self.max_auto_level.value):

183 return MiddlewareDecision(

184 allow=False, action="block",

185 reason=f"Tool '{ctx.tool_name}' risk {risk.level.value} exceeds auto limit {self.max_auto_level.value}",

186 blocked_by=self.name,

187 )

188

189 return MiddlewareDecision(allow=True)

190

191

192 class AuditLogMiddleware(AgentMiddleware):

193 """审计日志中间件：在所有阶段记录审计轨迹。"""

194

195 name = "audit_log"

196

197 @property

198 def phases(self) -> list[MiddlewarePhase]:

199 return [

200 MiddlewarePhase.ON_START, MiddlewarePhase.PRE_LLM,

201 MiddlewarePhase.PRE_TOOL, MiddlewarePhase.POST_TOOL,

202 MiddlewarePhase.ON_ERROR, MiddlewarePhase.ON_COMPLETE,

203 ]

204

205 async def process(self, ctx: MiddlewareContext) -> MiddlewareDecision:

206 import logging

207 logger = logging.getLogger("agentos.audit")

208 logger.info(

209 f"[{ctx.phase.value}] agent={ctx.agent_name} run={ctx.run_id} "

210 f"tool={ctx.tool_name or '-'}"

211 )

212 return MiddlewareDecision(allow=True)

213

214

215# ── 中间件管道 ──────────────────────────────────────────────────────────────

216

217 class MiddlewarePipeline:

218 """编排多个中间件按阶段执行。

219

220 每个阶段：

221 1. 筛选监听该阶段的中间件

222 2. 按注册顺序依次执行

223 3. 任一返回 allow=False 即阻断

224 4. 若返回 modified_context 则传递给后续中间件

225 """

226

227 def __init__(self, middlewares: Optional[list[AgentMiddleware]] = None):

228 self._middlewares: list[AgentMiddleware] = list(middlewares or [])

229

230 def add(self, middleware: AgentMiddleware) -> MiddlewarePipeline:

231 """添加中间件，返回自身以支持链式调用。"""

232 self._middlewares.append(middleware)

233 return self

234

235 def remove(self, name: str) -> None:

236 self._middlewares = [m for m in self._middlewares if m.name != name]

237

238 @property

239 def middleware_names(self) -> list[str]:

240 return [m.name for m in self._middlewares]

241

242 async def execute_phase(

243 self,

244 phase: MiddlewarePhase,

245 ctx: MiddlewareContext,

246 ) -> MiddlewareDecision:

247 """执行指定阶段的所有中间件。"""

248 current_ctx = ctx

249 for mw in self._middlewares:

250 if phase not in mw.phases:

251 continue

252 decision = await mw.process(current_ctx)

253 if not decision.allow:

254 decision.blocked_by = mw.name

255 return decision

256 if decision.modified_context:

257 current_ctx = decision.modified_context

258 return MiddlewareDecision(allow=True, modified_context=current_ctx)

259

260 async def on_start(self, ctx: MiddlewareContext) -> MiddlewareDecision:

261 return await self.execute_phase(MiddlewarePhase.ON_START, ctx)

262

263 async def pre_llm(self, ctx: MiddlewareContext) -> MiddlewareDecision:

264 return await self.execute_phase(MiddlewarePhase.PRE_LLM, ctx)

265

266 async def post_llm(self, ctx: MiddlewareContext) -> MiddlewareDecision:

267 return await self.execute_phase(MiddlewarePhase.POST_LLM, ctx)

268

269 async def pre_tool(self, ctx: MiddlewareContext) -> MiddlewareDecision:

270 return await self.execute_phase(MiddlewarePhase.PRE_TOOL, ctx)

271

272 async def post_tool(self, ctx: MiddlewareContext) -> MiddlewareDecision:

273 return await self.execute_phase(MiddlewarePhase.POST_TOOL, ctx)

274

275 async def on_error(self, ctx: MiddlewareContext) -> MiddlewareDecision:

276 return await self.execute_phase(MiddlewarePhase.ON_ERROR, ctx)

277

278 async def on_complete(self, ctx: MiddlewareContext) -> MiddlewareDecision:

279 return await self.execute_phase(MiddlewarePhase.ON_COMPLETE, ctx)