Coverage for agentos/core/feature_flags.py: 0%

175 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-08 13:14 +0800

1""" 

2AgentOS Feature Flags — Gradual Rollout & Experimentation Engine 

3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 

4 

5Production-grade feature flag system with: 

6 - Percentage-based gradual rollout 

7 - User/tenant targeting (allowlist/blocklist) 

8 - Time-based scheduling (start/end dates) 

9 - Kill-switch for emergency disable 

10 - Audit trail for flag changes 

11 - Backend-agnostic storage (in-memory / Redis / DB) 

12 

13Architecture: 

14 FlagStore (abstract) 

15 ├─ InMemoryFlagStore 

16 ├─ RedisFlagStore (coming) 

17 └─ DatabaseFlagStore (coming) 

18 

19 FeatureFlagManager 

20 ├─ is_enabled(flag_name, context) → bool 

21 ├─ get_variant(flag_name, context) → str 

22 └─ set_flag(...) / delete_flag(...) 

23""" 

24 

25from __future__ import annotations 

26 

27import builtins 

28import hashlib 

29import time 

30from abc import ABC, abstractmethod 

31from dataclasses import dataclass, field 

32from datetime import UTC, datetime 

33from enum import StrEnum 

34from typing import Any 

35 

36# --------------------------------------------------------------------------- 

37# Data Classes 

38# --------------------------------------------------------------------------- 

39 

40 

41class FlagType(StrEnum): 

42 """Type of feature flag.""" 

43 

44 BOOLEAN = "boolean" # Simple on/off toggle 

45 PERCENTAGE = "percentage" # Gradual rollout (0-100%) 

46 VARIANT = "variant" # A/B test variants 

47 SCHEDULED = "scheduled" # Time-based enable/disable 

48 

49 

50@dataclass 

51class FlagRule: 

52 """Targeting rule for a feature flag.""" 

53 

54 flag_type: FlagType = FlagType.BOOLEAN 

55 enabled: bool = False 

56 rollout_percentage: int = 0 # 0-100 for PERCENTAGE type 

57 variants: dict[str, int] = field(default_factory=dict) # variant_name → weight 

58 allowlist_users: set[str] = field(default_factory=set) 

59 allowlist_tenants: set[str] = field(default_factory=set) 

60 blocklist_users: set[str] = field(default_factory=set) 

61 blocklist_tenants: set[str] = field(default_factory=set) 

62 start_time: datetime | None = None 

63 end_time: datetime | None = None 

64 metadata: dict[str, Any] = field(default_factory=dict) 

65 

66 def to_dict(self) -> dict[str, Any]: 

67 return { 

68 "flag_type": self.flag_type.value, 

69 "enabled": self.enabled, 

70 "rollout_percentage": self.rollout_percentage, 

71 "variants": self.variants, 

72 "allowlist_users": list(self.allowlist_users), 

73 "allowlist_tenants": list(self.allowlist_tenants), 

74 "blocklist_users": list(self.blocklist_users), 

75 "blocklist_tenants": list(self.blocklist_tenants), 

76 "start_time": self.start_time.isoformat() if self.start_time else None, 

77 "end_time": self.end_time.isoformat() if self.end_time else None, 

78 "metadata": self.metadata, 

79 } 

80 

81 

82@dataclass 

83class FlagContext: 

84 """Context for feature flag evaluation.""" 

85 

86 user_id: str | None = None 

87 tenant_id: str | None = None 

88 attributes: dict[str, Any] = field(default_factory=dict) 

89 request_id: str | None = None 

90 

91 

92@dataclass 

93class FlagEvaluation: 

94 """Result of a feature flag evaluation.""" 

95 

96 flag_name: str 

97 enabled: bool 

98 variant: str | None = None 

99 reason: str = "" 

100 evaluated_at: float = field(default_factory=time.time) 

101 

102 

103# --------------------------------------------------------------------------- 

104# Flag Store Interface 

