Coverage for agentos/subagent/manager.py: 34%

1""" 

2子Agent管理 — Fork隔离 + Swarm并行 + A2A委派 + 父子通信。 

3基因来源: Claude Code (Fork) + Cursor (Swarm) 

4v1.3.15: +Parent-Child 通信（状态共享、心跳、生命周期） 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10import time 

11import uuid 

12from collections.abc import Awaitable, Callable 

13from dataclasses import dataclass, field 

14from enum import StrEnum 

15from typing import Any 

16 

17from .parent_child import ( 

18 ChildContext, 

19 ChildHandle, 

20 ChildStatus, 

21 SharedState, 

22) 

23 

24 

25class SubAgentMode(StrEnum): 

26 """子 Agent 模式枚举。""" 

27 

28 FORK = "fork" 

29 SWARM = "swarm" 

30 A2A = "a2a" 

31 

32 

33@dataclass 

34class SubAgentSpec: 

35 """子 Agent 规格。""" 

36 

37 id: str = field(default_factory=lambda: uuid.uuid4().hex[:8]) 

38 task: str = "" 

39 mode: SubAgentMode = SubAgentMode.FORK 

40 model: str = "kimi-k2.6" 

41 max_iterations: int = 50 

42 timeout: float | None = None 

43 heartbeat_interval: float = 2.0 

44 

45 

46@dataclass 

47class SubAgentResult: 

48 """子 Agent 执行结果。""" 

49 

50 agent_id: str 

51 output: str 

52 iterations: int 

53 error: str | None = None 

54 handle: ChildHandle | None = None 

55 

56 def summarize(self) -> str: 

57 if self.error: 

58 return f"[SubAgent {self.agent_id}] FAILED: {self.error}" 

59 return ( 

60 f"[SubAgent {self.agent_id}] Completed in {self.iterations} steps.\n" 

61 f"Result: {self.output[:500]}" 

62 ) 

63 

64 

65class SubAgentManager: 

66 """子Agent管理器 — Fork/Swarm/A2A + 父子通信。 

67 

68 用法:: 

69 

70 mgr = SubAgentManager() 

71 

72 # Fork 模式 

73 result = await mgr.spawn_fork("分析这份报告") 

74 

75 # Swarm 模式 

76 results = await mgr.spawn_swarm(["任务A", "任务B"]) 

77 

78 # 管控子Agent 

79 handle = mgr.get_handle(result.agent_id) 

80 await handle.pause() 

81 await handle.resume() 

82 await handle.cancel() 

83 status = handle.get_status() 

84 """ 

85 

86 MAX_SWARM_SIZE = 8 

87 

88 def __init__(self): 

89 self._agents: dict[str, ChildHandle] = {} 

90 self._shared_state = SharedState() # 全局共享状态 

91 

92 @property 

93 def shared_state(self) -> SharedState: 

94 """全局父子共享状态。""" 

95 return self._shared_state 

96 

97 @property 

98 def active_children(self) -> int: 

99 """当前活跃的子Agent数。""" 

100 return sum( 

101 1 

102 for h in self._agents.values() 

103 if h.status in (ChildStatus.RUNNING, ChildStatus.PAUSED) 

104 ) 

105 

106 def get_handle(self, agent_id: str) -> ChildHandle | None: 

107 """根据 agent_id 获取子Agent句柄。""" 

108 return self._agents.get(agent_id) 

109 

110 def list_children(self) -> list[dict[str, Any]]: 

111 """列出所有子Agent状态。""" 

112 return [h.get_status() for h in self._agents.values()] 

113 

114 async def cancel_all(self) -> None: 

115 """取消所有子Agent。""" 

116 tasks = [h.cancel() for h in self._agents.values()] 

117 await asyncio.gather(*tasks) 

118 

119 async def spawn_fork( 

120 self, 

121 task: str, 

122 model: str = "kimi-k2.6", 

123 run_func: Callable[[SubAgentSpec, ChildContext], Awaitable[tuple[str, int]]] | None = None, 

124 timeout: float | None = None, 

125 heartbeat_interval: float = 2.0, 

126 ) -> SubAgentResult: 

127 """Fork模式：子Agent在干净上下文中运行，父只拿摘要。 

128 

129 run_func(spec, ctx) -> (output, iterations) 

