1 from __future__ import annotations

2

3 import inspect

4 import typing as t

5 from contextlib import suppress

6 from dataclasses import dataclass, field

7

8 if t.TYPE_CHECKING:

9 from session_buddy.di import get_sync_typed

10 from session_buddy.di.config import SessionPaths

11

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

13

14

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

16 class ServiceSpec:

17 name: str

18 category: str

19 init: ServiceHook | None = None

20 health: ServiceHook | None = None

21 cleanup: ServiceHook | None = None

22

23

24 class 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

49 def get_service_registry() -> ServiceRegistry:

50 global _registry

51 if _registry is None:

52 _registry = ServiceRegistry()

53 _register_defaults(_registry)

54 return _registry

55

56

57 def _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

141 async 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

150 def _init_di_config() -> None:

151 from session_buddy.di import configure

152

153 configure()

154

155

156 def _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

165 def _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

181 def _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

192 def _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

205 def _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

216 async 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

223 def _health_reflection_db() -> bool:

224 from session_buddy.adapters.lifecycle import health_reflection_adapter

225

226 return health_reflection_adapter()

227

228

229 async 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

236 async 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

243 def _health_knowledge_graph() -> bool:

244 from session_buddy.adapters.lifecycle import health_knowledge_graph_adapter

245

246 return health_knowledge_graph_adapter()

247

248

249 def _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

256 def _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

263 def _health_storage_adapters() -> bool:

264 from session_buddy.adapters.lifecycle import health_storage_adapters

265

266 return health_storage_adapters()

267

268

269 async 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

276 def _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

283 def _health_cache_adapters() -> bool:

284 from session_buddy.adapters.lifecycle import health_cache_adapters

285

286 return health_cache_adapters()

287

288

289 async 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

296 def _health_tools_registry() -> bool:

297 with suppress(Exception):

298 import session_buddy.tools

299

300 return True

301 return False

302

303

304 def _init_logging() -> None:

305 from session_buddy.utils.logging import get_session_logger

306

307 get_session_logger()

308

309

310 def _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

319 def _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

335 def _noop() -> None:

336 return None

337

338

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