105# --------------------------------------------------------------------------- 

106 

107 

108class FlagStore(ABC): 

109 """Abstract storage backend for feature flags.""" 

110 

111 @abstractmethod 

112 async def get(self, flag_name: str) -> FlagRule | None: 

113 """Get a flag rule by name.""" 

114 

115 @abstractmethod 

116 async def set(self, flag_name: str, rule: FlagRule) -> None: 

117 """Set or update a flag rule.""" 

118 

119 @abstractmethod 

120 async def delete(self, flag_name: str) -> bool: 

121 """Delete a flag. Returns True if existed.""" 

122 

123 @abstractmethod 

124 async def list(self) -> builtins.list[str]: 

125 """List all flag names.""" 

126 

127 

128class InMemoryFlagStore(FlagStore): 

129 """In-memory flag store for development and testing.""" 

130 

131 def __init__(self): 

132 self._flags: dict[str, FlagRule] = {} 

133 

134 async def get(self, flag_name: str) -> FlagRule | None: 

135 return self._flags.get(flag_name) 

136 

137 async def set(self, flag_name: str, rule: FlagRule) -> None: 

138 self._flags[flag_name] = rule 

139 

140 async def delete(self, flag_name: str) -> bool: 

141 return self._flags.pop(flag_name, None) is not None 

142 

143 async def list(self) -> builtins.list[str]: 

144 return list(self._flags.keys()) 

145 

146 

147# --------------------------------------------------------------------------- 

148# Feature Flag Manager 

149# --------------------------------------------------------------------------- 

150 

151 

152class FeatureFlagManager: 

153 """ 

154 Production feature flag manager. 

155 

156 Supports: 

157 - Boolean toggles (on/off) 

158 - Percentage-based gradual rollout 

159 - A/B test variants 

160 - User/tenant targeting (allowlist/blocklist) 

161 - Time-based scheduling 

162 - Kill-switch (immediate disable) 

163 

164 Usage: 

165 manager = FeatureFlagManager(InMemoryFlagStore()) 

166 

167 # Register a flag 

168 await manager.set_flag("new_search", FlagRule( 

169 flag_type=FlagType.PERCENTAGE, 

170 enabled=True, 

171 rollout_percentage=10, 

172 )) 

173 

174 # Check in application code 

175 ctx = FlagContext(user_id="user_123", tenant_id="tenant_a") 

176 if await manager.is_enabled("new_search", ctx): 

177 use_new_search() 

178 """ 

179 

180 def __init__(self, store: FlagStore): 

181 self._store = store 

182 self._evaluation_log: list[FlagEvaluation] = [] 

183 

184 # ── Flag Management ──────────────────────────────────────────────── 

185 

186 async def set_flag(self, flag_name: str, rule: FlagRule) -> None: 

187 """Create or update a feature flag.""" 

188 await self._store.set(flag_name, rule) 

189 

190 async def delete_flag(self, flag_name: str) -> bool: 

191 """Delete a feature flag.""" 

192 return await self._store.delete(flag_name) 

193 

194 async def get_flag(self, flag_name: str) -> FlagRule | None: 

195 """Get a flag's rule.""" 

196 return await self._store.get(flag_name) 

197 

198 async def list_flags(self) -> list[str]: 

199 """List all registered flags.""" 

200 return await self._store.list() 

201 

202 async def kill_switch(self, flag_name: str) -> None: 

203 """Emergency disable a flag (kill-switch).""" 

204 rule = await self._store.get(flag_name) 

205 if rule: 

206 rule.enabled = False 

207 await self._store.set(flag_name, rule) 

208 

209 # ── Flag Evaluation ──────────────────────────────────────────────── 

210 

211 async def is_enabled(self, flag_name: str, context: FlagContext | None = None) -> bool: 

212 """Check if a feature flag is enabled for the given context.""" 

213 evaluation = await self.evaluate(flag_name, context) 

