server/opsguard_environment.py

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",
]