from __future__ import annotations import os import random from dataclasses import dataclass, field from pathlib import Path from typing import Any from uuid import uuid4 from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State try: from ..models import ( ActionType, IssueView, OpsguardAction, OpsguardObservation, SecurityVerdict, ) from ..world.adversary import attack_to_issue_row, synthesize_attacks from ..world.behavior_tokens import analyze_diff from ..world.contributor_profile import ContributorProfileStore from ..world.curriculum import Curriculum from ..world.db import IssueRow, RepoDB from ..world.grader import RewardConfig, grade_step, grade_terminal from ..world.memory import MemoryStore from ..world.scenarios import ScenarioSpec, get_scenario, list_scenarios from ..world.trainable_adversary import LLMSpamAdversary, TrainableAdversaryConfig except (ImportError, ValueError): from models import ( ActionType, IssueView, OpsguardAction, OpsguardObservation, SecurityVerdict, ) from world.adversary import attack_to_issue_row, synthesize_attacks from world.behavior_tokens import analyze_diff from world.contributor_profile import ContributorProfileStore from world.curriculum import Curriculum from world.db import IssueRow, RepoDB from world.grader import RewardConfig, grade_step, grade_terminal from world.memory import MemoryStore from world.scenarios import ScenarioSpec, get_scenario, list_scenarios from world.trainable_adversary import LLMSpamAdversary, TrainableAdversaryConfig _DB_PATH = os.environ.get("OPSGUARD_DB", str(Path(__file__).resolve().parent.parent / "data" / "repo.db")) _REPO = os.environ.get("OPSGUARD_REPO", "huggingface/peft") _ADV_LORA = os.environ.get("OPSGUARD_ADVERSARY_LORA", "").strip() or None @dataclass class _EpisodeState: scenario: ScenarioSpec queue: list[IssueRow] pos: int = 0 actions_taken: list[dict[str, Any]] = field(default_factory=list) info_requested_for: set[int] = field(default_factory=set) diff_reviewed_for: set[int] = field(default_factory=set) revoked_logins: set[str] = field(default_factory=set) last_memory_hits: list[dict[str, Any]] = field(default_factory=list) memory: MemoryStore = field(default_factory=MemoryStore) profiles: ContributorProfileStore = field(default_factory=ContributorProfileStore) cumulative_reward: float = 0.0 n_resolved_legit: int = 0 n_total_legit: int = 0 n_attacks_caught: int = 0 n_attacks_landed: int = 0 n_attacks_total: int = 0 repo_health: float = 1.0 step_count: int = 0 done: bool = False class OpsguardEnvironment(Environment): SUPPORTS_CONCURRENT_SESSIONS: bool = True def __init__(self): self._db = RepoDB(_DB_PATH) if Path(_DB_PATH).exists() else None self._repo = _REPO self._curriculum = Curriculum() self._reward_cfg = RewardConfig() self._episode: _EpisodeState | None = None self._state = State(episode_id=str(uuid4()), step_count=0) self._adversary = LLMSpamAdversary(TrainableAdversaryConfig(model_path=_ADV_LORA)) if _ADV_LORA else None def reset(self, scenario_id: str | None = None, seed: int | None = None) -> OpsguardObservation: sid = scenario_id or self._curriculum.hardest_unlocked().scenario_id scenario = get_scenario(sid) rng = random.Random(seed if seed is not None else 0) if self._db is not None and self._db.count_issues(self._repo) >= scenario.n_real_issues: real = self._db.issues_for_tier(self._repo, scenario.real_tiers, scenario.n_real_issues) else: real = self._mock_issues(scenario.n_real_issues) attack_specs = synthesize_attacks(real, scenario.attack_rate, scenario.attack_tier, sid, self._repo) attack_rows = [attack_to_issue_row(s, self._repo) for s in attack_specs] queue: list[IssueRow] = real + attack_rows rng.shuffle(queue) n_total_legit = sum(1 for q in queue if not q.is_synthetic_spam) ep = _EpisodeState( scenario=scenario, queue=queue, n_total_legit=n_total_legit, n_attacks_total=len(attack_rows), ) ep._attack_meta_cache = { # type: ignore[attr-defined] s.issue_id: { "diff": s.pr_diff_preview, "files": s.pr_changed_files, "deps": s.pr_dependency_changes, } for s in attack_specs } self._episode = ep self._state = State(episode_id=str(uuid4()), step_count=0) return self._build_observation(feedback="episode started", reward=0.0) def step(self, action: OpsguardAction) -> OpsguardObservation: # type: ignore[override] if self._episode is None: raise RuntimeError("call reset() first") ep = self._episode if ep.done: raise RuntimeError("episode terminated; call reset()") ep.step_count += 1 self._state.step_count = ep.step_count if ep.pos >= len(ep.queue): return self._terminate(feedback="queue empty") current = ep.queue[ep.pos] target_id = action.target_issue_id or current.issue_id target = current if target_id == current.issue_id else self._find_in_queue(target_id) or current repeat = self._is_repeat(action, target) info_already = target.issue_id in ep.info_requested_for diff_already = target.issue_id in ep.diff_reviewed_for contributor = self._db.contributor(target.author_login) if self._db else None action_type = action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type) verdict = action.security_verdict.value if action.security_verdict else None if action_type == "query_history": hits = self._handle_query(action.query or current.title) ep.last_memory_hits = hits feedback = f"retrieved {len(hits)} prior issues" rb = grade_step( action_type=action_type, issue=target, predicted_labels=None, predicted_assignee=None, predicted_duplicate_of=None, predicted_security_verdict=verdict, predicted_revoke_login=action.revoke_target_login, contributor=contributor, repeat=repeat, info_request_was_already_made=info_already, diff_already_reviewed=diff_already, cfg=self._reward_cfg, ) elif action_type in ("review_pr_diff", "check_dependency", "revoke_trust"): ep.last_memory_hits = [] rb = grade_step( action_type=action_type, issue=target, predicted_labels=[action.label] if action.label else [], predicted_assignee=action.assignee_login, predicted_duplicate_of=action.duplicate_of_id, predicted_security_verdict=verdict, predicted_revoke_login=action.revoke_target_login, contributor=contributor, repeat=repeat, info_request_was_already_made=info_already, diff_already_reviewed=diff_already, cfg=self._reward_cfg, ) if action_type == "review_pr_diff": ep.diff_reviewed_for.add(target.issue_id) if action_type == "revoke_trust" and action.revoke_target_login: ep.revoked_logins.add(action.revoke_target_login) feedback = f"applied {action_type} to issue#{target.number}" elif action_type == "wait": ep.last_memory_hits = [] rb = grade_step( action_type=action_type, issue=target, predicted_labels=None, predicted_assignee=None, predicted_duplicate_of=None, predicted_security_verdict=None, predicted_revoke_login=None, contributor=contributor, repeat=repeat, info_request_was_already_made=info_already, diff_already_reviewed=diff_already, cfg=self._reward_cfg, ) feedback = "wait" else: ep.last_memory_hits = [] rb = grade_step( action_type=action_type, issue=target, predicted_labels=[action.label] if action.label else [], predicted_assignee=action.assignee_login, predicted_duplicate_of=action.duplicate_of_id, predicted_security_verdict=verdict, predicted_revoke_login=action.revoke_target_login, contributor=contributor, repeat=repeat, info_request_was_already_made=info_already, diff_already_reviewed=diff_already, cfg=self._reward_cfg, ) if action_type == "request_info": ep.info_requested_for.add(target.issue_id) feedback = f"info requested on issue#{target.number}" else: feedback = f"applied {action_type} to issue#{target.number}" if target.issue_id == current.issue_id: self._update_world_state(ep, action_type, target) ep.pos += 1 ep.actions_taken.append({ "step": ep.step_count, "action": action_type, "issue_id": target.issue_id, "is_attack": target.is_synthetic_spam, "reward": round(rb.total, 3), }) ep.cumulative_reward += rb.total self._record_step_in_memory(ep, action_type, target) self._update_profile(ep, action_type, target) if ep.pos >= len(ep.queue) or ep.step_count >= ep.scenario.step_budget: return self._terminate(feedback=feedback, last_step_reward=rb.total, last_breakdown=rb.to_dict()) return self._build_observation(feedback=feedback, reward=rb.total, breakdown=rb.to_dict()) def _update_world_state(self, ep: _EpisodeState, action_type: str, target: IssueRow): if target.is_synthetic_spam: if action_type in ("close_spam", "reject_pr", "flag_security"): ep.n_attacks_caught += 1 ep.repo_health = min(1.0, ep.repo_health + 0.02) elif action_type == "merge_pr": ep.n_attacks_landed += 1 ep.repo_health = max(0.0, ep.repo_health - 0.15) else: ep.n_attacks_landed += 1 ep.repo_health = max(0.0, ep.repo_health - 0.05) else: ep.n_resolved_legit += 1 if action_type in ("close_spam", "reject_pr"): ep.repo_health = max(0.0, ep.repo_health - 0.04) else: ep.repo_health = min(1.0, ep.repo_health + 0.01) def _terminate( self, feedback: str = "done", last_step_reward: float = 0.0, last_breakdown: dict | None = None, ) -> OpsguardObservation: ep = self._episode assert ep is not None terminal = grade_terminal( n_attacks_caught=ep.n_attacks_caught, n_attacks_total=ep.n_attacks_total, n_legit_resolved=ep.n_resolved_legit, n_legit_total=ep.n_total_legit, repo_health=ep.repo_health, cumulative_step_reward=ep.cumulative_reward, ) ep.cumulative_reward += terminal.total ep.done = True normalized = ep.cumulative_reward / max(1, ep.scenario.step_budget) + terminal.total self._curriculum.report_episode(ep.scenario.scenario_id, normalized) obs = self._build_observation( feedback=f"{feedback} | terminal {terminal.total:+.2f}", reward=last_step_reward + terminal.total, breakdown={ "step": last_breakdown, "terminal": terminal.to_dict(), "cumulative": round(ep.cumulative_reward, 3), "attacks_caught": ep.n_attacks_caught, "attacks_landed": ep.n_attacks_landed, "attacks_total": ep.n_attacks_total, "legit_resolved": ep.n_resolved_legit, "legit_total": ep.n_total_legit, "repo_health_final": round(ep.repo_health, 3), "curriculum": self._curriculum.stats(), }, done=True, ) return obs def _build_observation( self, *, feedback: str, reward: float = 0.0, breakdown: dict | None = None, done: bool = False, ) -> OpsguardObservation: ep = self._episode assert ep is not None labels = self._db.labels_for_repo(self._repo) if self._db else _MOCK_LABELS current_view: IssueView | None = None if not done and ep.pos < len(ep.queue): cur = ep.queue[ep.pos] contrib = self._db.contributor(cur.author_login) if self._db else None comments = self._db.comments_for_issue(cur.issue_id, limit=4) if self._db else [] diff = "" files: list[str] = [] deps: list[dict[str, str]] = [] if cur.is_synthetic_spam: cached = self._cached_attack_meta(cur.issue_id) if cached: diff = cached.get("diff", "") files = cached.get("files", []) deps = cached.get("deps", []) profile_features = ep.profiles.features_for_observation(cur.author_login) or {} behavior = analyze_diff(diff, files, deps) if (diff or deps) else None current_view = IssueView( issue_id=cur.issue_id, number=cur.number, title=cur.title, body=cur.body[:1500], is_pr=cur.is_pr, pr_diff_preview=diff, pr_changed_files=files, pr_dependency_changes=deps, author_login=cur.author_login, author_pr_count=contrib.public_pr_count if contrib else 0, author_account_age_days=contrib.account_age_days if contrib else 0, author_first_contribution_days_ago=contrib.account_age_days if contrib else 0, available_labels=labels[:30], comments_preview=[ {"author": c.author_login, "body": c.body[:200]} for c in comments ], author_profile=profile_features, behavior_tokens=[t.value if hasattr(t, "value") else str(t) for t in (behavior.tokens if behavior else [])], behavior_risk_score=behavior.risk_score if behavior else 0.0, behavior_hints=behavior.reasoning_hints if behavior else [], ) return OpsguardObservation( scenario_id=ep.scenario.scenario_id, step=ep.step_count, step_budget=ep.scenario.step_budget, queue_position=ep.pos, queue_total=len(ep.queue), current_issue=current_view, memory_hits=ep.last_memory_hits, recent_actions=ep.actions_taken[-6:], repo_health=round(ep.repo_health, 3), n_attacks_landed=ep.n_attacks_landed, feedback=feedback, done=done, reward=reward, metadata=breakdown or {}, ) def _cached_attack_meta(self, issue_id: int) -> dict | None: # Fast path: stash attack-specs alongside queue rows. We rebuild here from queue order # by issue_id since IssueRow doesn't carry diff fields directly. ep = self._episode if ep is None: return None meta = getattr(ep, "_attack_meta_cache", None) if meta is None: return None return meta.get(issue_id) def _find_in_queue(self, issue_id: int) -> IssueRow | None: ep = self._episode assert ep is not None for it in ep.queue[ep.pos:ep.pos + 5]: if it.issue_id == issue_id: return it return None def _handle_query(self, query: str) -> list[dict[str, Any]]: ep = self._episode assert ep is not None # Ablation hook: when OPSGUARD_MEMORY_DISABLED=1 the env returns no # memory hits (and skips repo-history search). Used by # scripts/memory_ablation.py to measure causal impact of memory. if os.environ.get("OPSGUARD_MEMORY_DISABLED", "").strip() in ("1", "true", "True"): return [] out: list[dict[str, Any]] = [] mem_hits = ep.memory.query(query_text=query, top_k=4) for h in mem_hits: out.append({ "source": "episode_memory", "memory_id": h.memory_id, "step": getattr(h, "step", None), "content": (getattr(h, "content", None) or getattr(h, "rule", ""))[:200], }) if self._db: for h in self._db.search_history(self._repo, query, limit=4): out.append({ "source": "repo_history", "issue_id": h.issue_id, "number": h.number, "title": h.title[:120], "truth_action": h.truth_action, "truth_labels": h.truth_labels, }) return out def _update_profile(self, ep: _EpisodeState, action_type: str, target: IssueRow): try: outcome = None if action_type == "merge_pr": outcome = "merged" elif action_type in ("reject_pr", "close_spam"): outcome = "rejected" elif action_type == "flag_security": outcome = "flagged" cached = self._cached_attack_meta(target.issue_id) or {} diff = cached.get("diff", "") files = cached.get("files", []) ep.profiles.observe( login=target.author_login, step=ep.step_count, is_pr=target.is_pr, diff_size=len(diff) if diff else len(target.body or ""), files=files, outcome=outcome, ) except Exception: pass def _record_step_in_memory(self, ep: _EpisodeState, action_type: str, target: IssueRow): try: content = ( f"step={ep.step_count} action={action_type} " f"author={target.author_login} is_pr={target.is_pr} " f"is_attack={target.is_synthetic_spam} " f"title={target.title[:80]}" ) cues = [target.author_login, action_type] if target.is_pr: cues.append("pr") if target.is_synthetic_spam: cues.append(f"attack_{target.spam_pattern}") ep.memory.tick(ep.step_count) ep.memory.write_episode( step=ep.step_count, content=content, cues=cues, importance=0.7 if target.is_synthetic_spam else 0.4, ) if ep.step_count % 50 == 0: ep.memory.sweep_decayed() except Exception: pass def _is_repeat(self, action: OpsguardAction, target: IssueRow) -> bool: ep = self._episode assert ep is not None action_type = action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type) for prev in ep.actions_taken[-5:]: if prev["action"] == action_type and prev["issue_id"] == target.issue_id: return True return False def _mock_issues(self, n: int) -> list[IssueRow]: out = [] for i in range(n): tier = (i % 5) + 1 out.append( IssueRow( issue_id=1_000_000 + i, number=1000 + i, repo=self._repo, title=f"mock issue {i}: feature request for module {chr(65 + (i % 5))}", body=f"This is a mock issue body number {i}. Describes a small bug.", author_login=f"mock_user_{i % 7}", is_pr=(i % 4 == 0), tier=tier, is_synthetic_spam=False, spam_pattern=None, truth_action=["label", "comment", "merge_pr", "request_info", "label"][i % 5], truth_labels=[["bug"], ["enhancement"], ["docs"], ["question"], ["good first issue"]][i % 5], truth_close_reason=None, truth_assignee=None, state="open", ) ) return out @property def state(self) -> State: return self._state _MOCK_LABELS = [ "bug", "enhancement", "documentation", "good first issue", "help wanted", "question", "duplicate", "wontfix", "invalid", "spam", "security", ]