130 """ 

131 spec = SubAgentSpec( 

132 task=task, 

133 mode=SubAgentMode.FORK, 

134 model=model, 

135 timeout=timeout, 

136 heartbeat_interval=heartbeat_interval, 

137 ) 

138 

139 handle = ChildHandle( 

140 agent_id=spec.id, 

141 task=task, 

142 mode=spec.mode.value, 

143 timeout=timeout, 

144 heartbeat_interval=heartbeat_interval, 

145 ) 

146 self._agents[spec.id] = handle 

147 ctx = handle.create_context() 

148 handle.info.status = ChildStatus.RUNNING 

149 

150 if run_func: 

151 try: 

152 output, iterations = await run_func(spec, ctx) 

153 if handle._cancel_flag: 

154 handle.info.status = ChildStatus.CANCELLED 

155 return SubAgentResult( 

156 agent_id=spec.id, 

157 output=output, 

158 iterations=iterations, 

159 error="Cancelled by parent", 

160 handle=handle, 

161 ) 

162 await ctx.done(output) 

163 handle.info.output = output 

164 handle.info.iterations = iterations 

165 handle.info.status = ChildStatus.COMPLETED 

166 return SubAgentResult( 

167 agent_id=spec.id, 

168 output=output, 

169 iterations=iterations, 

170 handle=handle, 

171 ) 

172 except asyncio.CancelledError: 

173 handle.info.status = ChildStatus.CANCELLED 

174 return SubAgentResult( 

175 agent_id=spec.id, 

176 output="", 

177 iterations=handle.info.iterations, 

178 error="Cancelled by parent", 

179 handle=handle, 

180 ) 

181 except Exception as e: 

182 await ctx.fail(str(e)) 

183 handle.info.status = ChildStatus.FAILED 

184 handle.info.error = str(e) 

185 return SubAgentResult( 

186 agent_id=spec.id, 

187 output="", 

188 iterations=handle.info.iterations, 

189 error=str(e), 

190 handle=handle, 

191 ) 

192 

193 handle.info.status = ChildStatus.COMPLETED 

194 return SubAgentResult( 

195 agent_id=spec.id, 

196 output=f"Fork agent would process: {task}", 

197 iterations=0, 

198 handle=handle, 

199 ) 

200 

201 async def spawn_swarm( 

202 self, 

203 tasks: list[str], 

204 model: str = "kimi-k2.6", 

205 run_func: Callable[[SubAgentSpec, ChildContext], Awaitable[tuple[str, int]]] | None = None, 

206 timeout: float | None = None, 

207 heartbeat_interval: float = 2.0, 

208 ) -> list[SubAgentResult]: 

209 """Swarm模式：最多8个Agent并行处理。""" 

210 agents = [] 

211 for i, task in enumerate(tasks[: self.MAX_SWARM_SIZE]): 

212 spec = SubAgentSpec( 

213 task=task, 

214 mode=SubAgentMode.SWARM, 

215 model=model, 

216 timeout=timeout, 

217 heartbeat_interval=heartbeat_interval, 

218 ) 

219 agents.append(spec) 

220 

221 async def run_one(spec: SubAgentSpec) -> SubAgentResult: 

222 handle = ChildHandle( 

223 agent_id=spec.id, 

224 task=spec.task, 

225 mode=spec.mode.value, 

226 timeout=timeout, 

227 heartbeat_interval=heartbeat_interval, 

228 ) 

229 self._agents[spec.id] = handle 

230 ctx = handle.create_context() 

231 handle.info.status = ChildStatus.RUNNING 

232 

233 if run_func: 

234 try: 

235 output, iterations = await run_func(spec, ctx) 

236 if handle._cancel_flag: 

237 handle.info.status = ChildStatus.CANCELLED 

238 return SubAgentResult( 

239 agent_id=spec.id, 

240 output=output, 

241 iterations=iterations, 

242 error="Cancelled by parent", 

243 handle=handle, 

244 ) 

245 await ctx.done(output) 

246 handle.info.output = output 

247 handle.info.iterations = iterations 

248 handle.info.status = ChildStatus.COMPLETED 

249 return SubAgentResult( 

250 agent_id=spec.id, 

251 output=output, 

252 iterations=iterations, 

253 handle=handle, 

254 ) 

255 except Exception as e: 

256 await ctx.fail(str(e)) 

257 handle.info.status = ChildStatus.FAILED 

258 handle.info.error = str(e) 

259 return SubAgentResult( 

260 agent_id=spec.id, 

261 output="", 

262 iterations=handle.info.iterations, 

263 error=str(e), 

264 handle=handle, 

265 ) 

266 

267 handle.info.status = ChildStatus.COMPLETED 

268 return SubAgentResult( 

269 agent_id=spec.id, 

270 output=f"Swarm agent would process: {spec.task}", 

271 iterations=0, 

272 handle=handle, 

273 ) 

274 

275 return await asyncio.gather(*[run_one(a) for a in agents]) 

276 

277 def split_task(self, task: str) -> list[str]: 

278 """将复杂任务拆分为子任务。""" 

279 if "\n" in task: 

280 return [t.strip() for t in task.split("\n") if t.strip()] 

281 return [task] 

282 

283 async def monitor_heartbeats(self, interval: float = 1.0) -> None: 

284 """后台心跳监控协程，检测超时和失联子Agent。 

285 

286 用法:: 

287 

288 asyncio.create_task(mgr.monitor_heartbeats()) 

289 """ 

290 while True: 

291 await asyncio.sleep(interval) 

292 for agent_id, handle in list(self._agents.items()): 

293 running = handle.status in (ChildStatus.RUNNING, ChildStatus.PAUSED) 

294 if running and handle.check_timeout(): 

295 await handle.cancel() 

296 handle.info.status = ChildStatus.TIMEOUT 

297 handle.info.error = f"Timeout after {handle.info.timeout}s" 

298 elif running and handle.check_heartbeat_timeout(): 

299 handle.info.status = ChildStatus.FAILED 

300 handle.info.error = "Heartbeat lost — child agent unresponsive" 

301 

302 async def cleanup(self, max_age_seconds: float = 3600.0) -> int: 

303 """清理已完成/失败/取消且超过 max_age_seconds 的句柄。返回清理数。""" 

304 now = time.time() 

305 cleaned = 0 

306 terminal = ( 

307 ChildStatus.COMPLETED, 

308 ChildStatus.FAILED, 

309 ChildStatus.CANCELLED, 

310 ChildStatus.TIMEOUT, 

311 ) 

312 for agent_id, handle in list(self._agents.items()): 

313 if handle.status in terminal: 

314 age = now - handle.info.spawned_at 

315 if age > max_age_seconds: 

316 del self._agents[agent_id] 

317 cleaned += 1 

318 return cleaned