stonks-oracle/services/intelligence_pipeline_v3/canary/routing.py

"""Canary compatibility outputs — percentage routing and automatic rollback.

Enables v3 adapter outputs for non-trading consumers first, then
progressively routes more traffic. Automatic rollback triggers on
correctness, latency, queue, or availability thresholds. Rollback
preserves v3 audit records.
"""

from __future__ import annotations

import enum
import hashlib
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any
from uuid import UUID, uuid4


class RollbackReason(str, enum.Enum):
    """Reasons for automatic canary rollback."""

    CORRECTNESS_THRESHOLD = "correctness_threshold"
    LATENCY_THRESHOLD = "latency_threshold"
    QUEUE_SATURATION = "queue_saturation"
    AVAILABILITY_THRESHOLD = "availability_threshold"
    ERROR_RATE = "error_rate"
    MANUAL = "manual"


@dataclass(frozen=True)
class RollbackEvent:
    """Immutable record of a canary rollback.

    Rollback leaves v3 audit records intact — only routing changes.
    """

    event_id: UUID
    timestamp: datetime
    reason: RollbackReason
    previous_percentage: int
    metric_value: float
    threshold_value: float
    details: str = ""

    @classmethod
    def create(
        cls,
        reason: RollbackReason,
        previous_percentage: int,
        metric_value: float,
        threshold_value: float,
        details: str = "",
    ) -> RollbackEvent:
        return cls(
            event_id=uuid4(),
            timestamp=datetime.now(timezone.utc),
            reason=reason,
            previous_percentage=previous_percentage,
            metric_value=metric_value,
            threshold_value=threshold_value,
            details=details,
        )


@dataclass
class CanaryConfig:
    """Canary routing configuration with thresholds."""

    enabled: bool = False
    percentage: int = 0  # 0-100, percentage of docs using v3 outputs
    document_types: set[str] = field(default_factory=set)  # Types eligible for canary
    exclude_trading: bool = True  # Exclude trading consumers initially

    # Automatic rollback thresholds
    max_error_rate: float = 0.05
    max_p95_latency_ms: float = 5000.0
    max_queue_saturation: float = 0.90
    min_availability: float = 0.95
    min_correctness: float = 0.90

    # Rollback behavior
    rollback_to_percentage: int = 0  # Roll back to this percentage
    cooldown_minutes: int = 60  # Wait before re-enabling after rollback


@dataclass
class CanaryRouter:
    """Routes documents between v2 and v3 outputs at configurable percentages.

    Routing is deterministic per document_id to avoid inconsistent
    behavior on retries. Rollback preserves all v3 audit records.
    """

    config: CanaryConfig
    _rollback_events: list[RollbackEvent] = field(default_factory=list)
    _documents_routed_v3: int = 0
    _documents_routed_v2: int = 0
    _last_rollback: datetime | None = None

    def should_use_v3(
        self,
        document_id: str,
        document_type: str | None = None,
        is_trading_consumer: bool = False,
    ) -> bool:
        """Determine if a document should use v3 outputs.

        Deterministic per document_id for consistency.
        """
        if not self.config.enabled:
            return False

        # Respect trading exclusion
        if is_trading_consumer and self.config.exclude_trading:
            return False

        # Check if in cooldown after rollback
        if self._in_cooldown():
            return False

        # Document type filter
        if (
            self.config.document_types
            and document_type
            and document_type not in self.config.document_types
        ):
            return False

        # Percentage-based routing (deterministic hash)
        bucket = self._hash_to_bucket(document_id)
        use_v3 = bucket < self.config.percentage

        if use_v3:
            self._documents_routed_v3 += 1
        else:
            self._documents_routed_v2 += 1

        return use_v3

    def check_rollback(
        self,
        error_rate: float = 0.0,
        p95_latency_ms: float = 0.0,
        queue_saturation: float = 0.0,
        availability: float = 1.0,
        correctness: float = 1.0,
    ) -> RollbackEvent | None:
        """Check all rollback thresholds. Returns event if rollback triggered."""
        if not self.config.enabled or self.config.percentage == 0:
            return None

        checks: list[tuple[RollbackReason, float, float, str]] = [
            (
                RollbackReason.ERROR_RATE,
                error_rate,
                self.config.max_error_rate,
                f"Error rate {error_rate:.3f} > {self.config.max_error_rate}",
            ),
            (
                RollbackReason.LATENCY_THRESHOLD,
                p95_latency_ms,
                self.config.max_p95_latency_ms,
                f"P95 latency {p95_latency_ms:.0f}ms > {self.config.max_p95_latency_ms:.0f}ms",
            ),
            (
                RollbackReason.QUEUE_SATURATION,
                queue_saturation,
                self.config.max_queue_saturation,
                f"Queue saturation {queue_saturation:.2f} > {self.config.max_queue_saturation}",
            ),
        ]

        for reason, value, threshold, details in checks:
            if value > threshold:
                return self._trigger_rollback(reason, value, threshold, details)

        # These check for below threshold
        if availability < self.config.min_availability:
            return self._trigger_rollback(
                RollbackReason.AVAILABILITY_THRESHOLD,
                availability,
                self.config.min_availability,
                f"Availability {availability:.3f} < {self.config.min_availability}",
            )

        if correctness < self.config.min_correctness:
            return self._trigger_rollback(
                RollbackReason.CORRECTNESS_THRESHOLD,
                correctness,
                self.config.min_correctness,
                f"Correctness {correctness:.3f} < {self.config.min_correctness}",
            )

        return None

    def manual_rollback(self, details: str = "") -> RollbackEvent:
        """Trigger a manual rollback."""
        return self._trigger_rollback(
            RollbackReason.MANUAL,
            0.0,
            0.0,
            details or "Manual rollback requested",
        )

    def _trigger_rollback(
        self,
        reason: RollbackReason,
        metric_value: float,
        threshold_value: float,
        details: str,
    ) -> RollbackEvent:
        """Execute rollback — change routing but preserve audit data."""
        event = RollbackEvent.create(
            reason=reason,
            previous_percentage=self.config.percentage,
            metric_value=metric_value,
            threshold_value=threshold_value,
            details=details,
        )
        self.config.percentage = self.config.rollback_to_percentage
        self._rollback_events.append(event)
        self._last_rollback = datetime.now(timezone.utc)
        return event

    def _in_cooldown(self) -> bool:
        """Check if we're in cooldown after a rollback."""
        if self._last_rollback is None:
            return False
        from datetime import timedelta

        cooldown_end = self._last_rollback + timedelta(
            minutes=self.config.cooldown_minutes
        )
        return datetime.now(timezone.utc) < cooldown_end

    def _hash_to_bucket(self, document_id: str) -> int:
        """Deterministic hash to 0-99 bucket."""
        h = hashlib.sha256(f"canary:{document_id}".encode()).hexdigest()
        return int(h[:8], 16) % 100

    @property
    def rollback_events(self) -> list[RollbackEvent]:
        return list(self._rollback_events)

    @property
    def v3_traffic_ratio(self) -> float:
        total = self._documents_routed_v2 + self._documents_routed_v3
        if total == 0:
            return 0.0
        return self._documents_routed_v3 / total

    def summary(self) -> dict[str, Any]:
        return {
            "enabled": self.config.enabled,
            "percentage": self.config.percentage,
            "documents_v3": self._documents_routed_v3,
            "documents_v2": self._documents_routed_v2,
            "rollback_count": len(self._rollback_events),
            "in_cooldown": self._in_cooldown(),
        }