"""Internal mutable state for a single episode.""" from __future__ import annotations from typing import Any, Literal from env.customers import CustomerQueueManager from env.incident_history import IncidentHistoryStore from env.knowledge_base import KnowledgeBase, PersistentKnowledgeBase from env.policy_engine import PolicyEngine from env.runbooks import RunbookEngine from env.stakeholders import StakeholderManager from env.world import WorldState from graders.investigation_grader import ACTION_COSTS, EvidenceChain, SeverityReEvaluation from models.action import ApplyFixAction from models.incident import IncidentScenario from models.observation import ActionRecord, Observation, Phase from models.ticket import TicketData PHASE_VALID_ACTIONS: dict[Phase, frozenset[str]] = { "unclassified": frozenset(["classify"]), "classified": frozenset(["route", "escalate"]), "routed": frozenset(["respond", "escalate", "resolve", "request_info"]), "responding": frozenset(["respond", "escalate", "resolve", "request_info"]), "escalated": frozenset(["resolve"]), "resolved": frozenset(), } MAX_STEPS: dict[str, int] = { "easy": 8, "medium": 9, "hard": 10, } IncidentPhase = Literal["triage", "investigation", "response", "resolution"] IncidentSeverity = Literal["medium", "high", "critical", "P0"] INCIDENT_PHASE_VALID_ACTIONS: dict[IncidentPhase, frozenset[str]] = { "triage": frozenset( ["classify", "check_monitoring", "query_kb", "query_incident_history", "follow_runbook_step"] ), "investigation": frozenset( [ "check_monitoring", "probe_service", "fetch_logs", "fetch_user_data", "check_billing", "query_kb", "check_policy", "query_incident_history", "follow_runbook_step", "classify", "route", ] ), "response": frozenset( [ "apply_fix", "rollback_fix", "respond", "escalate", "request_info", "notify_stakeholders", "check_policy", "fetch_user_data", "check_billing", "query_kb", "follow_runbook_step", ] ), "resolution": frozenset( [ "verify_fix", "resolve", "respond", "write_postmortem", "update_kb", "notify_stakeholders", ] ), } SEVERITY_ESCALATION: dict[int, IncidentSeverity] = { 10: "high", 25: "critical", 40: "P0", } def compute_max_total_reward(ticket: TicketData) -> float: """Achievable maximum reward for an optimal agent on this ticket. Accounts for action rewards *and* any unavoidable SLA penalties incurred when the minimum optimal path exceeds the ticket's SLA deadline. """ total = 0.10 # classify base if ticket.gold_priority in ("critical", "high"): total += 0.10 # urgency bonus total += 0.10 # route min_steps = 3 # classify + route + resolve (always required) if ticket.difficulty in ("medium", "hard"): total += 0.20 # respond min_steps += 1 if ticket.partial_info: total += 0.05 # request_info bonus min_steps += 1 if ticket.requires_escalation: total += 0.15 # escalate min_steps += 1 total += 0.25 # resolve sla = ticket.effective_sla_steps for step_idx in range(min_steps): if step_idx >= sla: total -= 0.02 * (step_idx - sla + 1) return round(total, 4) class InternalState: """Tracks episode progress, phase transitions, and cumulative quality scores.""" __slots__ = ( "ticket", "phase", "steps_taken", "max_steps", "max_total_reward", "actions_log", "cumulative_reward", "classification_correct", "routing_correct", "urgency_handled", "response_quality_score", "resolution_quality_score", "escalation_score", "constraints_violated", "done", "last_action_json", # v2 additions "sla_steps", "urgency_penalty_accrued", "info_requested", # v3 — interpretability "last_reward_breakdown", ) def __init__(self, ticket: TicketData) -> None: self.ticket = ticket self.phase: Phase = "unclassified" self.steps_taken: int = 0 self.max_steps: int = MAX_STEPS[ticket.difficulty] self.max_total_reward: float = compute_max_total_reward(ticket) self.actions_log: list[ActionRecord] = [] self.cumulative_reward: float = 0.0 self.classification_correct: bool | None = None self.routing_correct: bool | None = None self.urgency_handled: bool = False self.response_quality_score: float | None = None self.resolution_quality_score: float | None = None self.escalation_score: float | None = None self.constraints_violated: list[str] = [] self.done: bool = False self.last_action_json: str | None = None # v2 self.sla_steps: int = ticket.effective_sla_steps self.urgency_penalty_accrued: float = 0.0 self.info_requested: bool = False # v3 — interpretability self.last_reward_breakdown: dict[str, float] = {} # ---- helpers -------------------------------------------------------- @property def available_actions(self) -> list[str]: actions = PHASE_VALID_ACTIONS[self.phase] if self.info_requested: actions = actions - frozenset(["request_info"]) return sorted(actions) def record_action( self, action_summary: str, feedback: str, reward: float ) -> None: self.actions_log.append( ActionRecord( step=self.steps_taken, action_taken=action_summary, env_feedback=feedback, reward_earned=round(reward, 4), ) ) self.cumulative_reward += reward self.steps_taken += 1 if self.steps_taken >= self.max_steps: self.done = True def to_observation(self) -> Observation: return Observation( ticket_id=self.ticket.ticket_id, ticket_text=self.ticket.ticket_text, customer_sentiment=self.ticket.customer_sentiment, customer_tier=self.ticket.customer_tier, category_hint=self.ticket.category_hint, history=list(self.actions_log), pending_tickets=0, current_step=self.steps_taken, max_steps=self.max_steps, constraints=list(self.ticket.constraints), available_actions=self.available_actions, phase=self.phase, sla_steps_remaining=max(0, self.sla_steps - self.steps_taken), customer_value=self.ticket.customer_value, max_total_reward=self.max_total_reward, ) def to_info(self) -> dict[str, Any]: mtr = self.max_total_reward info: dict[str, Any] = { "phase": self.phase, "steps_taken": self.steps_taken, "max_steps": self.max_steps, "cumulative_reward": round(self.cumulative_reward, 4), "max_total_reward": mtr, "normalized_score": round( min(max(self.cumulative_reward / mtr, 0.0), 1.0), 4 ), "classification_correct": self.classification_correct, "routing_correct": self.routing_correct, "urgency_handled": self.urgency_handled, "response_quality_score": self.response_quality_score, "resolution_quality_score": self.resolution_quality_score, "escalation_score": self.escalation_score, "constraints_violated": list(self.constraints_violated), "difficulty": self.ticket.difficulty, "sla_steps": self.sla_steps, "sla_overage": max(0, self.steps_taken - self.sla_steps), "urgency_penalty_accrued": round(self.urgency_penalty_accrued, 4), "customer_value": self.ticket.customer_value, "reward_breakdown": dict(self.last_reward_breakdown), } if self.done: info["final_score_breakdown"] = self._compute_final_breakdown() return info def _compute_final_breakdown(self) -> dict[str, float]: """Episode-level weighted breakdown mirroring ``grade_episode``.""" cls_s = 1.0 if self.classification_correct else 0.0 rte_s = 1.0 if self.routing_correct else 0.0 rsp_s = self.response_quality_score if self.response_quality_score is not None else 0.0 res_s = self.resolution_quality_score if self.resolution_quality_score is not None else 0.0 esc_s = max(0.0, self.escalation_score) if self.escalation_score is not None else 0.0 urg_s = 1.0 if self.urgency_handled else 0.0 eff_s = max(0.0, 1.0 - self.steps_taken / self.max_steps) if self.max_steps > 0 else 0.0 sla_overage = max(0, self.steps_taken - self.sla_steps) sla_s = max(0.0, 1.0 - sla_overage * 0.2) constraint_pen = len(self.constraints_violated) * 0.05 raw = ( 0.15 * cls_s + 0.10 * rte_s + 0.20 * rsp_s + 0.20 * res_s + 0.10 * esc_s + 0.10 * urg_s + 0.05 * eff_s + 0.10 * sla_s ) return { "classification": round(0.15 * cls_s, 4), "routing": round(0.10 * rte_s, 4), "response_quality": round(0.20 * rsp_s, 4), "resolution_quality": round(0.20 * res_s, 4), "escalation": round(0.10 * esc_s, 4), "urgency": round(0.10 * urg_s, 4), "efficiency": round(0.05 * eff_s, 4), "sla_compliance": round(0.10 * sla_s, 4), "constraint_penalty": round(-constraint_pen, 4), "total": round(max(0.0, min(raw - constraint_pen, 1.0)), 4), } class ResourceBudget: """Finite resources available during an incident episode.""" __slots__ = ( "max_fix_attempts", "max_escalations", "max_stakeholder_notifications", "remaining_fix_attempts", "remaining_escalations", "remaining_notifications", ) def __init__( self, max_fix_attempts: int = 3, max_escalations: int = 2, max_stakeholder_notifications: int = 5, ) -> None: self.max_fix_attempts = max_fix_attempts self.max_escalations = max_escalations self.max_stakeholder_notifications = max_stakeholder_notifications self.remaining_fix_attempts = max_fix_attempts self.remaining_escalations = max_escalations self.remaining_notifications = max_stakeholder_notifications def consume(self, resource: str) -> bool: """Consume one unit of a named resource if available.""" if resource == "fix_attempt": if self.remaining_fix_attempts <= 0: return False self.remaining_fix_attempts -= 1 return True if resource == "escalation": if self.remaining_escalations <= 0: return False self.remaining_escalations -= 1 return True if resource == "notification": if self.remaining_notifications <= 0: return False self.remaining_notifications -= 1 return True return False class AuditEntry: """Compliance audit row for one action.""" __slots__ = ( "step", "timestamp_simulated", "action_type", "target", "rationale_required", "policy_checked", "compliant", ) def __init__( self, *, step: int, timestamp_simulated: str, action_type: str, target: str, rationale_required: bool, policy_checked: bool, compliant: bool, ) -> None: self.step = step self.timestamp_simulated = timestamp_simulated self.action_type = action_type self.target = target self.rationale_required = rationale_required self.policy_checked = policy_checked self.compliant = compliant class AuditTrail: """Compliance audit trail across incident actions.""" __slots__ = ("entries",) def __init__(self) -> None: self.entries: list[AuditEntry] = [] def append(self, entry: AuditEntry) -> None: """Append a new audit entry.""" self.entries.append(entry) def grade_compliance(self) -> float: """Return required-policy-check compliance ratio.""" requiring = [ entry for entry in self.entries if entry.action_type in _POLICY_SENSITIVE_ACTIONS ] if not requiring: return 1.0 compliant = sum(1 for entry in requiring if entry.policy_checked and entry.compliant) return round(compliant / len(requiring), 4) class ChaosEvent: """Mid-episode injected failure event.""" __slots__ = ("step", "new_service", "reason", "alert_text") def __init__(self, step: int, new_service: str, reason: str, alert_text: str) -> None: self.step = step self.new_service = new_service self.reason = reason self.alert_text = alert_text class ChaosInjector: """Inject deterministic new failures during response phase.""" CHAOS_TRIGGERS: dict[str, dict[str, object]] = { "hard": { "trigger_step": 35, "probability": 0.5, "new_failure": { "service": "notifications", "mode": "queue_overflow", "reason": "Backpressure from payment retry storm", }, }, "nightmare": { "trigger_step": 25, "probability": 1.0, "new_failure": { "service": "analytics", "mode": "batch_job_runaway", "reason": "Error logging spike triggered batch reprocessing", }, }, } def maybe_inject(self, world: WorldState, step: int, difficulty: str) -> ChaosEvent | None: """Inject deterministic chaos event based on seed and step.""" config = self.CHAOS_TRIGGERS.get(difficulty) if config is None: return None trigger_step = int(config["trigger_step"]) probability = float(config["probability"]) if step < trigger_step: return None if not self._should_trigger(world.seed, step, probability): return None failure = config["new_failure"] service = str(failure["service"]) mode = str(failure["mode"]) reason = str(failure["reason"]) world.service_mesh.inject_failure(service, mode) return ChaosEvent( step=step, new_service=service, reason=reason, alert_text=f"NEW ALERT: {service} showing errors", ) @staticmethod def _should_trigger(seed: int, step: int, probability: float) -> bool: threshold = int(probability * 100) value = (seed * 31 + step * 17) % 100 return value < threshold class IncidentState: """Tracks incident lifecycle progression and integrated world state.""" __slots__ = ( "incident", "world", "incident_phase", "triage_complete", "investigation_complete", "response_complete", "episode_done", "root_cause_identified", "fix_applied", "fix_verified", "tickets_resolved", "tools_used_sequence", "policies_checked", "kb_queried", "logs_checked_for", "postmortem_written", "kb_updated", "steps_taken", "max_steps", "cumulative_reward", "last_action_json", "last_reward_breakdown", "known_facts", "active_policies", "tool_results", "active_alerts", "resource_budget", "audit_trail", "current_severity", "severity_re_eval", "_pending_reclassification", "chaos_injector", "has_escalated", "total_action_cost", "evidence_chain", "crm", "billing", "policy_engine", "history_store", "runbook_engine", "stakeholder_mgr", "customer_queue_mgr", "persistent_kb", "knowledge_base", "suggested_runbook", "actions_log", ) def __init__( self, incident: IncidentScenario, world: WorldState, *, crm: object, billing: object, policy_engine: PolicyEngine, history_store: IncidentHistoryStore, runbook_engine: RunbookEngine, stakeholder_mgr: StakeholderManager, customer_queue_mgr: CustomerQueueManager, persistent_kb: PersistentKnowledgeBase, knowledge_base: KnowledgeBase, suggested_runbook: dict[str, object] | None, ) -> None: self.incident = incident self.world = world self.incident_phase: IncidentPhase = "triage" self.triage_complete = False self.investigation_complete = False self.response_complete = False self.episode_done = False self.root_cause_identified = False self.fix_applied = False self.fix_verified = False self.tickets_resolved: list[str] = [] self.tools_used_sequence: list[str] = [] self.policies_checked: set[str] = set() self.kb_queried = False self.logs_checked_for: set[str] = set() self.postmortem_written = False self.kb_updated = False self.steps_taken = 0 self.max_steps = incident.max_steps self.cumulative_reward = 0.0 self.last_action_json: str | None = None self.last_reward_breakdown: dict[str, float] = {} self.known_facts: dict[str, object] = {} self.active_policies: dict[str, object] = {} self.tool_results: dict[str, object] | None = None self.active_alerts: list[str] = [] self.resource_budget = ResourceBudget() self.audit_trail = AuditTrail() self.current_severity = _severity_for_difficulty(incident.difficulty) self.severity_re_eval = SeverityReEvaluation() self._pending_reclassification: IncidentSeverity | None = None self.chaos_injector = ChaosInjector() self.has_escalated = False self.total_action_cost = 0.0 self.evidence_chain = EvidenceChain() self.crm = crm self.billing = billing self.policy_engine = policy_engine self.history_store = history_store self.runbook_engine = runbook_engine self.stakeholder_mgr = stakeholder_mgr self.customer_queue_mgr = customer_queue_mgr self.persistent_kb = persistent_kb self.knowledge_base = knowledge_base self.suggested_runbook = suggested_runbook self.actions_log: list[ActionRecord] = [] @property def available_actions(self) -> list[str]: """Return actions available in the current incident phase.""" return sorted(INCIDENT_PHASE_VALID_ACTIONS[self.incident_phase]) def record_action(self, action_type: str, feedback: str, reward: float) -> None: """Record action outcome in history and counters.""" self.actions_log.append( ActionRecord( step=self.steps_taken, action_taken=action_type, env_feedback=feedback, reward_earned=round(reward, 4), ) ) self.steps_taken += 1 self.cumulative_reward += reward self.tools_used_sequence.append(action_type) self.total_action_cost = round( self.total_action_cost + ACTION_COSTS.get(action_type, 0), 2 ) if self.steps_taken >= self.max_steps: self.episode_done = True def add_audit_entry( self, *, action_type: str, target: str, policy_checked: bool, compliant: bool, ) -> None: """Append compliance audit entry for current step.""" self.audit_trail.append( AuditEntry( step=self.steps_taken, timestamp_simulated=f"T+{self.steps_taken:03d}", action_type=action_type, target=target, rationale_required=action_type in {"apply_fix", "escalate"}, policy_checked=policy_checked, compliant=compliant, ) ) def advance_phase(self) -> None: """Advance phase when completion conditions are met.""" if self.incident_phase == "triage" and self.triage_complete: self.incident_phase = "investigation" return if self.incident_phase == "investigation" and self.investigation_complete: self.incident_phase = "response" return if self.incident_phase == "response" and self.response_complete: self.incident_phase = "resolution" def all_objectives_complete(self) -> bool: """Return True when incident is fully resolved.""" return self.fix_verified and self.postmortem_written and self.kb_updated def apply_severity_auto_escalation(self) -> None: """Escalate severity if unresolved for long durations.""" for step_trigger, new_severity in SEVERITY_ESCALATION.items(): if self.steps_taken >= step_trigger and not self.response_complete: self.current_severity = new_severity def maybe_update_reclassification(self) -> None: """Check whether evidence implies severity reclassification.""" reward, _, target = self.severity_re_eval.check_reclassification( evidence_chain=self.evidence_chain, current_step=self.steps_taken, current_severity=self.current_severity, ) if reward != 0.0: self._pending_reclassification = target def mark_fix_applied(self, action: ApplyFixAction, correct: bool) -> None: """Mark fix state and progression.""" self.fix_applied = True if correct: self.response_complete = True self.add_audit_entry( action_type="apply_fix", target=action.service_name, policy_checked="escalation" in self.policies_checked, compliant=correct, ) def to_observation(self) -> Observation: """Build incident-mode observation payload.""" ticket = self.world.support_queue[0] if self.world.support_queue else None ticket_text = ( getattr(ticket, "body", None) or getattr(ticket, "ticket_text", None) or "" ) if ticket is not None else self.incident.description alerts = self.world.service_mesh.generate_alerts(self.steps_taken) self.active_alerts = [alert.message for alert in alerts] return Observation( ticket_id=ticket.ticket_id if ticket is not None else self.incident.incident_id, ticket_text=ticket_text, customer_sentiment="frustrated", customer_tier="enterprise", customer_value="high", category_hint=None, constraints=[], phase="responding" if self.incident_phase in ("response", "resolution") else "classified", available_actions=self.available_actions, current_step=self.steps_taken, max_steps=self.max_steps, sla_steps_remaining=max(0, self.max_steps - self.steps_taken), history=list(self.actions_log), max_total_reward=self.incident.max_total_reward, incident_id=self.incident.incident_id, incident_title=self.incident.title, mode="incident", system_status=self.world.service_mesh.get_health_summary(), active_alerts=list(self.active_alerts), tool_results=self.tool_results, known_facts=dict(self.known_facts), active_policies=dict(self.active_policies), stakeholder_patience=self.stakeholder_mgr.get_patience_levels(), pending_customer_tickets=len(self.world.support_queue), incident_phase=self.incident_phase, suggested_runbook=self.suggested_runbook, total_incident_cost=round(self.world.total_downtime_cost + self.total_action_cost, 2), ) def to_info(self) -> dict[str, object]: """Build incident-mode diagnostics payload.""" return { "mode": "incident", "incident_id": self.incident.incident_id, "incident_phase": self.incident_phase, "steps_taken": self.steps_taken, "max_steps": self.max_steps, "cumulative_reward": round(self.cumulative_reward, 4), "known_facts": dict(self.known_facts), "active_policies": dict(self.active_policies), "total_incident_cost": round(self.world.total_downtime_cost + self.total_action_cost, 2), "reward_breakdown": dict(self.last_reward_breakdown), "resource_budget": { "remaining_fix_attempts": self.resource_budget.remaining_fix_attempts, "remaining_escalations": self.resource_budget.remaining_escalations, "remaining_notifications": self.resource_budget.remaining_notifications, }, "compliance_score": self.audit_trail.grade_compliance(), "current_severity": self.current_severity, } _POLICY_SENSITIVE_ACTIONS = frozenset( ["apply_fix", "escalate", "notify_stakeholders", "update_kb"] ) def _severity_for_difficulty(difficulty: str) -> IncidentSeverity: if difficulty == "easy": return "medium" if difficulty == "medium": return "high" return "critical"