Coverage for astrocyte/_policy.py: 98%
144 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Policy enforcer — access control, rate limiting, PII scanning, input validation."""
3from __future__ import annotations
5import threading
6from dataclasses import dataclass
7from typing import TYPE_CHECKING
9from astrocyte._validation import validate_bank_id
10from astrocyte.config import AstrocyteConfig, BarrierConfig, HomeostasisConfig
11from astrocyte.errors import AccessDenied, ConfigError, RateLimited
12from astrocyte.identity import (
13 BankResolver,
14 accessible_read_banks,
15 context_principal_label,
16 effective_permissions,
17)
18from astrocyte.policy.barriers import ContentValidator, MetadataSanitizer, PiiScanner
19from astrocyte.policy.escalation import CircuitBreaker, DegradedModeHandler
20from astrocyte.policy.homeostasis import QuotaTracker, RateLimiter
21from astrocyte.policy.signal_quality import NoisyBankDetector
22from astrocyte.types import (
23 AccessGrant,
24 AstrocyteContext,
25 PiiMatch,
26 RecallResult,
27 RetainResult,
28)
30if TYPE_CHECKING:
31 from astrocyte.types import Metadata
34@dataclass
35class _BankPolicy:
36 """Barrier + homeostasis enforcement built from one resolved config pair."""
38 homeostasis: HomeostasisConfig
39 barriers: BarrierConfig
40 pii_scanner: PiiScanner
41 content_validator: ContentValidator
42 metadata_sanitizer: MetadataSanitizer
43 rate_limiters: dict[str, RateLimiter]
44 quota_limits: dict[str, int | None]
46 @classmethod
47 def build(cls, homeostasis: HomeostasisConfig, barriers: BarrierConfig) -> _BankPolicy:
48 rl = homeostasis.rate_limits
49 per_minute = {
50 "retain": rl.retain_per_minute,
51 "recall": rl.recall_per_minute,
52 "reflect": rl.reflect_per_minute,
53 }
54 return cls(
55 homeostasis=homeostasis,
56 barriers=barriers,
57 pii_scanner=PiiScanner(
58 mode=barriers.pii.mode,
59 action=barriers.pii.action,
60 countries=barriers.pii.countries,
61 type_overrides=barriers.pii.type_overrides,
62 ),
63 content_validator=ContentValidator(
64 max_content_length=barriers.validation.max_content_length,
65 reject_empty=barriers.validation.reject_empty_content,
66 allowed_content_types=barriers.validation.allowed_content_types,
67 ),
68 metadata_sanitizer=MetadataSanitizer(
69 blocked_keys=barriers.metadata.blocked_keys,
70 max_size_bytes=barriers.metadata.max_metadata_size_bytes,
71 ),
72 rate_limiters={op: RateLimiter(limit) for op, limit in per_minute.items() if limit},
73 quota_limits={
74 "retain": homeostasis.quotas.retain_per_day,
75 "reflect": homeostasis.quotas.reflect_per_day,
76 },
77 )
80@dataclass(frozen=True)
81class NoisyBankVerdict:
82 """Outcome of a noisy-bank check: why the bank is flagged (empty = not),
83 whether that changed since its last check, and the configured action."""
85 reasons: tuple[str, ...]
86 changed: bool
87 action: str
90class PolicyEnforcer:
91 """Centralizes all policy enforcement: access control, rate limiting,
92 PII scanning, content validation, metadata sanitization, circuit breaking,
93 and input validation.
95 Barriers and homeostasis are per bank: a bank with its own
96 ``banks.<id>.homeostasis`` / ``barriers`` (or a bank ``profile``) gets its
97 own scanner, validator, sanitizer, rate limits and quotas; every other bank
98 shares the top-level set.
99 """
101 def __init__(self, config: AstrocyteConfig) -> None:
102 self._config = config
104 self._default_policy = _BankPolicy.build(config.homeostasis, config.barriers)
105 self._bank_policies: dict[str, _BankPolicy] = {
106 bank_id: _BankPolicy.build(config.bank_homeostasis(bank_id), config.bank_barriers(bank_id))
107 for bank_id, bank in (config.banks or {}).items()
108 if bank.homeostasis is not None or bank.barriers is not None
109 }
110 # Top-level components under their historical names (introspection).
111 self._pii_scanner = self._default_policy.pii_scanner
112 self._content_validator = self._default_policy.content_validator
113 self._metadata_sanitizer = self._default_policy.metadata_sanitizer
114 self._rate_limiters = self._default_policy.rate_limiters
115 self._quota_limits = self._default_policy.quota_limits
117 # Quota tracker (keyed by bank, so one tracker serves every bank's limits)
118 self._quota_tracker = QuotaTracker()
120 # Noisy-bank detection (signal_quality.noisy_bank, per bank)
121 self._noisy_banks = NoisyBankDetector()
123 # Atomic lock for rate + quota checks
124 self._rate_quota_lock = threading.Lock()
126 # Circuit breaker
127 cb = config.escalation.circuit_breaker
128 self._circuit_breaker = CircuitBreaker(
129 failure_threshold=cb.failure_threshold,
130 recovery_timeout_seconds=cb.recovery_timeout_seconds,
131 half_open_max_calls=cb.half_open_max_calls,
132 )
133 self._degraded_handler = DegradedModeHandler(mode=config.escalation.degraded_mode)
135 # Access control grants
136 self._access_grants: list[AccessGrant] = []
138 # -- Access grants --
140 def set_access_grants(self, grants: list[AccessGrant]) -> None:
141 """Configure access grants."""
142 self._access_grants = grants
144 @property
145 def access_grants(self) -> list[AccessGrant]:
146 return self._access_grants
148 # -- Access control --
150 def check_access(self, bank_id: str, permission: str, context: AstrocyteContext | None) -> None:
151 """Check access control. Raises AccessDenied if denied."""
152 if not self._config.access_control.enabled:
153 return
154 if context is None:
155 if self._config.access_control.default_policy == "open":
156 return
157 raise AccessDenied("anonymous", bank_id, permission)
159 eff = effective_permissions(context, self._access_grants, bank_id)
160 if permission in eff:
161 return
163 if self._config.access_control.default_policy == "open":
164 return
166 raise AccessDenied(context_principal_label(context), bank_id, permission)
168 # -- Bank resolution --
170 def make_bank_resolver(self) -> BankResolver:
171 i = self._config.identity
172 return BankResolver(
173 user_prefix=i.user_bank_prefix,
174 agent_prefix=i.agent_bank_prefix,
175 service_prefix=i.service_bank_prefix,
176 )
178 def resolve_read_bank_ids(
179 self,
180 bank_id: str | None,
181 banks: list[str] | None,
182 context: AstrocyteContext | None,
183 ) -> list[str]:
184 """Resolve bank list for recall/reflect; optional identity-driven auto-resolve."""
185 bank_ids = banks or ([bank_id] if bank_id else [])
186 if not bank_ids and self._config.identity.auto_resolve_banks and context is not None:
187 known = list((self._config.banks or {}).keys())
188 bank_ids = accessible_read_banks(
189 context,
190 self._access_grants,
191 known_bank_ids=known or None,
192 resolver=self.make_bank_resolver(),
193 )
194 if not bank_ids:
195 raise ConfigError("Either bank_id or banks must be provided")
196 for bid in bank_ids:
197 validate_bank_id(bid)
198 return bank_ids
200 # -- Rate limiting + quota --
202 def check_rate_and_quota(self, bank_id: str, operation: str) -> None:
203 """Atomically check rate limit and quota under a shared lock."""
204 with self._rate_quota_lock:
205 self._check_rate_limit(bank_id, operation)
206 self._check_quota(bank_id, operation)
208 def _for(self, bank_id: str | None) -> _BankPolicy:
209 return self._bank_policies.get(bank_id, self._default_policy) if bank_id else self._default_policy
211 def _check_rate_limit(self, bank_id: str, operation: str) -> None:
212 limiter = self._for(bank_id).rate_limiters.get(operation)
213 if limiter:
214 limiter.check_and_record(bank_id, operation)
216 def _check_quota(self, bank_id: str, operation: str) -> None:
217 limit = self._for(bank_id).quota_limits.get(operation)
218 if not self._quota_tracker.check(bank_id, operation, limit):
219 raise RateLimited(bank_id=bank_id, operation=operation)
221 def record_quota(self, bank_id: str, operation: str) -> None:
222 """Record a successful operation against the quota tracker."""
223 self._quota_tracker.record(bank_id, operation)
225 # -- PII scanning --
227 async def scan_pii(
228 self, content: str, mode: str | None = None, bank_id: str | None = None
229 ) -> tuple[str, list[PiiMatch]]:
230 """Scan content for PII under ``bank_id``'s barriers (top-level when
231 ``None``). ``mode`` defaults to that bank's ``barriers.pii.mode``.
232 Returns (possibly redacted content, matches)."""
233 policy = self._for(bank_id)
234 if (mode or policy.barriers.pii.mode) in ("llm", "rules_then_llm"):
235 return await policy.pii_scanner.apply_async(content)
236 return policy.pii_scanner.apply(content)
238 @property
239 def pii_action(self) -> str:
240 return self._pii_scanner.action
242 def pii_action_for(self, bank_id: str | None) -> str:
243 """``barriers.pii.action`` in effect for ``bank_id``."""
244 return self._for(bank_id).pii_scanner.action
246 def scan_pii_output(self, text: str) -> list[PiiMatch]:
247 """Scan text for PII matches (for DLP output scanning)."""
248 return self._pii_scanner.scan(text)
250 # -- Content validation --
252 def validate_content(self, content: str, content_type: str, bank_id: str | None = None) -> list[str]:
253 """Validate content under ``bank_id``'s barriers. Returns list of error strings (empty = valid)."""
254 return self._for(bank_id).content_validator.validate(content, content_type)
256 # -- Metadata sanitization --
258 def sanitize_metadata(
259 self, metadata: "Metadata | None", bank_id: str | None = None
260 ) -> tuple["Metadata | None", list[str]]:
261 """Sanitize metadata under ``bank_id``'s barriers. Returns (sanitized metadata, warnings)."""
262 return self._for(bank_id).metadata_sanitizer.sanitize(metadata)
264 # -- Noisy-bank detection --
266 def check_noisy_bank(self, bank_id: str) -> NoisyBankVerdict | None:
267 """Check ``bank_id`` against its ``signal_quality.noisy_bank`` settings.
268 ``None`` when detection is disabled for the bank."""
269 cfg = self._config.bank_signal_quality(bank_id).noisy_bank
270 if not cfg.enabled:
271 return None
272 reasons, changed = self._noisy_banks.check(
273 bank_id,
274 retain_spike_multiplier=cfg.retain_spike_multiplier,
275 min_avg_content_length=cfg.min_avg_content_length,
276 max_dedup_rate=cfg.max_dedup_rate,
277 )
278 return NoisyBankVerdict(reasons=reasons, changed=changed, action=cfg.action)
280 def record_retain_signal(self, bank_id: str, content_length: int, deduplicated: bool) -> None:
281 """Feed one processed retain to noisy-bank detection."""
282 self._noisy_banks.record(bank_id, content_length, deduplicated)
284 # -- Token budgets --
286 def token_budget(self, bank_ids: list[str], operation: str) -> int | None:
287 """``homeostasis.<operation>_max_tokens`` across ``bank_ids``: the
288 strictest budget any of them sets (``None`` when none sets one), since
289 a multi-bank recall or reflect returns one merged result."""
290 budgets = [getattr(self._config.bank_homeostasis(b), f"{operation}_max_tokens") for b in bank_ids]
291 set_budgets = [b for b in budgets if b is not None]
292 return min(set_budgets) if set_budgets else None
294 # -- Circuit breaker --
296 def check_circuit(self, provider_name: str) -> None:
297 """Check circuit breaker. Raises ProviderUnavailable if open."""
298 self._circuit_breaker.check(provider_name)
300 def record_success(self) -> None:
301 self._circuit_breaker.record_success()
303 def record_failure(self) -> None:
304 self._circuit_breaker.record_failure()
306 def handle_degraded_retain(self, provider_name: str) -> RetainResult:
307 self._degraded_handler.handle_retain(provider_name)
308 return RetainResult(stored=False, error="Provider unavailable (degraded mode)")
310 def handle_degraded_recall(self, provider_name: str) -> RecallResult:
311 return self._degraded_handler.handle_recall(provider_name)
313 # -- Input validation --
315 def validate_retain_input(self, content: str, tags: list[str] | None, bank_id: str | None = None) -> str | None:
316 """Validate retain input sizes under ``bank_id``'s homeostasis. Returns error string or None if valid."""
317 max_content_bytes = self._for(bank_id).homeostasis.retain_max_content_bytes
318 if max_content_bytes:
319 size = len(content.encode("utf-8"))
320 if size > max_content_bytes:
321 return f"Content exceeds maximum size ({size} > {max_content_bytes} bytes)"
322 if tags and len(tags) > 100:
323 return f"Too many tags ({len(tags)} > 100)"
324 return None