Coverage for agentos/subagent/parent_child.py: 47%

1""" 

2子Agent父子通信 — 状态共享、心跳、生命周期管理。 

3父Agent通过 ChildHandle 管控子Agent；子Agent通过 ChildContext 向父Agent报告。 

4""" 

5 

6from __future__ import annotations 

7 

8import asyncio 

9import time 

10from collections.abc import Awaitable, Callable 

11from dataclasses import dataclass, field 

12from enum import StrEnum 

13from typing import Any 

14 

15 

16class ChildStatus(StrEnum): 

17 """子Agent运行状态。""" 

18 

19 IDLE = "idle" 

20 RUNNING = "running" 

21 PAUSED = "paused" 

22 COMPLETED = "completed" 

23 FAILED = "failed" 

24 CANCELLED = "cancelled" 

25 TIMEOUT = "timeout" 

26 

27 

28@dataclass 

29class ChildHeartbeat: 

30 """子Agent心跳包。""" 

31 

32 agent_id: str 

33 status: ChildStatus = ChildStatus.RUNNING 

34 progress: float = 0.0 # 0.0 ~ 1.0 

35 current_step: str = "" 

36 message: str = "" 

37 iteration: int = 0 

38 timestamp: float = field(default_factory=time.time) 

39 

40 

41@dataclass 

42class ChildInfo: 

43 """子Agent元信息（父Agent侧）。""" 

44 

45 agent_id: str 

46 task: str 

47 mode: str 

48 status: ChildStatus = ChildStatus.IDLE 

49 spawned_at: float = field(default_factory=time.time) 

50 last_heartbeat: float = field(default_factory=time.time) 

51 heartbeat_interval: float = 2.0 # 期望心跳间隔（秒） 

52 timeout: float | None = None # 超时（秒），None=无超时 

53 progress: float = 0.0 

54 current_step: str = "" 

55 iterations: int = 0 

56 error: str | None = None 

57 output: str = "" 

58 

59 

60class SharedState: 

61 """父子共享状态（线程安全）。""" 

62 

63 def __init__(self): 

64 self._lock = asyncio.Lock() 

65 self._data: dict[str, Any] = {} 

66 

67 async def set(self, key: str, value: Any) -> None: 

68 async with self._lock: 

69 self._data[key] = value 

70 

71 async def get(self, key: str, default: Any = None) -> Any: 

72 async with self._lock: 

73 return self._data.get(key, default) 

74 

75 async def update(self, mapping: dict[str, Any]) -> None: 

76 async with self._lock: 

77 self._data.update(mapping) 

78 

79 async def snapshot(self) -> dict[str, Any]: 

80 async with self._lock: 

81 return dict(self._data) 

82 

83 def set_sync(self, key: str, value: Any) -> None: 

84 """同步写（非协程场景）。""" 

85 self._data[key] = value 

86 

87 def get_sync(self, key: str, default: Any = None) -> Any: 

88 """同步读（非协程场景）。""" 

89 return self._data.get(key, default) 

90 

91 

92class ChildContext: 

93 """子Agent视角 — 向父Agent报告状态、检查控制信号。""" 

94 

95 def __init__( 

96 self, 

97 agent_id: str, 

98 heartbeat_callback: Callable[[ChildHeartbeat], Awaitable[None]] | None = None, 

99 on_cancel: Callable[[], bool] | None = None, 

100 on_pause: Callable[[], Awaitable[None]] | None = None, 

101 shared_state: SharedState | None = None, 

102 ): 

103 self.agent_id = agent_id 

104 self._heartbeat_cb = heartbeat_callback 

105 self._cancel_check = on_cancel or (lambda: False) 

106 self._pause_cb = on_pause or (lambda: asyncio.sleep(0)) 

107 self.shared_state = shared_state or SharedState() 

108 self._cancelled = False 

109 self._paused = False 

110 self._progress = 0.0 

111 self._current_step = "" 

112 self._iteration = 0 

113 

114 @property 

115 def cancelled(self) -> bool: 

116 return self._cancelled 

117 

118 @property 

119 def paused(self) -> bool: 

120 return self._paused 

121 

122 @property 

123 def progress(self) -> float: 

124 return self._progress 

125 

126 async def report_progress( 

127 self, 

128 progress: float, 

129 step: str = "", 

130 message: str = "", 

131 ) -> None: 

132 """子Agent报告进度。""" 

133 self._progress = max(0.0, min(1.0, progress)) 

134 self._current_step = step 

135 if self._heartbeat_cb: 

136 await self._heartbeat_cb( 

137 ChildHeartbeat( 

138 agent_id=self.agent_id, 

139 status=ChildStatus.RUNNING, 

140 progress=self._progress, 

141 current_step=step, 

142 message=message, 

143 iteration=self._iteration, 

144 ) 

145 ) 

146 

147 async def step(self, iteration: int, step: str = "") -> None: 

148 """子Agent标记一个执行步。""" 

149 self._iteration = iteration 

150 self._current_step = step 

151 

152 async def check_control(self) -> ChildStatus: 

153 """检查父Agent控制信号，返回应执行的操作。""" 

154 if self._cancel_check(): 

155 self._cancelled = True 

156 return ChildStatus.CANCELLED 

157 if self._paused: 

158 await self._pause_cb() 