214 return evaluation.enabled 

215 

216 async def get_variant(self, flag_name: str, context: FlagContext | None = None) -> str | None: 

217 """Get the variant name for an A/B test flag.""" 

218 evaluation = await self.evaluate(flag_name, context) 

219 return evaluation.variant 

220 

221 async def evaluate(self, flag_name: str, context: FlagContext | None = None) -> FlagEvaluation: 

222 """Full evaluation of a feature flag with audit trail.""" 

223 ctx = context or FlagContext() 

224 rule = await self._store.get(flag_name) 

225 

226 # Flag not found → disabled 

227 if rule is None: 

228 evaluation = FlagEvaluation( 

229 flag_name=flag_name, 

230 enabled=False, 

231 reason="Flag not found", 

232 ) 

233 self._evaluation_log.append(evaluation) 

234 return evaluation 

235 

236 # Not enabled at rule level 

237 if not rule.enabled: 

238 evaluation = FlagEvaluation( 

239 flag_name=flag_name, 

240 enabled=False, 

241 reason="Flag disabled at rule level", 

242 ) 

243 self._evaluation_log.append(evaluation) 

244 return evaluation 

245 

246 # Blocklist check (highest priority) 

247 if ctx.user_id and ctx.user_id in rule.blocklist_users: 

248 evaluation = FlagEvaluation( 

249 flag_name=flag_name, 

250 enabled=False, 

251 reason="User in blocklist", 

252 ) 

253 self._evaluation_log.append(evaluation) 

254 return evaluation 

255 

256 if ctx.tenant_id and ctx.tenant_id in rule.blocklist_tenants: 

257 evaluation = FlagEvaluation( 

258 flag_name=flag_name, 

259 enabled=False, 

260 reason="Tenant in blocklist", 

261 ) 

262 self._evaluation_log.append(evaluation) 

263 return evaluation 

264 

265 # Allowlist check — exclusive: if allowlist is non-empty and user/tenant NOT in it, deny 

266 has_user_allowlist = bool(rule.allowlist_users) 

267 has_tenant_allowlist = bool(rule.allowlist_tenants) 

268 

269 if has_user_allowlist and (not ctx.user_id or ctx.user_id not in rule.allowlist_users): 

270 evaluation = FlagEvaluation( 

271 flag_name=flag_name, 

272 enabled=False, 

273 reason="User not in allowlist", 

274 ) 

275 self._evaluation_log.append(evaluation) 

276 return evaluation 

277 

278 if has_tenant_allowlist and ( 

279 not ctx.tenant_id or ctx.tenant_id not in rule.allowlist_tenants 

280 ): 

281 evaluation = FlagEvaluation( 

282 flag_name=flag_name, 

283 enabled=False, 

284 reason="Tenant not in allowlist", 

285 ) 

286 self._evaluation_log.append(evaluation) 

287 return evaluation 

288 

289 if ctx.user_id and ctx.user_id in rule.allowlist_users: 

290 return self._enabled_eval(flag_name, rule, "User in allowlist") 

291 

292 if ctx.tenant_id and ctx.tenant_id in rule.allowlist_tenants: 

293 return self._enabled_eval(flag_name, rule, "Tenant in allowlist") 

294 

295 # Time-based scheduling 

296 now = datetime.now(UTC) 

297 if rule.start_time and now < rule.start_time: 

298 evaluation = FlagEvaluation( 

299 flag_name=flag_name, 

300 enabled=False, 

301 reason="Before start_time", 

302 ) 

303 self._evaluation_log.append(evaluation) 

304 return evaluation 

305 

306 if rule.end_time and now > rule.end_time: 

307 evaluation = FlagEvaluation( 

308 flag_name=flag_name, 

309 enabled=False, 

310 reason="After end_time", 

311 ) 

312 self._evaluation_log.append(evaluation) 

313 return evaluation 

314 

315 # Type-specific evaluation 

