server-26/drb-c2-core/app/internal/feature_flags.py

"""
Global AI feature flags stored in Firestore at config/ai_features.

Defaults to all-on when the document does not exist yet.  Uses a short
in-memory TTL cache so flag reads don't add a Firestore round-trip to every
call upload.
"""
import time
from contextvars import ContextVar
from typing import Any, Optional
from app.internal.logger import logger
from app.internal import firestore as fstore

_COLLECTION = "config"
_DOC_ID     = "ai_features"
_TTL        = 30.0  # seconds before re-reading from Firestore

_DEFAULTS: dict[str, bool] = {
    "stt_enabled":                  True,
    "correlation_enabled":          True,
    "summaries_enabled":            True,
    "vocabulary_learning_enabled":  True,
    # Transcript correction runs inside transcribe_call and spends Gemini
    # tokens plus Places quota on every transcribed call.  Until server-26#76
    # it was reachable only through an env var and an ansible run, which meant
    # an "STT-only" evaluation window was never STT-only and its cost could
    # not be attributed (server-26#45).
    #
    # NOT a pure cost lever.  The corrector is also the noise gate: it is what
    # sets not_speech, and transcription.py returns nothing for a call it
    # flags.  _is_degenerate does not catch what the corrector catches, so with
    # this off, recogniser noise reaches extraction as a real transcript, comes
    # back with no units/tags/location, is judged thin, and auto-attaches to the
    # most recent incident on the talkgroup with no fit check.  Turning this off
    # while correlation_enabled is on therefore pushes over-merging -- do not do
    # it during an evaluation window.
    "transcript_correction_enabled": True,
}

# A replay run (app/internal/replay.py) states exactly which AI steps it runs,
# independent of the live switches — the whole point is re-running the pipeline
# while live AI is OFF. ContextVar so the override never reaches a live upload.
_forced: ContextVar[Optional[dict[str, bool]]] = ContextVar("drb_forced_flags", default=None)


def force_flags(flags: Optional[dict[str, bool]]):
    """Override resolve_flags() for the current context. Returns a reset token."""
    return _forced.set(flags)


def unforce_flags(token) -> None:
    _forced.reset(token)


_cache: dict[str, Any] = {}
_cache_ts: float = 0.0


async def get_flags() -> dict[str, bool]:
    """Return the current feature flags, using the TTL cache when fresh."""
    global _cache, _cache_ts

    now = time.monotonic()
    if _cache and (now - _cache_ts) < _TTL:
        return dict(_cache)

    try:
        doc = await fstore.doc_get(_COLLECTION, _DOC_ID)
        if doc:
            merged = {**_DEFAULTS, **{k: bool(v) for k, v in doc.items() if k in _DEFAULTS}}
        else:
            merged = dict(_DEFAULTS)
    except Exception as e:
        logger.warning(f"Feature flags: could not read from Firestore ({e}), using defaults")
        merged = dict(_DEFAULTS)

    _cache    = merged
    _cache_ts = now
    return dict(_cache)


async def _cascade_to_systems(clean: dict[str, bool]) -> tuple[list[dict], list[dict]]:
    """Clear per-system ``ai_flags`` overrides for the keys just set globally.

    Returns ``(changes, errors)``.

    Why clearing rather than overwriting with the new value: an override that
    stays present, merely agreeing with the global switch for now, defeats the
    NEXT flip exactly the same way. Removing it makes the system inherit, which
    is the same semantics the human-facing route already offers
    (``PUT /systems/{id}/ai-flags`` with null → "clear override, inherit
    global").

    Systems are discovered by scanning for documents that actually carry an
    ``ai_flags`` map — never a hardcoded id list. Two systems carry overrides
    today; a third added tomorrow would silently defeat a global shutoff if
    this were pinned to the current pair.
    """
    changes: list[dict] = []
    errors:  list[dict] = []

    systems = await fstore.collection_list("systems")
    for system in systems:
        sid = system.get("system_id")
        ai_flags = system.get("ai_flags")
        # Only documents that actually carry the map. A system with no
        # overrides already inherits, so there is nothing to cascade to.
        if not sid or not isinstance(ai_flags, dict) or not ai_flags:
            continue
        removed = {k: ai_flags[k] for k in clean if k in ai_flags}
        if not removed:
            continue
        remaining = {k: v for k, v in ai_flags.items() if k not in clean}
        try:
            await fstore.doc_update("systems", sid, {"ai_flags": remaining})
        except Exception as e:
            # Report rather than swallow: a half-applied cascade is the exact
            # failure mode this helper exists to prevent, so it must be visible
            # in the log and the audit entry.
            logger.error(f"Feature flags: cascade to system '{sid}' failed ({e})")
            errors.append({"system_id": sid, "error": str(e)})
            continue
        changes.append({
            "system_id": sid,
            "cleared_overrides": removed,
            "now_inherits": {k: clean[k] for k in removed},
        })

    return changes, errors


