Coverage for agentos/tests/test_subagent_parent_child.py: 0%

1"""测试 SubAgent 父子通信 — 状态共享、心跳、生命周期管理。""" 

2 

3import asyncio 

4import pytest 

5 

6pytestmark = pytest.mark.asyncio 

7from agentos.subagent import ( 

8 SubAgentManager, 

9 SubAgentSpec, 

10 ChildStatus, 

11 ChildHeartbeat, 

12 SharedState, 

13 ChildContext, 

14 ChildHandle, 

15) 

16 

17 

18class TestSharedState: 

19 async def test_set_get(self): 

20 ss = SharedState() 

21 await ss.set("key1", "val1") 

22 assert await ss.get("key1") == "val1" 

23 assert await ss.get("missing", "def") == "def" 

24 

25 async def test_update_snapshot(self): 

26 ss = SharedState() 

27 await ss.update({"a": 1, "b": 2}) 

28 snap = await ss.snapshot() 

29 assert snap == {"a": 1, "b": 2} 

30 

31 async def test_sync_ops(self): 

32 ss = SharedState() 

33 ss.set_sync("x", 42) 

34 assert ss.get_sync("x") == 42 

35 

36 async def test_concurrent_writes(self): 

37 ss = SharedState() 

38 

39 async def writer(key: str, n: int): 

40 for i in range(n): 

41 await ss.set(key, i) 

42 await asyncio.sleep(0) 

43 

44 await asyncio.gather(writer("a", 20), writer("b", 20)) 

45 assert await ss.get("a") == 19 

46 assert await ss.get("b") == 19 

47 

48 

49class TestChildContext: 

50 async def test_progress_report(self): 

51 hbs = [] 

52 

53 async def hb_cb(hb: ChildHeartbeat): 

54 hbs.append(hb) 

55 

56 ctx = ChildContext("test-1", heartbeat_callback=hb_cb) 

57 await ctx.report_progress(0.5, "step1", "half done") 

58 assert ctx.progress == 0.5 

59 assert len(hbs) == 1 

60 assert hbs[0].progress == 0.5 

61 assert hbs[0].current_step == "step1" 

62 

63 async def test_step_and_heartbeat(self): 

64 hbs = [] 

65 

66 async def hb_cb(hb: ChildHeartbeat): 

67 hbs.append(hb) 

68 

69 ctx = ChildContext("test-2", heartbeat_callback=hb_cb) 

70 await ctx.step(1, "init") 

71 await ctx.send_heartbeat("alive") 

72 assert len(hbs) == 1 

73 assert hbs[0].iteration == 1 

74 

75 async def test_done(self): 

76 hbs = [] 

77 

78 async def hb_cb(hb: ChildHeartbeat): 

79 hbs.append(hb) 

80 

81 ctx = ChildContext("test-3", heartbeat_callback=hb_cb) 

82 await ctx.done("all good") 

83 assert len(hbs) == 1 

84 assert hbs[0].status == ChildStatus.COMPLETED 

85 assert hbs[0].progress == 1.0 

86 assert hbs[0].message == "all good" 

87 

88 async def test_fail(self): 

89 hbs = [] 

90 

91 async def hb_cb(hb: ChildHeartbeat): 

92 hbs.append(hb) 

93 

94 ctx = ChildContext("test-4", heartbeat_callback=hb_cb) 

95 await ctx.fail("something broke") 

96 assert len(hbs) == 1 

97 assert hbs[0].status == ChildStatus.FAILED 

98 assert hbs[0].message == "something broke" 

99 

100 async def test_cancel_detection(self): 

101 cancelled = [False] 

102 

103 def on_cancel(): 

104 return cancelled[0] 

105 

106 ctx = ChildContext("test-5", on_cancel=on_cancel) 

107 assert not ctx.cancelled 

108 status = await ctx.check_control() 

109 assert status == ChildStatus.RUNNING 

110 

111 cancelled[0] = True 

112 status = await ctx.check_control() 

113 assert status == ChildStatus.CANCELLED 

114 assert ctx.cancelled 

115 

116 async def test_pause_resume(self): 

117 paused = [True] 

118 resume_triggered = [False] 

119 

120 async def on_pause(): 

121 resume_triggered[0] = True 

122 paused[0] = False 

123 

124 ctx = ChildContext("test-6", on_pause=on_pause) 

125 ctx._paused = paused[0] 

