Coverage for session_buddy / core / lifecycle / service_registry.py: 25.97%

1from __future__ import annotations 

2 

3import inspect 

4import typing as t 

5from contextlib import suppress 

6from dataclasses import dataclass, field 

7 

8if t.TYPE_CHECKING: 

9 from session_buddy.di import get_sync_typed 

10 from session_buddy.di.config import SessionPaths 

11 

12ServiceHook = t.Callable[[], t.Awaitable[t.Any] | t.Any] 

13 

14 

15@dataclass(frozen=True, slots=True) 

16class ServiceSpec: 

17 name: str 

18 category: str 

19 init: ServiceHook | None = None 

20 health: ServiceHook | None = None 

21 cleanup: ServiceHook | None = None 

22 

23 

24class ServiceRegistry: 

25 def __init__(self) -> None: 

26 self._services: list[ServiceSpec] = [] 

27 

28 def register(self, service: ServiceSpec) -> None: 

29 self._services.append(service) 

30 

31 async def init_all(self) -> None: 

32 for service in self._services: 

33 await _maybe_call(service.init) 

34 

35 async def health_all(self) -> dict[str, t.Any]: 

36 results: dict[str, t.Any] = {} 

37 for service in self._services: 

38 results[service.name] = await _maybe_call(service.health) 

39 return results 

40 

41 async def cleanup_all(self) -> None: 

42 for service in self._services: 

43 await _maybe_call(service.cleanup) 

44 

45 

46_registry: ServiceRegistry | None = None 

47 

48 

49def get_service_registry() -> ServiceRegistry: 

50 global _registry 

51 if _registry is None: 

52 _registry = ServiceRegistry() 

53 _register_defaults(_registry) 

54 return _registry 

55 

56 

57def _register_defaults(registry: ServiceRegistry) -> None: 

58 registry.register( 

59 ServiceSpec( 

60 name="core.di_config", 

61 category="core", 

62 init=_init_di_config, 

63 health=_health_di_config, 

64 cleanup=_noop, 

65 ) 

66 ) 

67 registry.register( 

68 ServiceSpec( 

69 name="core.permissions_manager", 

70 category="core", 

71 init=_init_permissions_manager, 

72 health=_health_permissions_manager, 

73 cleanup=_noop, 

74 ) 

75 ) 

76 registry.register( 

77 ServiceSpec( 

78 name="core.lifecycle_manager", 

79 category="core", 

80 init=_init_lifecycle_manager, 

81 health=_health_lifecycle_manager, 

82 cleanup=_noop, 

83 ) 

84 ) 

85 registry.register( 

86 ServiceSpec( 

87 name="memory.reflection_db", 

88 category="memory", 

89 init=_init_reflection_db, 

90 health=_health_reflection_db, 

91 cleanup=_cleanup_reflection_db, 

92 ) 

93 ) 

94 registry.register( 

95 ServiceSpec( 

96 name="memory.knowledge_graph", 

97 category="memory", 

98 init=_init_knowledge_graph, 

99 health=_health_knowledge_graph, 

100 cleanup=_cleanup_knowledge_graph, 

101 ) 

102 ) 

103 registry.register( 

104 ServiceSpec( 

105 name="adapters.storage", 

106 category="adapters", 

107 init=_init_storage_adapters, 

108 health=_health_storage_adapters, 

109 cleanup=_cleanup_storage_adapters, 

110 ) 

111 ) 

112 registry.register( 

113 ServiceSpec( 

114 name="adapters.caches", 

115 category="adapters", 

116 init=_init_cache_adapters, 

117 health=_health_cache_adapters, 

118 cleanup=_cleanup_cache_adapters, 

119 ) 

120 ) 

121 registry.register( 

122 ServiceSpec( 

123 name="tools.registry", 

124 category="tools", 

125 init=_noop, 

126 health=_health_tools_registry, 

127 cleanup=_noop, 

128 ) 

129 ) 

130 registry.register( 

131 ServiceSpec( 

132 name="utils.logging", 

133 category="utils", 

134 init=_init_logging, 

135 health=_health_logging, 

136 cleanup=_noop, 

137 ) 

138 ) 

139 

140 

141async def _maybe_call(hook: ServiceHook | None) -> t.Any: 

142 if hook is None: 

143 return None 

144 result = hook() 

145 if inspect.isawaitable(result): 

146 return await result 

147 return result 

148 

149 

150def _init_di_config() -> None: 

151 from session_buddy.di import configure 

152 

153 configure() 

154 

155 

156def _health_di_config() -> bool: 

157 from session_buddy.di.config import SessionPaths 

158 from session_buddy.di.container import depends 

159 

160 with suppress(Exception): 

161 return isinstance(depends.get_sync(SessionPaths), SessionPaths) 

162 return False 

163 

164 

165def _init_permissions_manager() -> None: 

166 from session_buddy.core.permissions import SessionPermissionsManager 

167 from session_buddy.di.config import SessionPaths 

168 from session_buddy.di.container import depends 

169 

170 with suppress(Exception): 

171 if isinstance( 

172 depends.get_sync(SessionPermissionsManager), SessionPermissionsManager 

173 ): 

174 return 

175 

176 paths = _ensure_session_paths() 