async def set_flags(
    updates: dict[str, bool],
    actor: tuple[str, str] | None = None,
    cascade: bool = False,
) -> dict[str, bool]:
    """Write flag updates to Firestore, invalidate the cache, and audit it.

    ``actor`` is ``(actor_uid, actor_email)`` — see auth.describe_actor. It is
    optional so existing callers keep working; an unattributed flip is logged
    as "unknown" rather than not logged at all.

    ``cascade`` also clears the matching per-system ``ai_flags`` overrides, so
    one call is a total flip. Defaults to False deliberately — see the route's
    comment in routers/admin.py.

    Returns the resulting global flags dict, unchanged in shape: the admin UI
    (drb-frontend/lib/c2api.ts setFeatureFlags) types the response as
    Record<string, boolean>, so cascade/audit detail goes to the log and the
    audit entry rather than into this payload.
    """
    global _cache_ts

    clean = {k: bool(v) for k, v in updates.items() if k in _DEFAULTS}
    if not clean:
        raise ValueError(f"No recognised flag keys in update: {list(updates)}")

    # Force a fresh read for the "before" side of the audit entry: the TTL
    # cache can be up to _TTL seconds stale, and a wrong previous value in an
    # audit log is worse than none.
    _cache_ts = 0.0
    before = await get_flags()

    await fstore.doc_set(_COLLECTION, _DOC_ID, clean)
    _cache_ts = 0.0  # force re-read on next get_flags()
    logger.info(f"Feature flags updated: {clean}")

    cascaded: list[dict] = []
    cascade_errors: list[dict] = []
    if cascade:
        cascaded, cascade_errors = await _cascade_to_systems(clean)
        logger.info(
            f"Feature flags: cascaded {list(clean)} to {len(cascaded)} system(s), "
            f"{len(cascade_errors)} error(s)"
        )

    after = await get_flags()

    # The audit entry is a record OF the write, never a precondition for it.
    # audit_log lives in the same Firestore that just accepted the flag write,
    # so a failure here is nearly always transient — losing the flip (or 500ing
    # a route that already succeeded, which invites a retry that flips it back)
    # would be a far worse outcome than an unrecorded flip that is still in the
    # service log above.
    try:
        # Deferred import: app.internal.audit pulls in firestore, and this
        # module is imported from router module scope.
        from app.internal import audit
        actor_uid, actor_email = actor or ("unknown", "")
        changed = {
            k: {"from": before.get(k), "to": after.get(k)}
            for k in clean
            if before.get(k) != after.get(k)
        }
        await audit.write_audit(
            actor_uid=actor_uid,
            actor_email=actor_email,
            action="feature_flags.update",
            details={
                "requested": clean,
                "changed": changed,
                "before": before,
                "after": after,
                "cascade": cascade,
                "cascaded_systems": cascaded,
                "cascade_errors": cascade_errors,
            },
        )
    except Exception as e:
        logger.error(f"Feature flags: audit write failed ({e}) — flag change stands")

    return after


async def resolve_flags(system_id: str | None):
    """
    Resolve the AI feature flags for one radio system.

    Returns ``(flags, flag)``: ``flags`` is the raw global config/ai_features
    document, and ``flag(name)`` layers the system's own ``ai_flags`` on top of
    it.  A system flag of False beats a global True, but a global False beats
    everything -- config/ai_features is the master switch, which is the whole
    point of having one (server-26#75, server-26#76).

    Every AI spend path resolves through here.  A path that reads ``flags``
    directly re-introduces #75; a path that reads neither re-introduces #76.
    """
    from app.internal import firestore as _fstore

    forced = _forced.get()
    if forced is not None:
        full = {k: bool(forced.get(k, False)) for k in _DEFAULTS}
        return full, lambda name: full.get(name, False)

    flags = await get_flags()

    system_ai_flags: dict = {}
    if system_id:
        sys_doc = await _fstore.doc_get_cached("systems", system_id)
        system_ai_flags = (sys_doc or {}).get("ai_flags") or {}

    def flag(name: str) -> bool:
        if not flags[name]:                      # global master off
            return False
        return system_ai_flags.get(name, True)   # system override, else inherit

    return flags, flag