126 status = await ctx.check_control() 

127 assert status == ChildStatus.PAUSED 

128 assert resume_triggered[0] 

129 

130 

131class TestChildHandle: 

132 async def test_create_context(self): 

133 handle = ChildHandle("h1", "do stuff", "fork") 

134 ctx = handle.create_context() 

135 assert ctx.agent_id == "h1" 

136 assert handle.context is ctx 

137 assert handle.shared_state is ctx.shared_state 

138 

139 async def test_pause_resume(self): 

140 handle = ChildHandle("h2", "task", "fork") 

141 handle.create_context() 

142 assert handle.status == ChildStatus.IDLE 

143 

144 await handle.pause() 

145 assert handle.status == ChildStatus.PAUSED 

146 

147 await handle.resume() 

148 assert handle.status == ChildStatus.RUNNING 

149 

150 async def test_cancel(self): 

151 handle = ChildHandle("h3", "task", "fork") 

152 handle.create_context() 

153 await handle.cancel() 

154 assert handle.status == ChildStatus.CANCELLED 

155 

156 async def test_get_status(self): 

157 handle = ChildHandle("h4", "analyze", "fork") 

158 handle.create_context() 

159 handle.info.progress = 0.7 

160 handle.info.current_step = "parsing" 

161 handle.info.iterations = 12 

162 

163 status = handle.get_status() 

164 assert status["agent_id"] == "h4" 

165 assert status["progress"] == 0.7 

166 assert status["current_step"] == "parsing" 

167 assert status["iterations"] == 12 

168 assert "elapsed" in status 

169 

170 async def test_timeout_detection(self): 

171 handle = ChildHandle("h5", "task", "fork", timeout=0.1) 

172 await asyncio.sleep(0.15) 

173 assert handle.check_timeout() 

174 

175 async def test_no_timeout_when_unset(self): 

176 handle = ChildHandle("h6", "task", "fork", timeout=None) 

177 assert not handle.check_timeout() 

178 

179 async def test_heartbeat_timeout(self): 

180 handle = ChildHandle("h7", "task", "fork", heartbeat_interval=0.1) 

181 await asyncio.sleep(0.35) 

182 assert handle.check_heartbeat_timeout() 

183 

184 async def test_heartbeat_updates_info(self): 

185 handle = ChildHandle("h8", "task", "fork") 

186 handle.create_context() 

187 await handle._receive_heartbeat(ChildHeartbeat( 

188 agent_id="h8", progress=0.5, current_step="s1", 

189 message="working", iteration=5, 

190 )) 

191 assert handle.info.progress == 0.5 

192 assert handle.info.current_step == "s1" 

193 assert handle.info.iterations == 5 

194 

195 async def test_shared_state_parent_child(self): 

196 handle = ChildHandle("h9", "task", "fork") 

197 ctx = handle.create_context() 

198 

199 await ctx.shared_state.set("data", [1, 2, 3]) 

200 val = await handle.shared_state.get("data") 

201 assert val == [1, 2, 3] 

202 

203 await handle.shared_state.set("status", "ok") 

204 assert await ctx.shared_state.get("status") == "ok" 

205 

206 

207class TestSubAgentManager: 

208 async def test_spawn_fork_with_child_context(self): 

209 hbs = [] 

210 

211 async def run_func(spec: SubAgentSpec, ctx: ChildContext): 

212 await ctx.report_progress(0.3, "init") 

213 await ctx.step(1, "load") 

214 await ctx.report_progress(0.7, "process") 

215 await ctx.done("success") 

216 return ("success", 2) 

217 

218 mgr = SubAgentManager() 

219 result = await mgr.spawn_fork("test task", run_func=run_func) 

220 assert result.output == "success" 

221 assert result.iterations == 2 

222 assert result.handle is not None 

223 assert result.handle.status == ChildStatus.COMPLETED 

224 

225 async def test_spawn_fork_failure(self): 

226 async def run_func(spec, ctx): 

227 await ctx.report_progress(0.1, "start") 

228 raise ValueError("boom") 

229 

230 mgr = SubAgentManager() 

231 result = await mgr.spawn_fork("bad task", run_func=run_func) 

232 assert result.error == "boom" 

233 assert result.handle.status == ChildStatus.FAILED 

234 assert result.handle.info.error == "boom" 

235 

236 async def test_spawn_fork_pause_resume_flow(self): 