159 return ChildStatus.PAUSED 

160 return ChildStatus.RUNNING 

161 

162 async def send_heartbeat(self, message: str = "") -> None: 

163 """子Agent发送心跳。""" 

164 if self._heartbeat_cb: 

165 await self._heartbeat_cb( 

166 ChildHeartbeat( 

167 agent_id=self.agent_id, 

168 status=ChildStatus.RUNNING, 

169 progress=self._progress, 

170 current_step=self._current_step, 

171 message=message, 

172 iteration=self._iteration, 

173 ) 

174 ) 

175 

176 async def done(self, output: str = "") -> None: 

177 """子Agent标记完成。""" 

178 if self._heartbeat_cb: 

179 await self._heartbeat_cb( 

180 ChildHeartbeat( 

181 agent_id=self.agent_id, 

182 status=ChildStatus.COMPLETED, 

183 progress=1.0, 

184 current_step=self._current_step, 

185 message=output, 

186 iteration=self._iteration, 

187 ) 

188 ) 

189 

190 async def fail(self, error: str) -> None: 

191 """子Agent报告失败。""" 

192 if self._heartbeat_cb: 

193 await self._heartbeat_cb( 

194 ChildHeartbeat( 

195 agent_id=self.agent_id, 

196 status=ChildStatus.FAILED, 

197 progress=self._progress, 

198 current_step=self._current_step, 

199 message=error, 

200 iteration=self._iteration, 

201 ) 

202 ) 

203 

204 

205class ChildHandle: 

206 """父Agent视角 — 管控一个子Agent。""" 

207 

208 def __init__( 

209 self, 

210 agent_id: str, 

211 task: str, 

212 mode: str, 

213 timeout: float | None = None, 

214 heartbeat_interval: float = 2.0, 

215 ): 

216 self.info = ChildInfo( 

217 agent_id=agent_id, 

218 task=task, 

219 mode=mode, 

220 heartbeat_interval=heartbeat_interval, 

221 timeout=timeout, 

222 ) 

223 self._cancel_flag = False 

224 self._pause_flag = False 

225 self._resume_event = asyncio.Event() 

226 self._resume_event.set() # 默认未暂停 

227 self.shared_state = SharedState() 

228 self.context: ChildContext | None = None 

229 

230 @property 

231 def agent_id(self) -> str: 

232 return self.info.agent_id 

233 

234 @property 

235 def status(self) -> ChildStatus: 

236 return self.info.status 

237 

238 def create_context(self) -> ChildContext: 

239 """为子Agent创建 ChildContext。""" 

240 ctx = ChildContext( 

241 agent_id=self.agent_id, 

242 heartbeat_callback=self._receive_heartbeat, 

243 on_cancel=self._is_cancelled, 

244 on_pause=self._wait_if_paused, 

245 shared_state=self.shared_state, 

246 ) 

247 self.context = ctx 

248 return ctx 

249 

250 async def _receive_heartbeat(self, hb: ChildHeartbeat) -> None: 

251 """接收子Agent心跳。""" 

252 self.info.last_heartbeat = time.time() 

253 self.info.status = hb.status 

254 self.info.progress = hb.progress 

255 self.info.current_step = hb.current_step 

256 self.info.iterations = hb.iteration 

257 if hb.status == ChildStatus.FAILED: 

258 self.info.error = hb.message 

259 elif hb.status == ChildStatus.COMPLETED: 

260 self.info.output = hb.message 

261 

262 def _is_cancelled(self) -> bool: 

263 return self._cancel_flag 

264 

265 async def _wait_if_paused(self) -> None: 

266 await self._resume_event.wait() 

267 

268 async def cancel(self) -> None: 

269 """取消子Agent。""" 

270 self._cancel_flag = True 

271 self.info.status = ChildStatus.CANCELLED 

272 

273 async def pause(self) -> None: 

274 """暂停子Agent。""" 

275 self._pause_flag = True 

276 self._resume_event.clear() 

277 self.info.status = ChildStatus.PAUSED 

278 if self.context: 

279 self.context._paused = True 

280 

281 async def resume(self) -> None: 

282 """恢复子Agent。""" 

283 self._pause_flag = False 

284 self._resume_event.set() 

285 self.info.status = ChildStatus.RUNNING 

286 if self.context: 

287 self.context._paused = False 

288 

289 def check_timeout(self) -> bool: 

290 """检查是否超时，返回 True 表示已超时。""" 

291 if self.info.timeout is None: 

292 return False 

293 elapsed = time.time() - self.info.spawned_at 

294 return elapsed > self.info.timeout 

295 

296 def check_heartbeat_timeout(self) -> bool: 

297 """检查心跳是否超时（3倍心跳间隔无响应视为失联）。""" 

298 elapsed = time.time() - self.info.last_heartbeat 

299 return elapsed > self.info.heartbeat_interval * 3 

300 

301 def get_status(self) -> dict[str, Any]: 

302 """获取子Agent状态摘要。""" 

303 return { 

304 "agent_id": self.info.agent_id, 

305 "status": self.info.status.value, 

306 "progress": self.info.progress, 

307 "current_step": self.info.current_step, 

308 "iterations": self.info.iterations, 

309 "elapsed": time.time() - self.info.spawned_at, 

310 "error": self.info.error, 

311 }