server/drift_injector.py

"""
drift_injector.py — Manages world state mutations for AdaptiveWorld.

DRIFT_CONFIGS holds one entry per scenario, defining:
  - initial: DYNAMIC_CONFIG values at episode start
  - mutated: DYNAMIC_CONFIG values after drift injection
  - drift_type: the category of drift (for belief scoring)

The DriftInjector is instantiated per episode and called by
AdaptiveWorldEnvironment at drift_trigger_step.
"""
from typing import Dict, Any


DRIFT_CONFIGS: Dict[str, Dict] = {

    # ── Easy: E1 — Field rename (e-commerce order) ─────────────────────────────
    "easy_field_rename": {
        "initial": {
            "order_field": "qty",
            "required_extra": None,
            "order_status": "confirmed",
            "noisy_errors": False,
        },
        "mutated": {
            "order_field": "quantity",
            "required_extra": "customer_id",
            "order_status": "confirmed",
            "noisy_errors": False,
        },
        "drift_type": "field_rename",
    },

    # ── Easy: E2 — Endpoint version (hotel booking) ────────────────────────────
    "easy_endpoint_version": {
        "initial": {
            "rooms_endpoint": "/mock_api/rooms/book",
            "max_rooms_per_request": 0,
            "bulk_booking_allowed": True,
        },
        "mutated": {
            "rooms_endpoint": "/mock_api/v2/rooms/book",
            "max_rooms_per_request": 0,
            "bulk_booking_allowed": True,
        },
        "drift_type": "endpoint_version",
    },

    # ── Easy: E3 — Auth scheme (flight search) ─────────────────────────────────
    "easy_auth_format": {
        "initial": {
            "auth_scheme": "Bearer",
        },
        "mutated": {
            "auth_scheme": "ApiKey",
        },
        "drift_type": "policy_change",
    },

    # ── Medium: M1 — Double drift (insurance claim) ────────────────────────────
    "medium_double_drift": {
        "initial": {
            "claims_endpoint": "/mock_api/claims",
            "claims_id_field": "policy_id",
            "claims_amount_field": "amount",
            "claims_require_date": False,
        },
        "mutated": {
            "claims_endpoint": "/mock_api/claims/v2",
            "claims_id_field": "policy_number",
            "claims_amount_field": "claim_amount",
            "claims_require_date": True,
        },
        "drift_type": "field_rename",
    },

    # ── Medium: M2 — Policy flip (e-commerce discount) ─────────────────────────
    "medium_policy_flip": {
        "initial": {
            "discount_requires_membership": False,
        },
        "mutated": {
            "discount_requires_membership": True,
        },
        "drift_type": "policy_change",
    },

    # ── Medium: M3 — Rate limit rule change (hotel conference) ─────────────────
    "medium_rate_limit": {
        "initial": {
            "rooms_endpoint": "/mock_api/rooms/book",
            "max_rooms_per_request": 0,
            "bulk_booking_allowed": True,
        },
        "mutated": {
            "rooms_endpoint": "/mock_api/rooms/book",
            "max_rooms_per_request": 1,
            "bulk_booking_allowed": False,
        },
        "drift_type": "policy_change",
    },

    # ── Hard: H1 — Silent semantic drift (order status) ────────────────────────
    "hard_status_meaning": {
        "initial": {
            "order_field": "qty",
            "required_extra": None,
            "order_status": "confirmed",
            "noisy_errors": False,
        },
        "mutated": {
            "order_field": "qty",
            "required_extra": None,
            "order_status": "approved",    # API returns 200 OK but meaning changed
            "noisy_errors": False,
        },
        "drift_type": "silent_semantic",
    },

    # ── Hard: H2 — Response structure drift (product search) ───────────────────
    "hard_response_structure": {
        "initial": {
            "product_key": "products",
            "product_id_field": "id",
            "product_price_field": "price",
        },
        "mutated": {
            "product_key": "items",
            "product_id_field": "product_id",
            "product_price_field": "cost",
        },
        "drift_type": "silent_semantic",
    },

    # ── Hard: H3 — Cascading invalidation (category field) ─────────────────────
    "hard_cascading": {
        "initial": {
            "product_category": "electronics",
            "discount_requires_membership": False,
        },
        "mutated": {
            "product_category": "tech",       # category change
            "discount_requires_membership": True,  # coupon eligibility invalidated
        },
        "drift_type": "silent_semantic",
    },

    # ── Expert: EX1 — Transient error then real drift ─────────────────────────
    # Note: transient 503 is handled via secondary_drift_step logic in environment
    "expert_transient_vs_real": {
        "initial": {
            "order_field": "qty",
            "required_extra": None,
            "order_status": "confirmed",
            "noisy_errors": False,
        },
        "mutated": {
            "order_field": "quantity",
            "required_extra": None,
            "order_status": "confirmed",
            "noisy_errors": False,
        },
        "drift_type": "field_rename",
        "transient_error_step": 2,   # step 2 returns 503 (not a real drift)
    },

    # ── Expert: EX2 — Cross-service drift ─────────────────────────────────────
    "expert_cross_service": {
        "initial": {
            "auth_scheme": "Bearer",
            "rooms_endpoint": "/mock_api/rooms/book",
        },
        "mutated": {
            "auth_scheme": "ApiKey",               # flight drift at step 3
            "rooms_endpoint": "/mock_api/rooms/book",
        },
        "secondary_mutated": {
            "auth_scheme": "ApiKey",
            "rooms_endpoint": "/mock_api/v2/rooms/book",  # hotel drift at step 6
        },
        "drift_type": "field_rename",
        "secondary_drift_step": 6,
    },

    # ── Expert: EX3 — Deprecation red herring ─────────────────────────────────
    "expert_red_herring": {
        "initial": {
            "claims_endpoint": "/mock_api/claims",
            "claims_id_field": "policy_id",
            "claims_amount_field": "amount",
            "claims_require_date": False,
        },
        "mutated": {
            "claims_endpoint": "/mock_api/claims",    # warning added but still same endpoint
            "claims_id_field": "policy_id",
            "claims_amount_field": "amount",
            "claims_require_date": False,
            "_deprecated_notice": "field X will be removed in v3",  # red herring
        },
        "final_mutated": {
            "claims_endpoint": "/mock_api/claims/v2",  # real drift happens at step 6
            "claims_id_field": "policy_number",
            "claims_amount_field": "amount",
            "claims_require_date": False,
        },
        "drift_type": "endpoint_version",
        "secondary_drift_step": 6,
    },
}


