""" RepairEnv V2 — Main OpenEnv-compliant environment class. HARD CONSTRAINTS ENFORCED: - Hidden state: Observations only contain symptoms (logs + metrics). - Strict step order: Action → Cascade → Degrade → Recovery Ticks → Symptoms → Reward - Delayed recovery: repair_root_cause takes 2 steps to fully heal. - Inspection cost: 3 free, then -0.1 reward per inspect. """ from __future__ import annotations import copy import random from typing import Any, Dict, List, Optional, Tuple from models_v2 import ( Action, ActionType, ComponentMetrics, ComponentStatus, ComponentTarget, EpisodeResultV2, InspectResult, SymptomObservation, Reward, ACTION_COSTS ) from symptom_engine import SymptomEngine # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- SEED = 42 STEP_PENALTY = -0.2 INSPECT_PENALTY = -0.1 MAX_FREE_INSPECTIONS = 3 COMPONENTS = ["database", "api", "cache", "queue"] # --------------------------------------------------------------------------- # RepairEnvV2 # --------------------------------------------------------------------------- class RepairEnvV2: def __init__(self, seed: int = SEED) -> None: """OpenEnv-compatible constructor. Accepts optional seed (default 42).""" self._init_seed = seed self._task_config: Optional[Dict[str, Any]] = None # Hidden State self._true_state: Dict[str, Dict] = {} self._root_causes: List[str] = [] self._recovery_cooldowns: Dict[str, int] = {} self._step_count: int = 0 self._max_steps: int = 0 self._task_id: str = "" self._done: bool = False self._total_reward: float = 0.0 self._inspections_used: int = 0 self._fixed_symptom_before_root: bool = False self._root_cause_fixed: bool = False self._rng: random.Random = random.Random(seed) self._symptom_engine: SymptomEngine = SymptomEngine(seed) # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def reset( self, task_config: Optional[Dict[str, Any]] = None, task_id: str = "easy", ) -> SymptomObservation: """ Reset the environment. OpenEnv calls this with no arguments → defaults to 'easy' task. server_v2.py calls it with (task_config=config, task_id=task_id) → still works. """ # If no task_config provided, load from tasks_v2 by task_id if task_config is None: from tasks_v2 import TASKS if task_id not in TASKS: task_id = "easy" task_config = TASKS[task_id] self._task_config = copy.deepcopy(task_config) self._task_id = task_id self._max_steps = self._task_config["max_steps"] self._step_count = 0 self._done = False self._total_reward = 0.0 self._inspections_used = 0 self._fixed_symptom_before_root = False self._root_causes_fixed: Dict[str, bool] = {} # Initialize Hidden State root = self._task_config["root_cause"] self._root_causes = root if isinstance(root, list) else [root] self._root_causes_fixed = {r: False for r in self._root_causes} self._true_state = copy.deepcopy(self._task_config["initial_state"]) self._recovery_cooldowns = {c: 0 for c in COMPONENTS} # Ensures rng is deterministic per run seed = self._task_config.get("seed", SEED) self._rng.seed(seed) self._symptom_engine = SymptomEngine(seed) return self._generate_symptoms() def step(self, action: Action) -> Tuple[SymptomObservation, Reward, bool, Dict[str, Any]]: """ Hard Rule #4: Strict Step Order 1. Action 2. Cascading 3. Degradation 4. Recovery Ticks 5. Symptoms (Metrics + Logs) 6. Reward computation """ if self._done: raise RuntimeError("Episode is done. Call reset().") # Snapshot for reward computation old_health = sum(1 for c in self._true_state.values() if c["status"] == "healthy") # 1. Action action_reward_modifier, inspect_info = self._apply_action(action) # 2. Cascading (Side effects driven by root cause) self._apply_cascading_effects() # 3. Degradation (Passive worsening) self._apply_degradation() # 4. Recovery Ticks self._tick_recovery_cooldowns(action) # 5. Generate Symptoms obs = self._generate_symptoms() self._step_count += 1 # Check terminal condition all_healthy = all(c["status"] == "healthy" for c in self._true_state.values()) if all_healthy or self._step_count >= self._max_steps: self._done = True # 6. Compute Reward new_health = sum(1 for c in self._true_state.values() if c["status"] == "healthy") reward = self._compute_reward( action, old_health, new_health, all_healthy, action_reward_modifier ) self._total_reward = round(self._total_reward + reward.value, 4) # Info dict (transparency for grading, completely hidden from agent observation) info = { "step": self._step_count, "true_state_snapshot": copy.deepcopy(self._true_state), "root_causes": self._root_causes, "cooldowns": self._recovery_cooldowns.copy(), "inspect_result": inspect_info.model_dump() if inspect_info else None, "reward_reason": reward.reason } return obs, reward, self._done, info def get_episode_result(self) -> EpisodeResultV2: final_health = sum(1 for c in self._true_state.values() if c["status"] == "healthy") health_ratio = final_health / 4.0 steps = self._step_count optimal = self._task_config.get("optimal_steps", 5) efficiency = max(0.0, min(1.0, 1.0 - (steps - optimal) * 0.05)) all_roots_fixed = all(self._root_causes_fixed.values()) correct_root_fix_ratio = 1.0 if (all_roots_fixed and not self._fixed_symptom_before_root) else 0.0 score = (0.5 * health_ratio) + (0.3 * efficiency) + (0.2 * correct_root_fix_ratio) return EpisodeResultV2( task_id=self._task_id, steps_taken=steps, total_reward=self._total_reward, score=round(score, 4), success=(health_ratio == 1.0), final_health_ratio=health_ratio, efficiency=efficiency, correct_root_fix_ratio=correct_root_fix_ratio ) # ------------------------------------------------------------------ # Private Implementation # ------------------------------------------------------------------ def _apply_action(self, action: Action) -> Tuple[float, Optional[InspectResult]]: """Handles both repair and inspect actions. Returns (reward_mod, inspect_result).""" target = action.target atype = action.action_type # Apply strict action costs reward_modifier = ACTION_COSTS.get(atype, 0.0) inspect_res = None if atype.startswith("inspect_"): self._inspections_used += 1 if self._inspections_used > MAX_FREE_INSPECTIONS: reward_modifier += INSPECT_PENALTY investigating_root = (target in self._root_causes) # If it's the root cause, high confidence. if investigating_root: conf = round(self._rng.uniform(0.75, 1.0), 2) hint = f"Critical failures traced back to {target}. High correlation with active symptoms." is_likely = True else: # Symptom. Low confidence, but might randomly say True (false positive). conf = round(self._rng.uniform(0.2, 0.55), 2) is_likely = self._rng.random() < 0.2 # 20% false positive hint = f"Anomalies detected in {target}, but could be upstream cascade." if self._true_state[target]["status"] == "healthy": conf = round(self._rng.uniform(0.8, 1.0), 2) is_likely = False hint = f"{target} appears healthy. No immediate faults found." inspect_res = InspectResult( component=target, hint=hint, confidence=conf, is_likely_root_cause=is_likely ) elif atype == ActionType.repair_database and target == "database": # Delayed recovery logic. if self._true_state["database"]["status"] != "healthy": if "database" in self._root_causes: self._recovery_cooldowns["database"] = 2 else: self._recovery_cooldowns["database"] = 1 elif atype == ActionType.restart_service: # Short term symptom fix. Will be overwritten by cascade if root still broken. if self._true_state[target]["status"] == "down": self._true_state[target]["status"] = "degraded" elif self._true_state[target]["status"] == "degraded": self._true_state[target]["status"] = "healthy" active_roots = [r for r in self._root_causes if self._true_state[r]["status"] != "healthy"] if target not in self._root_causes and active_roots: # They are fixing a symptom before the root cause self._fixed_symptom_before_root = True if target in self._root_causes: self._recovery_cooldowns[target] = 1 # Restarts are slightly faster but less reliable than dedicated tools # Other actions omitted for brevity, mapping easily to V1 equivalents where needed. elif atype == ActionType.clear_cache and target == "cache": if self._true_state["cache"]["status"] != "healthy": self._true_state["cache"]["status"] = "healthy" return reward_modifier, inspect_res def _apply_cascading_effects(self): """Root-cause driven cascades. Overwrites temporary symptom fixes.""" if "database" in self._root_causes and self._true_state["database"]["status"] != "healthy": # DB brings down the API self._true_state["api"]["status"] = "degraded" self._true_state["queue"]["status"] = "degraded" if "cache" in self._root_causes and self._true_state["cache"]["status"] != "healthy": # Cache pressure ruins queue self._true_state["queue"]["status"] = "degraded" def _apply_degradation(self): """Passive worsening.""" pass # Currently handled implicitly by SymptomEngine metric generation, but placeholder for hard state drops def _tick_recovery_cooldowns(self, action: Action): for c in COMPONENTS: if self._recovery_cooldowns[c] > 0: # Don't tick the cooldown if we *just* started it this exact step. if action.target == c and action.action_type in [ActionType.repair_database, ActionType.restart_service]: continue self._recovery_cooldowns[c] -= 1 if self._recovery_cooldowns[c] == 0: self._true_state[c]["status"] = "healthy" if c in self._root_causes: self._root_causes_fixed[c] = True def _generate_symptoms(self) -> SymptomObservation: metrics = self._symptom_engine.generate_metrics(self._true_state) logs = self._symptom_engine.generate_logs(self._true_state, self._step_count) inspections = max(0, MAX_FREE_INSPECTIONS - self._inspections_used) return SymptomObservation( metrics=metrics, logs=logs, step_count=self._step_count, max_steps=self._max_steps, inspections_remaining=inspections ) def _compute_reward(self, action: Action, old_health: int, new_health: int, all_healthy: bool, modifier: float) -> Reward: base = STEP_PENALTY + modifier reason = [] if modifier < 0: reason.append(f"Inspection cost ({modifier})") if action.action_type == ActionType.no_op: return Reward(value=base - 0.5, reason="Useless actions: no_op") # Root cause fix detection is delayed by cooldown, so we check if cooldown was just set # (meaning the correct action was taken this turn) root_cooldown_started = ( action.target in self._root_causes and self._recovery_cooldowns[action.target] > 0 and action.action_type in [ActionType.repair_database, ActionType.restart_service, ActionType.clear_cache] ) if root_cooldown_started: base += 1.5 reason.append(f"Root cause addressed ({action.target})") elif new_health > old_health: base += 0.5 reason.append("Symptom improved") elif action.target not in self._root_causes and not action.action_type.startswith("inspect"): # They tried to fix a symptom that is immediately cascaded away (or just wrong) active_roots = [r for r in self._root_causes if self._true_state[r]["status"] != "healthy"] if active_roots: base -= 1.0 reason.append("Wrong component targeted (root cause still active)") else: base -= 0.5 reason.append("Component action had no net effect") if all_healthy: base += 5.0 reason.append("Terminal Bonus: All healthy") return Reward(value=round(base, 4), reason=" | ".join(reason) if reason else "Step tick")