237 """模拟父子协作：父暂停→子暂停→父恢复→子继续→完成。""" 

238 state = {"phase": "init"} 

239 

240 async def run_func(spec, ctx: ChildContext): 

241 state["phase"] = "running" 

242 await ctx.report_progress(0.2, "step1") 

243 

244 # 检查控制信号 

245 status = await ctx.check_control() 

246 if status == ChildStatus.PAUSED: 

247 state["phase"] = "paused" 

248 

249 # 再次检查（模拟恢复后继续） 

250 status = await ctx.check_control() 

251 if status == ChildStatus.RUNNING: 

252 state["phase"] = "resumed" 

253 

254 await ctx.done("ok") 

255 return ("ok", 3) 

256 

257 mgr = SubAgentManager() 

258 

259 # 启动 

260 task = asyncio.create_task( 

261 mgr.spawn_fork("pause test", run_func=run_func) 

262 ) 

263 

264 await asyncio.sleep(0.05) # 让子Agent跑到 step 

265 handle = mgr.get_handle(task.result().handle.agent_id) if hasattr(task, 'result') else None 

266 

267 # 等task完成 

268 result = await task 

269 assert result.error is None or result.error == "" 

270 assert state["phase"] in ("running", "paused", "resumed") 

271 

272 async def test_swarm_parallel(self): 

273 results_log = [] 

274 

275 async def run_func(spec: SubAgentSpec, ctx: ChildContext): 

276 await ctx.report_progress(0.5, spec.task) 

277 await asyncio.sleep(0.01) 

278 await ctx.done(f"done_{spec.task}") 

279 results_log.append(spec.task) 

280 return (f"done_{spec.task}", 1) 

281 

282 mgr = SubAgentManager() 

283 results = await mgr.spawn_swarm( 

284 ["A", "B", "C"], run_func=run_func 

285 ) 

286 assert len(results) == 3 

287 assert len(results_log) == 3 

288 for r in results: 

289 assert r.handle is not None 

290 assert r.handle.status == ChildStatus.COMPLETED 

291 

292 async def test_cancel_all(self): 

293 async def run_func(spec, ctx: ChildContext): 

294 await ctx.report_progress(0.1, "init") 

295 for i in range(50): 

296 await ctx.step(i, f"step_{i}") 

297 status = await ctx.check_control() 

298 if status == ChildStatus.CANCELLED: 

299 return ("cancelled", i) 

300 await asyncio.sleep(0.01) 

301 return ("done", 50) 

302 

303 mgr = SubAgentManager() 

304 t1 = asyncio.create_task( 

305 mgr.spawn_fork("long task 1", run_func=run_func) 

306 ) 

307 t2 = asyncio.create_task( 

308 mgr.spawn_fork("long task 2", run_func=run_func) 

309 ) 

310 

311 await asyncio.sleep(0.05) 

312 await mgr.cancel_all() 

313 

314 r1, r2 = await asyncio.gather(t1, t2) 

315 assert r1.handle.status == ChildStatus.CANCELLED 

316 assert r2.handle.status == ChildStatus.CANCELLED 

317 

318 async def test_list_children(self): 

319 mgr = SubAgentManager() 

320 r = await mgr.spawn_fork("task1") 

321 children = mgr.list_children() 

322 assert len(children) == 1 

323 assert children[0]["agent_id"] == r.agent_id 

324 

325 async def test_cleanup(self): 

326 mgr = SubAgentManager() 

327 r = await mgr.spawn_fork("cleanup test") 

328 assert len(mgr._agents) == 1 

329 

330 cleaned = await mgr.cleanup(max_age_seconds=-1.0) 

331 assert cleaned == 1 

332 assert len(mgr._agents) == 0 

333 

334 async def test_heartbeat_monitoring(self): 

335 mgr = SubAgentManager() 

336 handle = ChildHandle("hb-test", "task", "fork", timeout=0.05) 

337 mgr._agents["hb-test"] = handle 

338 handle.info.status = ChildStatus.RUNNING 

339 

340 monitor = asyncio.create_task(mgr.monitor_heartbeats(interval=0.02)) 

341 await asyncio.sleep(0.1) 

342 monitor.cancel() 

343 try: 

344 await monitor 

345 except asyncio.CancelledError: 

346 pass 

347 

348 assert handle.status in (ChildStatus.TIMEOUT, ChildStatus.RUNNING)