class DriftInjector:
    """
    Manages world state for a single episode.
    Instantiated per episode by AdaptiveWorldEnvironment.
    """

    def __init__(self, scenario_id: str):
        if scenario_id not in DRIFT_CONFIGS:
            raise ValueError(f"Unknown scenario: {scenario_id}. "
                             f"Available: {list(DRIFT_CONFIGS.keys())}")
        self.scenario_id = scenario_id
        self.config = DRIFT_CONFIGS[scenario_id]
        self.current_world: dict = dict(self.config["initial"])
        self.drifted: bool = False
        self.secondary_drifted: bool = False

    def inject(self) -> Dict[str, Any]:
        """
        Inject primary drift. Returns the mutation dict for the mock API.
        Should be passed to /_admin/mutate.
        """
        self.current_world = dict(self.config["mutated"])
        self.drifted = True
        return self.current_world

    def inject_secondary(self) -> Dict[str, Any]:
        """
        Inject secondary drift (expert scenarios with two drift events).
        Returns mutation dict, or empty dict if no secondary drift defined.
        """
        key = "secondary_mutated" if "secondary_mutated" in self.config else "final_mutated"
        if key not in self.config:
            return {}
        self.current_world = dict(self.config[key])
        self.secondary_drifted = True
        return self.current_world

    def get_initial(self) -> Dict[str, Any]:
        return dict(self.config["initial"])

    def get_world(self) -> Dict[str, Any]:
        return self.current_world

    def get_drift_type(self) -> str:
        return self.config["drift_type"]

    def get_world_truth(self) -> Dict[str, Any]:
        """
        Ground truth for belief scoring.
        Uses secondary/final mutated world for expert scenarios.
        """
        for key in ("final_mutated", "secondary_mutated", "mutated"):
            if key in self.config:
                return dict(self.config[key])
        return dict(self.config["mutated"])

    def get_transient_error_step(self) -> int:
        """Returns the step at which a transient 503 should be simulated (EX1)."""
        return self.config.get("transient_error_step", -1)

    def get_secondary_drift_step(self) -> int:
        """Returns the step for secondary drift injection (EX2, EX3)."""
        return self.config.get("secondary_drift_step", -1)