"""Core environment implementing the OpenEnv step() / reset() / state() API. This module owns all mutable episode state. It is deliberately a single- episode, per-session environment — the FastAPI layer maintains one instance per session ID, ensuring concurrent agents never share state. """ from __future__ import annotations import copy import random import time from collections import deque from datetime import datetime, timezone from typing import Any, Deque, Dict, List, Optional, Tuple from src.graders import grade from src.models import ( Action, ActionType, EnvironmentState, GraderResult, IncidentSeverity, IncidentStatus, Observation, Reward, ServiceMetrics, StepResult, ) from src.rewards import compute_step_reward from src.scenarios import Scenario, apply_blast_radius, get_scenario class IncidentResponseEnv: """Incident Response Triage environment. Lifecycle: env = IncidentResponseEnv() obs = env.reset("severity_classification") while not done: result = env.step(action) obs, reward, done, info = result.observation, result.reward, result.done, result.info grader_result = env.grade() """ # How many per-step metric snapshots to retain per service _TSDB_MAX_SAMPLES: int = 64 def __init__(self) -> None: self._scenario: Optional[Scenario] = None self._task_id: Optional[str] = None self._step: int = 0 self._done: bool = True self._cumulative_reward: float = 0.0 # Progressive state self._investigated: List[str] = [] self._severity_classified: Optional[IncidentSeverity] = None self._diagnosis: Optional[str] = None self._remediations: List[str] = [] self._escalations: List[str] = [] self._communications: List[str] = [] self._actions_history: List[Dict[str, Any]] = [] self._incident_status: IncidentStatus = IncidentStatus.OPEN self._last_message: str = "" # Logs / metrics revealed so far self._revealed_logs: Dict[str, list] = {} self._revealed_metrics: Dict[str, Any] = {} # TSDB ring buffer: service -> deque of (unix_timestamp, ServiceMetrics) # Populated after every step so /prometheus/query_range returns real history. self._metric_history: Dict[str, Deque[Tuple[float, ServiceMetrics]]] = {} # ------------------------------------------------------------------ # reset() # ------------------------------------------------------------------ def reset(self, task_id: str, variant_seed: int = 0) -> Observation: """Reset the environment for a new episode on the given task. Args: task_id: Task to run. variant_seed: Scenario variant index (default 0 = primary scenario). """ scenario = get_scenario(task_id, variant_seed=variant_seed) self._scenario = scenario self._task_id = task_id self._step = 0 self._done = False self._cumulative_reward = 0.0 self._investigated = [] self._severity_classified = None self._diagnosis = None self._remediations = [] self._escalations = [] self._communications = [] self._actions_history = [] self._incident_status = IncidentStatus.OPEN self._last_message = "Incident opened. Review the alerts and begin your investigation." self._revealed_logs = {} self._revealed_metrics = {} self._metric_history = {} # Record step-0 snapshot so range queries have at least one data point self._record_metric_snapshot() return self._build_observation() # ------------------------------------------------------------------ # step() # ------------------------------------------------------------------ def step(self, action: Action) -> StepResult: """Process one agent action and return the result.""" if self._done: raise RuntimeError("Episode is done. Call reset() first.") if self._scenario is None: raise RuntimeError("Environment not initialised. Call reset() first.") self._step += 1 scenario = self._scenario # Record action self._actions_history.append(action.model_dump()) # Process action effects self._process_action(action, scenario) # Compute reward reward = compute_step_reward( action=action, scenario=scenario, step_number=self._step, already_investigated=self._investigated, already_classified=self._severity_classified, already_diagnosed=self._diagnosis, already_remediated=self._remediations, already_escalated=self._escalations, already_communicated=self._communications, actions_history=self._actions_history, ) self._cumulative_reward += reward.value # Apply action state changes (after reward so duplicates are penalised first) self._apply_state_changes(action, scenario) # Record metric snapshot into TSDB ring buffer self._record_metric_snapshot() # Check episode termination done = self._check_done(scenario) self._done = done obs = self._build_observation() info: Dict[str, Any] = { "cumulative_reward": round(self._cumulative_reward, 4), "steps_remaining": max(0, scenario.max_steps - self._step), } if done: info["grader"] = self.grade().model_dump() return StepResult(observation=obs, reward=reward, done=done, info=info) # ------------------------------------------------------------------ # state() # ------------------------------------------------------------------ def state(self) -> EnvironmentState: """Return the full internal state snapshot.""" return EnvironmentState( task_id=self._task_id or "", scenario_id=self._scenario.scenario_id if self._scenario else "", step_number=self._step, max_steps=self._scenario.max_steps if self._scenario else 0, incident_status=self._incident_status, done=self._done, cumulative_reward=round(self._cumulative_reward, 4), total_steps_taken=self._step, alerts=list(self._scenario.initial_alerts) if self._scenario else [], actions_history=copy.deepcopy(self._actions_history), severity_classified=self._severity_classified, diagnosis=self._diagnosis, remediations_applied=list(self._remediations), escalations_made=list(self._escalations), communications_sent=list(self._communications), investigated_services=list(self._investigated), ) # ------------------------------------------------------------------ # grade() # ------------------------------------------------------------------ def grade(self) -> GraderResult: """Grade the current episode. Can be called mid-episode or after done.""" if self._scenario is None or self._task_id is None: raise RuntimeError("No episode in progress.") return grade(self._task_id, self.state(), self._scenario) def live_metrics(self) -> Dict[str, ServiceMetrics]: """Return service metrics with blast-radius degradation at the current step. Safe to call at any point (including before any actions are taken). Returns an empty dict when no episode is in progress. This is the same numerical data the agent would eventually see via investigate actions, but served here without consuming an action slot — analogous to a Prometheus scrape that is always available passively. """ if self._scenario is None: return {} return apply_blast_radius(self._scenario, self._step) def metric_history( self, start: float, end: float, step_seconds: float = 1.0, ) -> Dict[str, List[Tuple[float, ServiceMetrics]]]: """Return per-service metric history in the [start, end] time window. This powers the Prometheus-compatible ``/prometheus/query_range`` endpoint. Each entry is a ``(unix_timestamp, ServiceMetrics)`` tuple, sampled once per environment step. The ``step_seconds`` parameter is accepted for API compatibility but does not resample — the ring buffer already stores one sample per episode step. Returns an empty dict when no episode is in progress. """ if not self._metric_history: return {} result: Dict[str, List[Tuple[float, ServiceMetrics]]] = {} for svc, dq in self._metric_history.items(): samples = [(ts, m) for ts, m in dq if start <= ts <= end] if samples: result[svc] = samples return result # ------------------------------------------------------------------ # Private helpers # ------------------------------------------------------------------ def _record_metric_snapshot(self) -> None: """Append the current blast-radius metrics to the per-service ring buffer.""" if self._scenario is None: return ts = time.time() live = apply_blast_radius(self._scenario, self._step) for svc, metrics in live.items(): if svc not in self._metric_history: self._metric_history[svc] = deque(maxlen=self._TSDB_MAX_SAMPLES) self._metric_history[svc].append((ts, metrics)) def _process_action(self, action: Action, scenario: Scenario) -> None: """Process action effects (messages, state transitions).""" if action.action_type == ActionType.INVESTIGATE: target = (action.target or "").strip() if target not in scenario.available_services: self._last_message = ( f"Unknown service '{target}'. " f"Available: {', '.join(scenario.available_services)}" ) elif target in self._investigated: self._last_message = f"Already investigated {target}. Logs and metrics available." else: # Reveal logs (always static — logs are historical records) if target in scenario.service_logs: self._revealed_logs[target] = [ entry.model_dump() for entry in scenario.service_logs[target] ] # Reveal LIVE metrics with blast-radius degradation applied live_metrics = apply_blast_radius(scenario, self._step) if target in live_metrics: self._revealed_metrics[target] = live_metrics[target].model_dump() self._last_message = ( f"Investigation of {target} complete. Logs and live metrics now available." ) if self._incident_status == IncidentStatus.OPEN: self._incident_status = IncidentStatus.INVESTIGATING elif action.action_type == ActionType.CLASSIFY: severity_str = action.parameters.get("severity", "") try: sev = IncidentSeverity(severity_str) self._last_message = f"Incident classified as {sev.value}." except ValueError: self._last_message = ( f"Invalid severity '{severity_str}'. Use P1, P2, P3, or P4." ) elif action.action_type == ActionType.DIAGNOSE: root_cause = action.parameters.get("root_cause", "") target_svc = (action.target or "").strip() self._last_message = ( f"Diagnosis recorded: root cause in {target_svc} — {root_cause[:120]}" ) elif action.action_type == ActionType.REMEDIATE: rem_action = action.parameters.get("action", "") target_svc = (action.target or "").strip() if not rem_action or not target_svc: self._last_message = "Remediation requires 'action' parameter and 'target' service." else: self._last_message = ( f"Remediation '{rem_action}' applied to {target_svc}." ) if self._incident_status in (IncidentStatus.OPEN, IncidentStatus.INVESTIGATING): self._incident_status = IncidentStatus.MITIGATING elif action.action_type == ActionType.ESCALATE: team = (action.target or "").strip() priority = action.parameters.get("priority", "high") message = action.parameters.get("message", "") self._last_message = ( f"Escalated to {team} (priority: {priority}). " f"Message: {message[:80]}" ) elif action.action_type == ActionType.COMMUNICATE: channel = (action.target or "status_page").strip() message = action.parameters.get("message", "") self._last_message = ( f"Status update posted to {channel}: {message[:100]}" ) def _apply_state_changes(self, action: Action, scenario: Scenario) -> None: """Persist state changes after reward is computed.""" if action.action_type == ActionType.INVESTIGATE: target = (action.target or "").strip() if target in scenario.available_services and target not in self._investigated: self._investigated.append(target) elif action.action_type == ActionType.CLASSIFY: severity_str = action.parameters.get("severity", "") try: self._severity_classified = IncidentSeverity(severity_str) except ValueError: pass elif action.action_type == ActionType.DIAGNOSE: if self._diagnosis is None: self._diagnosis = action.parameters.get("root_cause", "") elif action.action_type == ActionType.REMEDIATE: rem_action = action.parameters.get("action", "") target_svc = (action.target or "").strip() if rem_action and target_svc: key = f"{rem_action}:{target_svc}" if key not in self._remediations: self._remediations.append(key) elif action.action_type == ActionType.ESCALATE: team = (action.target or "").strip() if team and team not in self._escalations: self._escalations.append(team) elif action.action_type == ActionType.COMMUNICATE: message = action.parameters.get("message", "") if message: self._communications.append(message[:200]) def _check_done(self, scenario: Scenario) -> bool: """Episode ends when max steps reached or incident resolved.""" if self._step >= scenario.max_steps: self._last_message += " [Episode ended: max steps reached.]" return True # For easy task: done once classified if scenario.task_id == "severity_classification" and self._severity_classified is not None: # Give agent a chance to investigate first, but if classified, we're done # Actually let them keep going for a few more steps if they want if self._step >= 2 or self._severity_classified is not None: # Check if the last action was classify if (self._actions_history and self._actions_history[-1].get("action_type") == ActionType.CLASSIFY.value): self._incident_status = IncidentStatus.RESOLVED self._last_message += " [Episode complete: severity classified.]" return True # For medium: done once diagnosed AND remediated if scenario.task_id == "root_cause_analysis": if self._diagnosis and len(self._remediations) > 0: self._incident_status = IncidentStatus.RESOLVED self._last_message += " [Episode complete: diagnosed and remediated.]" return True # For hard: done once all major actions taken if scenario.task_id == "full_incident_management": has_severity = self._severity_classified is not None has_diagnosis = self._diagnosis is not None has_remediation = len(self._remediations) > 0 has_escalation = len(self._escalations) > 0 has_communication = len(self._communications) > 0 if all([has_severity, has_diagnosis, has_remediation, has_escalation, has_communication]): self._incident_status = IncidentStatus.RESOLVED self._last_message += " [Episode complete: full incident management workflow completed.]" return True return False def _build_observation(self) -> Observation: """Construct the current observation from internal state.""" scenario = self._scenario if scenario is None: raise RuntimeError("No scenario loaded.") # Reconstruct typed logs from revealed data from src.models import LogEntry, ServiceMetrics typed_logs = {} for svc, entries in self._revealed_logs.items(): typed_logs[svc] = [LogEntry(**e) for e in entries] typed_metrics = {} for svc, m in self._revealed_metrics.items(): typed_metrics[svc] = ServiceMetrics(**m) return Observation( incident_id=scenario.incident_id, timestamp=datetime.now(timezone.utc).isoformat(), step_number=self._step, max_steps=scenario.max_steps, task_id=scenario.task_id, task_description=scenario.description, alerts=list(scenario.initial_alerts), available_services=list(scenario.available_services), investigated_services=list(self._investigated), logs=typed_logs, metrics=typed_metrics, incident_status=self._incident_status, severity_classified=self._severity_classified, diagnosis=self._diagnosis, actions_taken=[ f"Step {i+1}: {a.get('action_type', '?')} → {a.get('target', 'N/A')}" for i, a in enumerate(self._actions_history) ], remediations_applied=list(self._remediations), escalations_made=list(self._escalations), communications_sent=list(self._communications), message=self._last_message, )