from __future__ import annotations from typing import Any, Dict, List, Tuple from models import DataRecord, EraseOpsAction, EraseOpsObservation, EraseOpsReward, EraseRequest from server.tasks import get_task class EraseOpsEnvironment: """Stateful OpenEnv-compatible GDPR erasure environment.""" def __init__(self) -> None: self._records: Dict[str, DataRecord] = {} self._request: EraseRequest | None = None self._task_name: str | None = None self._observation_order: List[str] = [] self._history: List[Dict[str, Any]] = [] self._step: int = 0 self._step_budget: int = 0 self._cumulative_penalty: float = 0.0 self._last_action_result: str = "Environment not initialized" self._legal_hold_violations: int = 0 self._dependency_violations: int = 0 self._wrong_user_deletions: int = 0 self._contamination_violations: int = 0 self._requests_completed: int = 0 self._escalated_legal_holds: set[str] = set() self._blocked_since: Dict[str, int] = {} self._done: bool = False self._terminated_reason: str = "" def reset(self, task_name: str, seed: int = 42) -> EraseOpsObservation: task = get_task(task_name=task_name, seed=seed) self._task_name = task.task_name self._request = task.request.model_copy(deep=True) self._records = {record.record_id: record.model_copy(deep=True) for record in task.records} self._observation_order = list(task.observation_order) self._history = [] self._step = 0 self._step_budget = task.step_budget self._cumulative_penalty = 0.0 self._last_action_result = "Environment reset" self._legal_hold_violations = 0 self._dependency_violations = 0 self._wrong_user_deletions = 0 self._contamination_violations = 0 self._requests_completed = 0 self._escalated_legal_holds = set() self._blocked_since = {} self._done = False self._terminated_reason = "" return self._build_observation() def step(self, action: EraseOpsAction) -> Tuple[EraseOpsObservation, EraseOpsReward, bool, Dict[str, Any]]: if self._request is None: raise RuntimeError("Environment must be reset before calling step") if self._done: obs = self._build_observation() reward = EraseOpsReward( total=-1.0, deletion_reward=0.0, step_penalty=0.0, legal_hold_penalty=0.0, dependency_penalty=0.0, contamination_penalty=0.0, over_deletion_penalty=0.0, completion_bonus=0.0, ) info = { "error": "Episode already terminated. Call reset before taking new actions.", "terminated_reason": self._terminated_reason, } return obs, reward, True, info self._step += 1 deletion_reward = 0.0 step_penalty = -0.01 legal_hold_penalty = 0.0 dependency_penalty = 0.0 contamination_penalty = 0.0 over_deletion_penalty = 0.0 completion_bonus = 0.0 reward_total = step_penalty record = self._records.get(action.record_id) if action.record_id else None if action.action_type == "noop": self._last_action_result = "No action taken" elif action.action_type == "inspect": if record is None: reward_total -= 1.0 self._last_action_result = "ERROR: record not found" else: self._last_action_result = f"Inspected {action.record_id}: {record.model_dump()}" elif action.action_type == "delete": if record is None: reward_total -= 1.0 self._last_action_result = "ERROR: record not found" elif record.status != "active": reward_total -= 0.5 self._last_action_result = "ERROR: already deleted" elif record.legal_hold: legal_hold_penalty -= 5.0 reward_total += legal_hold_penalty self._legal_hold_violations += 1 self._cumulative_penalty += 5.0 self._last_action_result = "VIOLATION: legal hold record cannot be deleted" elif record.user_id != self._request.user_id: contamination_penalty -= 2.0 reward_total += contamination_penalty self._wrong_user_deletions += 1 self._contamination_violations += 1 self._cumulative_penalty += 2.0 self._last_action_result = "ERROR: wrong user - GDPR scope violation" else: unresolved_deps = self._unresolved_dependencies(record) if unresolved_deps: dependency_penalty -= 3.0 reward_total += dependency_penalty self._dependency_violations += 1 self._cumulative_penalty += 3.0 self._blocked_since.setdefault(record.record_id, self._step) self._last_action_result = f"VIOLATION: must delete dependencies first: {unresolved_deps}" elif record.data_type != "PII": over_deletion_penalty -= 2.0 reward_total += over_deletion_penalty record.status = "deleted" self._blocked_since.pop(record.record_id, None) self._last_action_result = "WARNING: over-deletion - this record is not PII" else: deletion_reward += 1.0 reward_total += deletion_reward record.status = "deleted" self._blocked_since.pop(record.record_id, None) self._requests_completed += 1 self._last_action_result = f"SUCCESS: deleted {record.record_id} from {record.system}" elif action.action_type == "anonymize": if record is None: reward_total -= 1.0 self._last_action_result = "ERROR: record not found" elif record.status != "active": reward_total -= 0.5 self._last_action_result = "ERROR: already deleted" elif record.legal_hold: legal_hold_penalty -= 5.0 reward_total += legal_hold_penalty self._legal_hold_violations += 1 self._cumulative_penalty += 5.0 self._last_action_result = "VIOLATION: legal hold record cannot be anonymized" elif record.user_id != self._request.user_id: contamination_penalty -= 2.0 reward_total += contamination_penalty self._wrong_user_deletions += 1 self._contamination_violations += 1 self._cumulative_penalty += 2.0 self._last_action_result = "ERROR: wrong user - GDPR scope violation" else: unresolved_deps = self._unresolved_dependencies(record) if unresolved_deps: dependency_penalty -= 3.0 reward_total += dependency_penalty self._dependency_violations += 1 self._cumulative_penalty += 3.0 self._blocked_since.setdefault(record.record_id, self._step) self._last_action_result = f"VIOLATION: must delete dependencies first: {unresolved_deps}" elif record.data_type == "PII": reward_total -= 1.0 self._last_action_result = "ERROR: anonymization is insufficient for PII erasure" elif record.data_type in {"derived", "anonymized"}: reward_total += 0.3 record.status = "anonymized" self._blocked_since.pop(record.record_id, None) self._last_action_result = f"SUCCESS: anonymized {record.record_id}" else: reward_total -= 1.0 self._last_action_result = "ERROR: invalid anonymization target" elif action.action_type == "escalate": if record is None: reward_total -= 1.0 self._last_action_result = "ERROR: record not found" elif record.legal_hold: if record.record_id in self._escalated_legal_holds: reward_total -= 0.1 self._last_action_result = f"ERROR: {record.record_id} already escalated" else: reward_total += 0.2 self._escalated_legal_holds.add(record.record_id) self._last_action_result = f"ESCALATED: {record.record_id} flagged for legal review" else: reward_total -= 0.5 self._last_action_result = "ERROR: escalation only valid for legal hold records" else: reward_total -= 1.0 self._last_action_result = f"ERROR: unsupported action type '{action.action_type}'" done = False pii_remaining = self._pii_remaining() if pii_remaining == 0: done = True self._terminated_reason = "all_pii_deleted" if self._step < self._step_budget: completion_bonus = 0.5 reward_total += completion_bonus elif self._step >= self._step_budget: done = True self._terminated_reason = "step_budget_exhausted" elif self._cumulative_penalty > 15.0: done = True self._terminated_reason = "catastrophic_penalty" self._done = done reward = EraseOpsReward( total=round(reward_total, 4), deletion_reward=round(deletion_reward, 4), step_penalty=round(step_penalty, 4), legal_hold_penalty=round(legal_hold_penalty, 4), dependency_penalty=round(dependency_penalty, 4), contamination_penalty=round(contamination_penalty, 4), over_deletion_penalty=round(over_deletion_penalty, 4), completion_bonus=round(completion_bonus, 4), ) obs = self._build_observation() info: Dict[str, Any] = { "task_name": self._task_name, "terminated_reason": self._terminated_reason if done else "", "wrong_user_deletions": self._wrong_user_deletions, "contamination_violations": self._contamination_violations, } self._history.append( { "step": self._step, "action": action.model_dump(), "reward": reward.model_dump(), "done": done, "info": info, "observation": { "step": obs.step, "pii_remaining": obs.pii_remaining, "legal_hold_violations": obs.legal_hold_violations, "dependency_violations": obs.dependency_violations, "contamination_violations": obs.contamination_violations, "requests_completed": obs.requests_completed, "cumulative_penalty": obs.cumulative_penalty, "last_action_result": obs.last_action_result, "deadline_pressure": obs.deadline_pressure, "escalation_queue_size": obs.escalation_queue_size, "blocked_dependency_count": obs.blocked_dependency_count, }, } ) return obs, reward, done, info def state(self) -> Dict[str, Any]: return { "task_name": self._task_name, "request": self._request.model_dump() if self._request else None, "step": self._step, "step_budget": self._step_budget, "done": self._done, "terminated_reason": self._terminated_reason, "counters": { "cumulative_penalty": self._cumulative_penalty, "legal_hold_violations": self._legal_hold_violations, "dependency_violations": self._dependency_violations, "wrong_user_deletions": self._wrong_user_deletions, "contamination_violations": self._contamination_violations, "requests_completed": self._requests_completed, "pii_remaining": self._pii_remaining(), }, "last_action_result": self._last_action_result, "history_length": len(self._history), "escalated_legal_holds": sorted(self._escalated_legal_holds), "records": [ self._records[record_id].model_dump() for record_id in sorted(self._records.keys()) ], "observation_order": list(self._observation_order), } def metrics(self, include_trace: bool = False) -> Dict[str, Any]: current_state = self.state() payload: Dict[str, Any] = { "task_name": self._task_name, "step": self._step, "step_budget": self._step_budget, "done": self._done, "terminated_reason": self._terminated_reason, "pii_remaining": self._pii_remaining(), "legal_hold_violations": self._legal_hold_violations, "dependency_violations": self._dependency_violations, "wrong_user_deletions": self._wrong_user_deletions, "contamination_violations": self._contamination_violations, "requests_completed": self._requests_completed, "cumulative_penalty": self._cumulative_penalty, "last_action_result": self._last_action_result, "deadline_pressure": self._deadline_pressure(), "escalation_queue_size": self._escalation_queue_stats()[0], "escalation_queue_max_age": self._escalation_queue_stats()[1], "blocked_dependency_count": self._blocked_dependency_stats()[0], "blocked_dependency_max_age": self._blocked_dependency_stats()[1], "state": current_state, } if include_trace: payload["trace"] = list(self._history) return payload def replay(self, limit: int | None = None) -> Dict[str, Any]: trace = self._history if limit is None else self._history[-max(limit, 0) :] return { "task_name": self._task_name, "step_budget": self._step_budget, "done": self._done, "terminated_reason": self._terminated_reason, "trace_length": len(self._history), "trace": list(trace), "final_state": self.state(), } def visualization(self) -> str: status_counts = {"active": 0, "deleted": 0, "anonymized": 0} for record in self._records.values(): status_counts[record.status] = status_counts.get(record.status, 0) + 1 lines = [ f"GDPR EraseOps | task={self._task_name} | step={self._step}/{self._step_budget}", f"pii_remaining={self._pii_remaining()} | legal_hold_violations={self._legal_hold_violations} | dependency_violations={self._dependency_violations}", f"requests_completed={self._requests_completed} | cumulative_penalty={self._cumulative_penalty:.2f} | done={self._done}", f"record_statuses={status_counts}", f"last_action_result={self._last_action_result}", ] return "\n".join(lines) def _build_observation(self) -> EraseOpsObservation: if self._request is None: raise RuntimeError("Environment is not initialized") ordered_records: List[DataRecord] = [] seen: set[str] = set() for record_id in self._observation_order: if record_id in self._records: ordered_records.append(self._records[record_id].model_copy(deep=True)) seen.add(record_id) for record_id, record in self._records.items(): if record_id not in seen: ordered_records.append(record.model_copy(deep=True)) escalation_queue_size, escalation_queue_max_age = self._escalation_queue_stats() blocked_dependency_count, blocked_dependency_max_age = self._blocked_dependency_stats() return EraseOpsObservation( records=ordered_records, request=self._request.model_copy(deep=True), step=self._step, step_budget=self._step_budget, cumulative_penalty=round(self._cumulative_penalty, 4), last_action_result=self._last_action_result, pii_remaining=self._pii_remaining(), legal_hold_violations=self._legal_hold_violations, dependency_violations=self._dependency_violations, contamination_violations=self._contamination_violations, requests_completed=self._requests_completed, deadline_pressure=self._deadline_pressure(), escalation_queue_size=escalation_queue_size, escalation_queue_max_age=escalation_queue_max_age, blocked_dependency_count=blocked_dependency_count, blocked_dependency_max_age=blocked_dependency_max_age, ) def _unresolved_dependencies(self, record: DataRecord) -> List[str]: unresolved: List[str] = [] for dep_id in record.dependencies: dep = self._records.get(dep_id) if dep is None: unresolved.append(dep_id) continue if dep.status == "deleted": continue # Escalated legal-hold dependencies are treated as policy-resolved. if dep.legal_hold and dep_id in self._escalated_legal_holds: continue unresolved.append(dep_id) return unresolved def _deadline_pressure(self) -> float: if self._request is None or self._request.deadline_steps <= 0: return 0.0 return round(min(1.0, self._step / float(self._request.deadline_steps)), 4) def _escalation_queue_stats(self) -> Tuple[int, int]: if self._request is None: return 0, 0 pending = [ record for record in self._records.values() if record.user_id == self._request.user_id and record.legal_hold and record.status == "active" and record.record_id not in self._escalated_legal_holds ] if not pending: return 0, 0 age = max(0, self._step - int(self._request.submitted_at_step or 0)) return len(pending), age def _blocked_dependency_stats(self) -> Tuple[int, int]: if self._request is None: return 0, 0 blocked_count = 0 max_age = 0 for record in self._records.values(): if record.user_id != self._request.user_id or record.data_type != "PII" or record.status != "active": continue unresolved = self._unresolved_dependencies(record) if not unresolved: self._blocked_since.pop(record.record_id, None) continue blocked_count += 1 first_seen = self._blocked_since.setdefault(record.record_id, self._step) max_age = max(max_age, max(0, self._step - first_seen)) return blocked_count, max_age def _pii_remaining(self) -> int: if self._request is None: return 0 return sum( 1 for record in self._records.values() if record.user_id == self._request.user_id and record.data_type == "PII" and record.status == "active" )