Coverage for agentos/core/feature_flags.py: 41%
174 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 08:01 +0800
1"""
2AgentOS Feature Flags — Gradual Rollout & Experimentation Engine
3━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
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)
13Architecture:
14 FlagStore (abstract)
15 ├─ InMemoryFlagStore
16 ├─ RedisFlagStore (coming)
17 └─ DatabaseFlagStore (coming)
19 FeatureFlagManager
20 ├─ is_enabled(flag_name, context) → bool
21 ├─ get_variant(flag_name, context) → str
22 └─ set_flag(...) / delete_flag(...)
23"""
25from __future__ import annotations
27import hashlib
28import time
29from abc import ABC, abstractmethod
30from dataclasses import dataclass, field
31from datetime import datetime, timezone
32from enum import Enum
33from typing import Any, Dict, List, Optional, Set
36# ---------------------------------------------------------------------------
37# Data Classes
38# ---------------------------------------------------------------------------
41class FlagType(str, Enum):
42 """Type of feature flag."""
43 BOOLEAN = "boolean" # Simple on/off toggle
44 PERCENTAGE = "percentage" # Gradual rollout (0-100%)
45 VARIANT = "variant" # A/B test variants
46 SCHEDULED = "scheduled" # Time-based enable/disable
49@dataclass
50class FlagRule:
51 """Targeting rule for a feature flag."""
52 flag_type: FlagType = FlagType.BOOLEAN
53 enabled: bool = False
54 rollout_percentage: int = 0 # 0-100 for PERCENTAGE type
55 variants: Dict[str, int] = field(default_factory=dict) # variant_name → weight
56 allowlist_users: Set[str] = field(default_factory=set)
57 allowlist_tenants: Set[str] = field(default_factory=set)
58 blocklist_users: Set[str] = field(default_factory=set)
59 blocklist_tenants: Set[str] = field(default_factory=set)
60 start_time: Optional[datetime] = None
61 end_time: Optional[datetime] = None
62 metadata: Dict[str, Any] = field(default_factory=dict)
64 def to_dict(self) -> Dict[str, Any]:
65 return {
66 "flag_type": self.flag_type.value,
67 "enabled": self.enabled,
68 "rollout_percentage": self.rollout_percentage,
69 "variants": self.variants,
70 "allowlist_users": list(self.allowlist_users),
71 "allowlist_tenants": list(self.allowlist_tenants),
72 "blocklist_users": list(self.blocklist_users),
73 "blocklist_tenants": list(self.blocklist_tenants),
74 "start_time": self.start_time.isoformat() if self.start_time else None,
75 "end_time": self.end_time.isoformat() if self.end_time else None,
76 "metadata": self.metadata,
77 }
80@dataclass
81class FlagContext:
82 """Context for feature flag evaluation."""
83 user_id: Optional[str] = None
84 tenant_id: Optional[str] = None
85 attributes: Dict[str, Any] = field(default_factory=dict)
86 request_id: Optional[str] = None
89@dataclass
90class FlagEvaluation:
91 """Result of a feature flag evaluation."""
92 flag_name: str
93 enabled: bool
94 variant: Optional[str] = None
95 reason: str = ""
96 evaluated_at: float = field(default_factory=time.time)
99# ---------------------------------------------------------------------------
100# Flag Store Interface
101# ---------------------------------------------------------------------------
104class FlagStore(ABC):
105 """Abstract storage backend for feature flags."""
107 @abstractmethod
108 async def get(self, flag_name: str) -> Optional[FlagRule]:
109 """Get a flag rule by name."""
111 @abstractmethod
112 async def set(self, flag_name: str, rule: FlagRule) -> None:
113 """Set or update a flag rule."""
115 @abstractmethod
116 async def delete(self, flag_name: str) -> bool:
117 """Delete a flag. Returns True if existed."""
119 @abstractmethod
120 async def list(self) -> List[str]:
121 """List all flag names."""
124class InMemoryFlagStore(FlagStore):
125 """In-memory flag store for development and testing."""
127 def __init__(self):
128 self._flags: Dict[str, FlagRule] = {}
130 async def get(self, flag_name: str) -> Optional[FlagRule]:
131 return self._flags.get(flag_name)
133 async def set(self, flag_name: str, rule: FlagRule) -> None:
134 self._flags[flag_name] = rule
136 async def delete(self, flag_name: str) -> bool:
137 return self._flags.pop(flag_name, None) is not None
139 async def list(self) -> List[str]:
140 return list(self._flags.keys())
143# ---------------------------------------------------------------------------
144# Feature Flag Manager
145# ---------------------------------------------------------------------------
148class FeatureFlagManager:
149 """
150 Production feature flag manager.
152 Supports:
153 - Boolean toggles (on/off)
154 - Percentage-based gradual rollout
155 - A/B test variants
156 - User/tenant targeting (allowlist/blocklist)
157 - Time-based scheduling
158 - Kill-switch (immediate disable)
160 Usage:
161 manager = FeatureFlagManager(InMemoryFlagStore())
163 # Register a flag
164 await manager.set_flag("new_search", FlagRule(
165 flag_type=FlagType.PERCENTAGE,
166 enabled=True,
167 rollout_percentage=10,
168 ))
170 # Check in application code
171 ctx = FlagContext(user_id="user_123", tenant_id="tenant_a")
172 if await manager.is_enabled("new_search", ctx):
173 use_new_search()
174 """
176 def __init__(self, store: FlagStore):
177 self._store = store
178 self._evaluation_log: List[FlagEvaluation] = []
180 # ── Flag Management ────────────────────────────────────────────────
182 async def set_flag(self, flag_name: str, rule: FlagRule) -> None:
183 """Create or update a feature flag."""
184 await self._store.set(flag_name, rule)
186 async def delete_flag(self, flag_name: str) -> bool:
187 """Delete a feature flag."""
188 return await self._store.delete(flag_name)
190 async def get_flag(self, flag_name: str) -> Optional[FlagRule]:
191 """Get a flag's rule."""
192 return await self._store.get(flag_name)
194 async def list_flags(self) -> List[str]:
195 """List all registered flags."""
196 return await self._store.list()
198 async def kill_switch(self, flag_name: str) -> None:
199 """Emergency disable a flag (kill-switch)."""
200 rule = await self._store.get(flag_name)
201 if rule:
202 rule.enabled = False
203 await self._store.set(flag_name, rule)
205 # ── Flag Evaluation ────────────────────────────────────────────────
207 async def is_enabled(self, flag_name: str, context: Optional[FlagContext] = None) -> bool:
208 """Check if a feature flag is enabled for the given context."""
209 evaluation = await self.evaluate(flag_name, context)
210 return evaluation.enabled
212 async def get_variant(self, flag_name: str, context: Optional[FlagContext] = None) -> Optional[str]:
213 """Get the variant name for an A/B test flag."""
214 evaluation = await self.evaluate(flag_name, context)
215 return evaluation.variant
217 async def evaluate(
218 self, flag_name: str, context: Optional[FlagContext] = None
219 ) -> FlagEvaluation:
220 """Full evaluation of a feature flag with audit trail."""
221 ctx = context or FlagContext()
222 rule = await self._store.get(flag_name)
224 # Flag not found → disabled
225 if rule is None:
226 evaluation = FlagEvaluation(
227 flag_name=flag_name,
228 enabled=False,
229 reason="Flag not found",
230 )
231 self._evaluation_log.append(evaluation)
232 return evaluation
234 # Not enabled at rule level
235 if not rule.enabled:
236 evaluation = FlagEvaluation(
237 flag_name=flag_name,
238 enabled=False,
239 reason="Flag disabled at rule level",
240 )
241 self._evaluation_log.append(evaluation)
242 return evaluation
244 # Blocklist check (highest priority)
245 if ctx.user_id and ctx.user_id in rule.blocklist_users:
246 evaluation = FlagEvaluation(
247 flag_name=flag_name,
248 enabled=False,
249 reason="User in blocklist",
250 )
251 self._evaluation_log.append(evaluation)
252 return evaluation
254 if ctx.tenant_id and ctx.tenant_id in rule.blocklist_tenants:
255 evaluation = FlagEvaluation(
256 flag_name=flag_name,
257 enabled=False,
258 reason="Tenant in blocklist",
259 )
260 self._evaluation_log.append(evaluation)
261 return evaluation
263 # Allowlist check — exclusive: if allowlist is non-empty and user/tenant NOT in it, deny
264 has_user_allowlist = bool(rule.allowlist_users)
265 has_tenant_allowlist = bool(rule.allowlist_tenants)
267 if has_user_allowlist and (not ctx.user_id or ctx.user_id not in rule.allowlist_users):
268 evaluation = FlagEvaluation(
269 flag_name=flag_name,
270 enabled=False,
271 reason="User not in allowlist",
272 )
273 self._evaluation_log.append(evaluation)
274 return evaluation
276 if has_tenant_allowlist and (not ctx.tenant_id or ctx.tenant_id not in rule.allowlist_tenants):
277 evaluation = FlagEvaluation(
278 flag_name=flag_name,
279 enabled=False,
280 reason="Tenant not in allowlist",
281 )
282 self._evaluation_log.append(evaluation)
283 return evaluation
285 if ctx.user_id and ctx.user_id in rule.allowlist_users:
286 return self._enabled_eval(flag_name, rule, "User in allowlist")
288 if ctx.tenant_id and ctx.tenant_id in rule.allowlist_tenants:
289 return self._enabled_eval(flag_name, rule, "Tenant in allowlist")
291 # Time-based scheduling
292 now = datetime.now(timezone.utc)
293 if rule.start_time and now < rule.start_time:
294 evaluation = FlagEvaluation(
295 flag_name=flag_name,
296 enabled=False,
297 reason="Before start_time",
298 )
299 self._evaluation_log.append(evaluation)
300 return evaluation
302 if rule.end_time and now > rule.end_time:
303 evaluation = FlagEvaluation(
304 flag_name=flag_name,
305 enabled=False,
306 reason="After end_time",
307 )
308 self._evaluation_log.append(evaluation)
309 return evaluation
311 # Type-specific evaluation
312 if rule.flag_type == FlagType.BOOLEAN:
313 return self._enabled_eval(flag_name, rule, "Boolean: enabled=True")
315 elif rule.flag_type == FlagType.PERCENTAGE:
316 hash_val = self._hash_context(flag_name, ctx)
317 bucket = hash_val % 100
318 if bucket < rule.rollout_percentage:
319 return self._enabled_eval(
320 flag_name, rule,
321 f"Percentage: bucket {bucket} < {rule.rollout_percentage}%"
322 )
323 else:
324 evaluation = FlagEvaluation(
325 flag_name=flag_name,
326 enabled=False,
327 reason=f"Percentage: bucket {bucket} >= {rule.rollout_percentage}%",
328 )
329 self._evaluation_log.append(evaluation)
330 return evaluation
332 elif rule.flag_type == FlagType.VARIANT:
333 variant = self._select_variant(flag_name, ctx, rule)
334 if variant:
335 evaluation = FlagEvaluation(
336 flag_name=flag_name,
337 enabled=True,
338 variant=variant,
339 reason=f"Variant selected: {variant}",
340 )
341 self._evaluation_log.append(evaluation)
342 return evaluation
343 else:
344 evaluation = FlagEvaluation(
345 flag_name=flag_name,
346 enabled=False,
347 reason="Variant: no variant selected",
348 )
349 self._evaluation_log.append(evaluation)
350 return evaluation
352 elif rule.flag_type == FlagType.SCHEDULED:
353 return self._enabled_eval(flag_name, rule, "Scheduled: within time window")
355 # Fallback
356 evaluation = FlagEvaluation(
357 flag_name=flag_name,
358 enabled=False,
359 reason="Unknown flag type",
360 )
361 self._evaluation_log.append(evaluation)
362 return evaluation
364 # ── Internal Helpers ───────────────────────────────────────────────
366 def _hash_context(self, flag_name: str, ctx: FlagContext) -> int:
367 """Deterministic hash for percentage-based rollout."""
368 seed = f"{flag_name}:{ctx.user_id or ''}:{ctx.tenant_id or ''}"
369 return int(hashlib.md5(seed.encode()).hexdigest(), 16)
371 def _select_variant(
372 self, flag_name: str, ctx: FlagContext, rule: FlagRule
373 ) -> Optional[str]:
374 """Select a variant based on weighted distribution."""
375 if not rule.variants:
376 return None
378 hash_val = self._hash_context(flag_name, ctx)
379 bucket = hash_val % 100
380 cumulative = 0
381 for variant_name, weight in rule.variants.items():
382 cumulative += weight
383 if bucket < cumulative:
384 return variant_name
385 return None
387 def _enabled_eval(
388 self, flag_name: str, rule: FlagRule, reason: str
389 ) -> FlagEvaluation:
390 evaluation = FlagEvaluation(
391 flag_name=flag_name,
392 enabled=True,
393 reason=reason,
394 )
395 self._evaluation_log.append(evaluation)
396 return evaluation
398 # ── Audit ──────────────────────────────────────────────────────────
400 def get_evaluation_log(self, limit: int = 100) -> List[FlagEvaluation]:
401 """Get recent flag evaluations for audit."""
402 return self._evaluation_log[-limit:]
404 def clear_evaluation_log(self) -> None:
405 self._evaluation_log.clear()
408# ---------------------------------------------------------------------------
409# Convenience
410# ---------------------------------------------------------------------------
412def create_flag_manager() -> FeatureFlagManager:
413 """Create a FeatureFlagManager with in-memory store."""
414 return FeatureFlagManager(InMemoryFlagStore())