from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass from datetime import UTC, datetime from typing import Any, Protocol from tenant_engine.domain import ( CapabilityRole, PlanAssignment, RoleGrant, Tenant, TenantLifecycle, TenantRetiredError, ) class TenantNotFoundError(KeyError): pass class VersionConflictError(RuntimeError): """The caller's `If-Match` version is not the tenant's current version.""" def __init__(self, *, expected: int, actual: int) -> None: super().__init__(f"expected version {expected}, current version is {actual}") self.expected = expected self.actual = actual class IdempotencyConflictError(RuntimeError): """An `Idempotency-Key` was reused for a materially different request.""" class TenantAlreadyExistsError(ValueError): pass class GrantNotFoundError(KeyError): pass class StoreUnavailableError(RuntimeError): """The store could not answer -- callers on a privileged decision path (flex-auth's live lookup) must treat this as deny, never as "zero roles". """ @dataclass(frozen=True, slots=True) class DomainEvent: """Boundary contract's Audit Correlation Contract, in event form.""" event_type: str tenant_id: str at: datetime payload: dict[str, Any] class TenantStore(Protocol): """Swappable persistence seam -- domain/ and api/ depend on this, not a backend. Every method taking a `tenant_id` accepts either the tenant's internal id (caller-chosen at creation, e.g. from an admin tool) or its profile identifier (e.g. "tenant:friendly:binky") -- external callers like key-cape and flex-auth only ever have the identifier (it's the IAM Profile `tenant` claim value), never tenant-engine's internal id. See `_resolve_tenant_id` for why this had to be added after real cross-service testing caught the gap (found integrating key-cape's tenant_roles claim, KEY-WP-0005-T02). """ def create_tenant(self, tenant: Tenant) -> None: ... def get_tenant(self, tenant_id: str) -> Tenant: ... def grant_role(self, grant: RoleGrant) -> None: ... def revoke_role(self, *, tenant_id: str, grant_id: str, at: datetime) -> RoleGrant: ... def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]: ... def assign_plan(self, assignment: PlanAssignment) -> None: ... def events(self) -> list[DomainEvent]: ... def mutate_tenant( self, *, tenant_id: str, expected_version: int, mutate: Callable[[Tenant], Tenant], event_type: str, evidence: dict[str, Any], idempotency_key: str, request_fingerprint: str, ) -> tuple[Tenant, bool]: """Atomically compare-and-swap a tenant record. Returns (tenant, replayed). One method carries all four concerns -- idempotency replay, version CAS, the mutation itself, and the audit event -- because they have to commit or fail together. Splitting them across store calls would leave a window where a crash yields a bumped version with no receipt (a retry then double-applies) or a receipt with no mutation. Order matters: a replayed `idempotency_key` short-circuits *before* the version check, because a genuine retry of an already-applied mutation necessarily carries a now-stale `If-Match`. """ ... class InMemoryTenantStore: def __init__(self) -> None: self._tenants: dict[str, Tenant] = {} self._by_identifier: dict[str, str] = {} self._grants: dict[str, dict[str, RoleGrant]] = {} self._plans: dict[str, PlanAssignment] = {} self._events: list[DomainEvent] = [] # (tenant_id, idempotency_key) -> (request_fingerprint, result snapshot) self._receipts: dict[tuple[str, str], tuple[str, Tenant]] = {} def create_tenant(self, tenant: Tenant) -> None: if tenant.tenant_id in self._tenants: raise TenantAlreadyExistsError(tenant.tenant_id) if tenant.identifier in self._by_identifier: raise TenantAlreadyExistsError(tenant.identifier) self._tenants[tenant.tenant_id] = tenant self._by_identifier[tenant.identifier] = tenant.tenant_id self._grants[tenant.tenant_id] = {} self._emit( "tenant_created", tenant.tenant_id, {"identifier": tenant.identifier, "grouping": tenant.grouping}, ) def get_tenant(self, tenant_id: str) -> Tenant: return self._tenants[self._resolve(tenant_id)] def grant_role(self, grant: RoleGrant) -> None: resolved = self._resolve(grant.tenant_id) self._require_active(resolved, "grant a role") self._grants[resolved][grant.grant_id] = grant self._emit( "role_granted", resolved, { "grant_id": grant.grant_id, "role": grant.role.value, "grant_reason": grant.grant_reason, "correlation_id": grant.correlation_id, }, ) def revoke_role(self, *, tenant_id: str, grant_id: str, at: datetime) -> RoleGrant: resolved = self._resolve(tenant_id) try: grant = self._grants[resolved][grant_id] except KeyError: raise GrantNotFoundError(grant_id) from None revoked = grant.revoke(at=at) self._grants[resolved][grant_id] = revoked self._emit( "role_revoked", resolved, {"grant_id": grant_id, "role": revoked.role.value}, ) return revoked def active_roles(self, tenant_id: str) -> frozenset[CapabilityRole]: resolved = self._resolve(tenant_id) return frozenset( grant.role for grant in self._grants.get(resolved, {}).values() if grant.active ) def assign_plan(self, assignment: PlanAssignment) -> None: resolved = self._resolve(assignment.tenant_id) self._require_active(resolved, "assign a plan") self._plans[resolved] = assignment self._emit("plan_assigned", resolved, {"plan_id": assignment.plan_id}) def events(self) -> list[DomainEvent]: return list(self._events) def mutate_tenant( self, *, tenant_id: str, expected_version: int, mutate: Callable[[Tenant], Tenant], event_type: str, evidence: dict[str, Any], idempotency_key: str, request_fingerprint: str, ) -> tuple[Tenant, bool]: resolved = self._resolve(tenant_id) receipt = self._receipts.get((resolved, idempotency_key)) if receipt is not None: fingerprint, snapshot = receipt if fingerprint != request_fingerprint: raise IdempotencyConflictError(idempotency_key) return snapshot, True current = self._tenants[resolved] if current.version != expected_version: raise VersionConflictError(expected=expected_version, actual=current.version) updated = mutate(current) self._tenants[resolved] = updated self._receipts[(resolved, idempotency_key)] = (request_fingerprint, updated) self._emit(event_type, resolved, {**evidence, "version": updated.version}) return updated, False def _require_active(self, resolved_id: str, what: str) -> None: tenant = self._tenants[resolved_id] if tenant.lifecycle is not TenantLifecycle.ACTIVE: raise TenantRetiredError(f"cannot {what} on a retired tenant") def _resolve(self, tenant_id: str) -> str: """Resolve an internal tenant_id or a profile identifier to the canonical internal tenant_id every other private dict is keyed by. External callers (key-cape, flex-auth) only ever have the identifier (the IAM Profile `tenant` claim value) -- they have no way to know a tenant's internal tenant_id, which is caller-chosen at creation time and otherwise opaque. Found the hard way: a real cross-service check (key-cape's KEY-WP-0005-T02 against a live tenant-engine) returned tenant_not_found for a tenant that genuinely existed, because the caller only had the identifier. """ resolved = self._by_identifier.get(tenant_id, tenant_id) if resolved not in self._tenants: raise TenantNotFoundError(tenant_id) return resolved def _emit(self, event_type: str, tenant_id: str, payload: dict[str, Any]) -> None: self._events.append( DomainEvent(event_type=event_type, tenant_id=tenant_id, at=datetime.now(UTC), payload=payload) )