""" IncidentOps — Environment Core ================================ Manages episode state, action execution, and step-level reward shaping. """ from server.models import ( IncidentAction, IncidentObservation, IncidentReward, ActionType, MetricSeries, DeployEvent, ) from server.scenario_engine import ScenarioEngine from typing import Dict, Any, Tuple import copy class IncidentOpsEnv: """ Core IncidentOps environment. Lifecycle: env = IncidentOpsEnv(task="alert-triage", seed=42) obs = env.reset() obs, reward, done, info = env.step(action) state = env.state() """ TASK_TEMPLATES = { "alert-triage": "mixed_alerts", "root-cause-analysis": "bad_deploy_memory_leak", "full-incident-response": "db_pool_exhaustion", } TASK_MAX_STEPS = { "alert-triage": 5, "root-cause-analysis": 12, "full-incident-response": 20, } def __init__(self, task: str, seed: int = 42): if task not in self.TASK_TEMPLATES: raise ValueError( f"Unknown task '{task}'. Valid: {list(self.TASK_TEMPLATES)}" ) self.task = task self.seed = seed self.max_steps = self.TASK_MAX_STEPS[task] self.scenario_engine = ScenarioEngine() template_id = self.TASK_TEMPLATES[task] self.scenario = self.scenario_engine.generate(template_id, seed) # Episode state self.step_count = 0 self.agent_triages: list = [] self.agent_hypothesis: str | None = None self.agent_fix: str | None = None self.state_data = self._build_initial_state() # ── Public API ──────────────────────────────────────────────────────────── def reset(self) -> IncidentObservation: """Reset the environment to a clean initial state (same task + seed).""" self.__init__(self.task, self.seed) return IncidentObservation(**self.state_data, step=0) def step( self, action: IncidentAction ) -> Tuple[IncidentObservation, IncidentReward, bool, Dict]: """ Execute one action and return (observation, reward, done, info). - Reward is always clamped to [0.0, 1.0]. - `done` becomes True when max_steps reached OR task-specific completion criteria are met (see _check_done). """ self.step_count += 1 result_msg = self._execute_action(action) action_desc = f"{action.action_type.value}: {json_safe(action.parameters)}" self.state_data["last_action_result"] = result_msg self.state_data["incident_timeline"].append( f"Step {self.step_count}: {action_desc} → {result_msg}" ) reward_value, breakdown = self._compute_reward(action) done = self._check_done() self.state_data["done"] = done obs = IncidentObservation(**self.state_data, step=self.step_count) reward = IncidentReward(value=reward_value, breakdown=breakdown, done=done) return obs, reward, done, {} def state(self) -> Dict[str, Any]: """Return the full current state (debugging / /state endpoint).""" return { **self.state_data, "step": self.step_count, "task": self.task, "max_steps": self.max_steps, } # ── Internal helpers ────────────────────────────────────────────────────── def _build_initial_state(self) -> Dict[str, Any]: """Strip grader-internal fields from alerts before exposing to agent.""" obs_alerts = [] for a in self.scenario["alerts"]: adict = copy.deepcopy(a) adict.pop("true_severity", None) adict.pop("true_team", None) obs_alerts.append(adict) metrics = self.scenario["metrics"] return { "alerts": obs_alerts, "recent_logs": self.scenario["logs"][:50], "metrics": {k: v.model_dump() for k, v in metrics.items()}, "deploy_history": self.scenario["deploy_history"], "incident_timeline": [], "resolution_status": "open", "runbook_draft": "", "comms_draft": "", "last_action_result": "Incident started. Review alerts and begin investigation.", "done": False, } def _check_done(self) -> bool: """Return True when episode should terminate.""" # Hard cap always fires if self.step_count >= self.max_steps: return True # Task-specific early-done criteria if self.task == "alert-triage": # Done when agent has triaged every real alert OR step cap real_alert_count = sum( 1 for a in self.scenario["alerts"] if "true_severity" in a ) if len(self.agent_triages) >= real_alert_count: return True elif self.task == "root-cause-analysis": # Done as soon as a hypothesis is submitted if self.agent_hypothesis: return True elif self.task == "full-incident-response": # Done when incident is explicitly closed if self.state_data["resolution_status"] == "closed": return True return False def _execute_action(self, action: IncidentAction) -> str: at = action.action_type p = action.parameters if at == ActionType.TRIAGE_ALERT: triage = { "alert_id": p.get("alert_id"), "severity": str(p.get("severity", "")), "team": str(p.get("team", "")), } # Avoid double-triaging the same alert already = [t for t in self.agent_triages if t["alert_id"] == triage["alert_id"]] if already: return f"Alert {triage['alert_id']} already triaged." self.agent_triages.append(triage) return ( f"Alert {triage['alert_id']} classified as {triage['severity']} " f"→ routed to {triage['team']} team." ) elif at == ActionType.QUERY_LOGS: srv = p.get("service", "") lvl = p.get("level", "") filtered = [ log for log in self.scenario["logs"] if (not srv or log["service"] == srv) and (not lvl or log["level"] == lvl) ] self.state_data["recent_logs"] = filtered[-50:] return ( f"Found {len(filtered)} log lines for service='{srv}' level='{lvl}'. " f"Showing last {min(50, len(filtered))}." ) elif at == ActionType.QUERY_METRICS: srv = p.get("service", "") keys = [k for k in self.scenario["metrics"] if k.startswith(srv)] return f"Metrics for '{srv}': {', '.join(keys) or 'none found'}." elif at == ActionType.PAGE_TEAM: return f"Paged {p.get('team', 'unknown')} team. They have been notified." elif at == ActionType.HYPOTHESIZE: self.agent_hypothesis = p.get("hypothesis", "") return f"Root cause hypothesis recorded: «{self.agent_hypothesis[:120]}»" elif at == ActionType.APPLY_FIX: fix_id = p.get("fix_id", "") self.agent_fix = fix_id if fix_id in self.scenario.get("valid_fixes", []): self.state_data["resolution_status"] = "mitigated" return ( f"✅ Fix '{fix_id}' applied successfully. " "Service metrics are recovering. Status → mitigated." ) return ( f"⚠️ Fix '{fix_id}' applied but no improvement observed. " "Check your root cause hypothesis." ) elif at == ActionType.WRITE_RUNBOOK: text = p.get("text", "").strip() self.state_data["runbook_draft"] += ("\n\n" if self.state_data["runbook_draft"] else "") + text word_count = len(self.state_data["runbook_draft"].split()) return f"Runbook updated ({word_count} words total)." elif at == ActionType.WRITE_COMMS: text = p.get("text", "").strip() self.state_data["comms_draft"] += ("\n\n" if self.state_data["comms_draft"] else "") + text return "Customer comms draft updated." elif at == ActionType.CLOSE_INCIDENT: if self.state_data["resolution_status"] != "mitigated": return "Cannot close incident — must apply a valid fix first (status must be 'mitigated')." self.state_data["resolution_status"] = "closed" return "Incident marked as closed." return "Action executed." def _compute_reward(self, action: IncidentAction) -> Tuple[float, Dict[str, float]]: """ Step-level reward shaping. All values are contribution percentages of the final [0, 1] scale. The true final score is always computed by the grader at episode end. """ reward = 0.0 breakdown = {} at = action.action_type p = action.parameters if at == ActionType.TRIAGE_ALERT: alert_id = p.get("alert_id") alert = next( (a for a in self.scenario["alerts"] if a.get("alert_id") == alert_id), None, ) if alert and "true_severity" in alert: ag_sev = str(p.get("severity", "")) gt_sev = alert["true_severity"] ag_tm = str(p.get("team", "")) gt_tm = alert["true_team"] sev_ok = ag_sev == gt_sev # Partial credit for one severity step off sev_ranks = ["P0", "P1", "P2"] near_sev = ( not sev_ok and ag_sev in sev_ranks and gt_sev in sev_ranks and abs(sev_ranks.index(ag_sev) - sev_ranks.index(gt_sev)) == 1 ) breakdown["triage_severity"] = 0.04 if sev_ok else (0.02 if near_sev else 0.0) breakdown["triage_team"] = 0.04 if (ag_tm == gt_tm) else 0.0 reward += breakdown["triage_severity"] + breakdown["triage_team"] elif at == ActionType.QUERY_LOGS: relevant = self.scenario.get("relevant_services", []) srv = p.get("service", "") if srv in relevant: breakdown["relevant_investigation"] = 0.03 reward += 0.03 # Penalise exact duplicate queries dup_count = sum( 1 for line in self.state_data["incident_timeline"] if "query_logs" in line.lower() and srv in line ) if dup_count > 1: breakdown["duplicate_query_penalty"] = -0.03 reward -= 0.03 elif at == ActionType.APPLY_FIX: if p.get("fix_id") in self.scenario.get("valid_fixes", []): breakdown["valid_fix"] = 0.20 reward += 0.20 else: breakdown["invalid_fix"] = -0.05 reward -= 0.05 elif at == ActionType.CLOSE_INCIDENT: if self.state_data["resolution_status"] == "closed": breakdown["clean_close"] = 0.05 reward += 0.05 # Always clamp to [0.0, 1.0] — never return negative per spec final = round(max(0.0, min(reward, 1.0)), 4) return final, breakdown # ── Utility ─────────────────────────────────────────────────────────────────── def json_safe(obj: Any) -> str: """Convert parameters dict to a short string safely.""" try: import json return json.dumps(obj, default=str)[:120] except Exception: return str(obj)[:120] # Allow `from server.environment import *` imports to work cleanly __all__ = ["IncidentOpsEnv"]