316 if rule.flag_type == FlagType.BOOLEAN: 

317 return self._enabled_eval(flag_name, rule, "Boolean: enabled=True") 

318 

319 elif rule.flag_type == FlagType.PERCENTAGE: 

320 hash_val = self._hash_context(flag_name, ctx) 

321 bucket = hash_val % 100 

322 if bucket < rule.rollout_percentage: 

323 return self._enabled_eval( 

324 flag_name, rule, f"Percentage: bucket {bucket} < {rule.rollout_percentage}%" 

325 ) 

326 else: 

327 evaluation = FlagEvaluation( 

328 flag_name=flag_name, 

329 enabled=False, 

330 reason=f"Percentage: bucket {bucket} >= {rule.rollout_percentage}%", 

331 ) 

332 self._evaluation_log.append(evaluation) 

333 return evaluation 

334 

335 elif rule.flag_type == FlagType.VARIANT: 

336 variant = self._select_variant(flag_name, ctx, rule) 

337 if variant: 

338 evaluation = FlagEvaluation( 

339 flag_name=flag_name, 

340 enabled=True, 

341 variant=variant, 

342 reason=f"Variant selected: {variant}", 

343 ) 

344 self._evaluation_log.append(evaluation) 

345 return evaluation 

346 else: 

347 evaluation = FlagEvaluation( 

348 flag_name=flag_name, 

349 enabled=False, 

350 reason="Variant: no variant selected", 

351 ) 

352 self._evaluation_log.append(evaluation) 

353 return evaluation 

354 

355 elif rule.flag_type == FlagType.SCHEDULED: 

356 return self._enabled_eval(flag_name, rule, "Scheduled: within time window") 

357 

358 # Fallback 

359 evaluation = FlagEvaluation( 

360 flag_name=flag_name, 

361 enabled=False, 

362 reason="Unknown flag type", 

363 ) 

364 self._evaluation_log.append(evaluation) 

365 return evaluation 

366 

367 # ── Internal Helpers ─────────────────────────────────────────────── 

368 

369 def _hash_context(self, flag_name: str, ctx: FlagContext) -> int: 

370 """Deterministic hash for percentage-based rollout.""" 

371 seed = f"{flag_name}:{ctx.user_id or ''}:{ctx.tenant_id or ''}" 

372 return int(hashlib.md5(seed.encode()).hexdigest(), 16) 

373 

374 def _select_variant(self, flag_name: str, ctx: FlagContext, rule: FlagRule) -> str | None: 

375 """Select a variant based on weighted distribution.""" 

376 if not rule.variants: 

377 return None 

378 

379 hash_val = self._hash_context(flag_name, ctx) 

380 bucket = hash_val % 100 

381 cumulative = 0 

382 for variant_name, weight in rule.variants.items(): 

383 cumulative += weight 

384 if bucket < cumulative: 

385 return variant_name 

386 return None 

387 

388 def _enabled_eval(self, flag_name: str, rule: FlagRule, reason: str) -> FlagEvaluation: 

389 evaluation = FlagEvaluation( 

390 flag_name=flag_name, 

391 enabled=True, 

392 reason=reason, 

393 ) 

394 self._evaluation_log.append(evaluation) 

395 return evaluation 

396 

397 # ── Audit ────────────────────────────────────────────────────────── 

398 

399 def get_evaluation_log(self, limit: int = 100) -> list[FlagEvaluation]: 

400 """Get recent flag evaluations for audit.""" 

401 return self._evaluation_log[-limit:] 

402 

403 def clear_evaluation_log(self) -> None: 

404 self._evaluation_log.clear() 

405 

406 

407# --------------------------------------------------------------------------- 

408# Convenience 

409# --------------------------------------------------------------------------- 

410 

411 

412def create_flag_manager() -> FeatureFlagManager: 

413 """Create a FeatureFlagManager with in-memory store.""" 

414 return FeatureFlagManager(InMemoryFlagStore())