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

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 dataclasses import dataclass, field 

13from enum import Enum 

14from typing import Any, Callable, Awaitable 

15 

16from .parent_child import ( 

17 ChildStatus, 

18 SharedState, 

19 ChildContext, 

20 ChildHandle, 

21) 

22 

23 

24class SubAgentMode(str, Enum): 

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

26 FORK = "fork" 

27 SWARM = "swarm" 

28 A2A = "a2a" 

29 

30 

31@dataclass 

32class SubAgentSpec: 

33 """子 Agent 规格。""" 

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

35 task: str = "" 

36 mode: SubAgentMode = SubAgentMode.FORK 

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

38 max_iterations: int = 50 

39 timeout: float | None = None 

40 heartbeat_interval: float = 2.0 

41 

42 

43@dataclass 

44class SubAgentResult: 

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

46 agent_id: str 

47 output: str 

48 iterations: int 

49 error: str | None = None 

50 handle: ChildHandle | None = None 

51 

52 def summarize(self) -> str: 

53 if self.error: 

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

55 return ( 

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

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

58 ) 

59 

60 

61class SubAgentManager: 

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

63 

64 用法:: 

65 

66 mgr = SubAgentManager() 

67 

68 # Fork 模式 

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

70 

71 # Swarm 模式 

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

73 

74 # 管控子Agent 

75 handle = mgr.get_handle(result.agent_id) 

76 await handle.pause() 

77 await handle.resume() 

78 await handle.cancel() 

79 status = handle.get_status() 

80 """ 

81 

82 MAX_SWARM_SIZE = 8 

83 

84 def __init__(self): 

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

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

87 

88 @property 

89 def shared_state(self) -> SharedState: 

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

91 return self._shared_state 

92 

93 @property 

94 def active_children(self) -> int: 

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

96 return sum( 

97 1 for h in self._agents.values() 

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

99 ) 

100 

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

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

103 return self._agents.get(agent_id) 

104 

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

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

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

108 

109 async def cancel_all(self) -> None: 

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

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

112 await asyncio.gather(*tasks) 

113 

114 async def spawn_fork( 

115 self, 

116 task: str, 

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

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

119 timeout: float | None = None, 

120 heartbeat_interval: float = 2.0, 

121 ) -> SubAgentResult: 

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

123 

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

125 """ 

126 spec = SubAgentSpec( 

127 task=task, 

128 mode=SubAgentMode.FORK, 

129 model=model, 

130 timeout=timeout, 

131 heartbeat_interval=heartbeat_interval, 

132 ) 

133 

134 handle = ChildHandle( 

135 agent_id=spec.id, 

136 task=task, 

137 mode=spec.mode.value, 

138 timeout=timeout, 

139 heartbeat_interval=heartbeat_interval, 

140 ) 

141 self._agents[spec.id] = handle 

142 ctx = handle.create_context() 

143 handle.info.status = ChildStatus.RUNNING 

144 

145 if run_func: 

146 try: 

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

148 if handle._cancel_flag: 

149 handle.info.status = ChildStatus.CANCELLED 

150 return SubAgentResult( 

151 agent_id=spec.id, 

152 output=output, 

153 iterations=iterations, 

154 error="Cancelled by parent", 

155 handle=handle, 

156 ) 

157 await ctx.done(output) 

158 handle.info.output = output 

159 handle.info.iterations = iterations 

160 handle.info.status = ChildStatus.COMPLETED 

161 return SubAgentResult( 

162 agent_id=spec.id, 

163 output=output, 

164 iterations=iterations, 

165 handle=handle, 

166 ) 

167 except asyncio.CancelledError: 

168 handle.info.status = ChildStatus.CANCELLED 

169 return SubAgentResult( 

170 agent_id=spec.id, 

171 output="", 

172 iterations=handle.info.iterations, 

173 error="Cancelled by parent", 

174 handle=handle, 

175 ) 

176 except Exception as e: 

177 await ctx.fail(str(e)) 

178 handle.info.status = ChildStatus.FAILED 

179 handle.info.error = str(e) 

180 return SubAgentResult( 

181 agent_id=spec.id, 

182 output="", 

183 iterations=handle.info.iterations, 

184 error=str(e), 

185 handle=handle, 

186 ) 

187 

188 handle.info.status = ChildStatus.COMPLETED 

189 return SubAgentResult( 

190 agent_id=spec.id, 

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

192 iterations=0, 

193 handle=handle, 

194 ) 

195 

196 async def spawn_swarm( 

197 self, 

198 tasks: list[str], 

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

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

201 timeout: float | None = None, 

202 heartbeat_interval: float = 2.0, 

203 ) -> list[SubAgentResult]: 

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

205 agents = [] 

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

207 spec = SubAgentSpec( 

208 task=task, 

209 mode=SubAgentMode.SWARM, 

210 model=model, 

211 timeout=timeout, 

212 heartbeat_interval=heartbeat_interval, 

213 ) 

214 agents.append(spec) 

215 

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

217 handle = ChildHandle( 

218 agent_id=spec.id, 

219 task=spec.task, 

220 mode=spec.mode.value, 

221 timeout=timeout, 

222 heartbeat_interval=heartbeat_interval, 

223 ) 

224 self._agents[spec.id] = handle 

225 ctx = handle.create_context() 

226 handle.info.status = ChildStatus.RUNNING 

227 

228 if run_func: 

229 try: 

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

231 if handle._cancel_flag: 

232 handle.info.status = ChildStatus.CANCELLED 

233 return SubAgentResult( 

234 agent_id=spec.id, 

235 output=output, 

236 iterations=iterations, 

237 error="Cancelled by parent", 

238 handle=handle, 

239 ) 

240 await ctx.done(output) 

241 handle.info.output = output 

242 handle.info.iterations = iterations 

243 handle.info.status = ChildStatus.COMPLETED 

244 return SubAgentResult( 

245 agent_id=spec.id, 

246 output=output, 

247 iterations=iterations, 

248 handle=handle, 

249 ) 

250 except Exception as e: 

251 await ctx.fail(str(e)) 

252 handle.info.status = ChildStatus.FAILED 

253 handle.info.error = str(e) 

254 return SubAgentResult( 

255 agent_id=spec.id, 

256 output="", 

257 iterations=handle.info.iterations, 

258 error=str(e), 

259 handle=handle, 

260 ) 

261 

262 handle.info.status = ChildStatus.COMPLETED 

263 return SubAgentResult( 

264 agent_id=spec.id, 

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

266 iterations=0, 

267 handle=handle, 

268 ) 

269 

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

271 

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

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

274 if "\n" in task: 

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

276 return [task] 

277 

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

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

280 

281 用法:: 

282 

283 asyncio.create_task(mgr.monitor_heartbeats()) 

284 """ 

285 while True: 

286 await asyncio.sleep(interval) 

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

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

289 if running and handle.check_timeout(): 

290 await handle.cancel() 

291 handle.info.status = ChildStatus.TIMEOUT 

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

293 elif running and handle.check_heartbeat_timeout(): 

294 handle.info.status = ChildStatus.FAILED 

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

296 

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

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

299 now = time.time() 

300 cleaned = 0 

301 terminal = (ChildStatus.COMPLETED, ChildStatus.FAILED, 

302 ChildStatus.CANCELLED, ChildStatus.TIMEOUT) 

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

304 if handle.status in terminal: 

305 age = now - handle.info.spawned_at 

306 if age > max_age_seconds: 

307 del self._agents[agent_id] 

308 cleaned += 1 

309 return cleaned