177 manager = SessionPermissionsManager(paths.claude_dir) 

178 depends.set(SessionPermissionsManager, manager) 

179 

180 

181def _health_permissions_manager() -> bool: 

182 from session_buddy.core.permissions import SessionPermissionsManager 

183 from session_buddy.di.container import depends 

184 

185 with suppress(Exception): 

186 return isinstance( 

187 depends.get_sync(SessionPermissionsManager), SessionPermissionsManager 

188 ) 

189 return False 

190 

191 

192def _init_lifecycle_manager() -> None: 

193 from session_buddy.core.session_manager import SessionLifecycleManager 

194 from session_buddy.di.container import depends 

195 

196 with suppress(Exception): 

197 if isinstance( 

198 depends.get_sync(SessionLifecycleManager), SessionLifecycleManager 

199 ): 

200 return 

201 

202 depends.set(SessionLifecycleManager, SessionLifecycleManager()) 

203 

204 

205def _health_lifecycle_manager() -> bool: 

206 from session_buddy.core.session_manager import SessionLifecycleManager 

207 from session_buddy.di.container import depends 

208 

209 with suppress(Exception): 

210 return isinstance( 

211 depends.get_sync(SessionLifecycleManager), SessionLifecycleManager 

212 ) 

213 return False 

214 

215 

216async def _init_reflection_db() -> None: 

217 from session_buddy.adapters.lifecycle import init_reflection_adapter 

218 

219 with suppress(Exception): 

220 await init_reflection_adapter() 

221 

222 

223def _health_reflection_db() -> bool: 

224 from session_buddy.adapters.lifecycle import health_reflection_adapter 

225 

226 return health_reflection_adapter() 

227 

228 

229async def _cleanup_reflection_db() -> None: 

230 from session_buddy.adapters.lifecycle import cleanup_reflection_adapter 

231 

232 with suppress(Exception): 

233 await cleanup_reflection_adapter() 

234 

235 

236async def _init_knowledge_graph() -> None: 

237 from session_buddy.adapters.lifecycle import init_knowledge_graph_adapter 

238 

239 with suppress(Exception): 

240 await init_knowledge_graph_adapter() 

241 

242 

243def _health_knowledge_graph() -> bool: 

244 from session_buddy.adapters.lifecycle import health_knowledge_graph_adapter 

245 

246 return health_knowledge_graph_adapter() 

247 

248 

249def _cleanup_knowledge_graph() -> None: 

250 from session_buddy.adapters.lifecycle import cleanup_knowledge_graph_adapter 

251 

252 with suppress(Exception): 

253 cleanup_knowledge_graph_adapter() 

254 

255 

256def _init_storage_adapters() -> None: 

257 from session_buddy.adapters.lifecycle import init_storage_adapters 

258 

259 with suppress(Exception): 

260 init_storage_adapters() 

261 

262 

263def _health_storage_adapters() -> bool: 

264 from session_buddy.adapters.lifecycle import health_storage_adapters 

265 

266 return health_storage_adapters() 

267 

268 

269async def _cleanup_storage_adapters() -> None: 

270 from session_buddy.adapters.lifecycle import cleanup_storage_adapters 

271 

272 with suppress(Exception): 

273 await cleanup_storage_adapters() 

274 

275 

276def _init_cache_adapters() -> None: 

277 from session_buddy.adapters.lifecycle import init_cache_adapters 

278 

279 with suppress(Exception): 

280 init_cache_adapters() 

281 

282 

283def _health_cache_adapters() -> bool: 

284 from session_buddy.adapters.lifecycle import health_cache_adapters 

285 

286 return health_cache_adapters() 

287 

288 

289async def _cleanup_cache_adapters() -> None: 

290 from session_buddy.adapters.lifecycle import cleanup_cache_adapters 

291 

292 with suppress(Exception): 

293 await cleanup_cache_adapters() 

294 

295 

296def _health_tools_registry() -> bool: 

297 with suppress(Exception): 

298 import session_buddy.tools 

299 

300 return True 

301 return False 

302 

303 

304def _init_logging() -> None: 

305 from session_buddy.utils.logging import get_session_logger 

306 

307 get_session_logger() 

308 

309 

310def _health_logging() -> bool: 

311 from session_buddy.di.container import depends 

312 from session_buddy.utils.logging import SessionLogger 

313 

314 with suppress(Exception): 

315 return isinstance(depends.get_sync(SessionLogger), SessionLogger) 

316 return False 

317 

318 

319def _ensure_session_paths() -> SessionPaths: 

320 from session_buddy.di import get_sync_typed 

321 from session_buddy.di.config import SessionPaths 

322 from session_buddy.di.container import depends 

323 

324 with suppress(Exception): 

325 paths = get_sync_typed(SessionPaths) # type: ignore[no-any-return] 

326 if isinstance(paths, SessionPaths): 

327 return paths 

328 

329 paths = SessionPaths.from_home() 

330 paths.ensure_directories() 

331 depends.set(SessionPaths, paths) 

332 return paths 

333 

334 

335def _noop() -> None: 

336 return None 

337 

338 

339__all__ = ["ServiceRegistry", "ServiceSpec", "get_service_registry"]