org_sim/server/org_environment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# BSD-style license

"""
Organization Simulation Environment.

A multi-agent environment simulating hierarchical organization with
teams, escalation, and resource management.
"""

import uuid
from typing import Any

from openenv.core.env_server.interfaces import Environment

from ..models import OrgAction, OrgObservation, OrgState


TASK_SCENARIOS = {
    "solo_bug_fix": {
        "description": "Fix a high-priority bug in the engineering team",
        "tasks": [
            {
                "id": "task_bug",
                "team": "engineering",
                "type": "fix_bug",
                "priority": "high",
                "deadline": 10,
                "difficulty": 1,
            }
        ],
    },
    "cross_team_launch": {
        "description": "Engineering and Sales coordinate a product launch",
        "tasks": [
            {
                "id": "task_eng",
                "team": "engineering",
                "type": "implement_feature",
                "priority": "high",
                "deadline": 15,
                "difficulty": 2,
            },
            {
                "id": "task_sales",
                "team": "sales",
                "type": "prepare_proposal",
                "priority": "medium",
                "deadline": 12,
                "difficulty": 2,
            },
        ],
    },
    "startup_crisis": {
        "description": "Simultaneous incident + resource-gated feature + cross-team client meeting",
        "tasks": [
            {
                "id": "task_incident",
                "team": "engineering",
                "type": "fix_bug",
                "priority": "critical",
                "deadline": 5,
                "difficulty": 3,
            },
            {
                "id": "task_feature",
                "team": "engineering",
                "type": "implement_feature",
                "priority": "high",
                "deadline": 20,
                "difficulty": 2,
                "requires_resource": "senior_engineer",
            },
            {
                "id": "task_client",
                "team": "sales",
                "type": "client_meeting",
                "priority": "critical",
                "deadline": 8,
                "difficulty": 2,
            },
        ],
    },
}

TASK_IDS = list(TASK_SCENARIOS.keys())


class OrgSimEnvironment(Environment):
    """Organization simulation environment."""

    def __init__(self, max_steps: int = 50):
        super().__init__()
        self.max_steps = max_steps
        self._current_task_id = "cross_team_launch"

        self._state = OrgState(episode_id="", step_count=0)
        self._current_agent = "eng_swe1"
        self._tasks: dict[str, dict] = {}
        self._messages: list[dict] = []
        self._episode_resources: dict[str, bool] = {}

        self._init_organization()

    def _init_organization(self):
        """Initialize static organization structure."""
        self.teams = {
            "engineering": {
                "lead": "eng_lead",
                "members": ["eng_swe1", "eng_swe2", "eng_swe3"],
            },
            "sales": {
                "lead": "sales_mgr",
                "members": ["sales_rep1", "sales_rep2", "sales_rep3"],
            },
            "operations": {
                "lead": "ops_lead",
                "members": ["ops_analyst1", "ops_analyst2", "ops_analyst3"],
            },
        }

    def reset(self, agent_id: str = "eng_swe1", task_id: str = "cross_team_launch", **kwargs: Any) -> OrgObservation:
        """Reset environment for new episode.

        Args:
            agent_id: The agent playing this episode.
            task_id: One of 'solo_bug_fix', 'cross_team_launch', 'startup_crisis'.
        """
        self._current_task_id = task_id if task_id in TASK_SCENARIOS else "cross_team_launch"
        self._state.episode_id = str(uuid.uuid4())
        self._state.step_count = 0
        self._current_agent = agent_id
        self._messages = []

        scenario = TASK_SCENARIOS[self._current_task_id]

        self._tasks = {
            t["id"]: {
                **t,
                "status": "pending",
                "assignee": None,
                "progress": 0.0,
            }
            for t in scenario["tasks"]
        }

        # Reset shared resources per episode
        self._episode_resources = {
            "senior_engineer": True,
            "cloud_instances": True,
            "budget_approval": True,
        }

        return self._get_observation()

    def step(self, action: OrgAction) -> OrgObservation:
        """Execute action and return observation."""
        self._state.step_count += 1

        reward = 0.0

        if action.action_type == "REQUEST_TASK":
            reward = self._handle_request_task(action)
        elif action.action_type == "ACCEPT_TASK":
            reward = self._handle_accept_task(action)
        elif action.action_type == "COMPLETE_TASK":
            reward = self._handle_complete_task(action)
        elif action.action_type == "REQUEST_HELP":
            reward = self._handle_request_help(action)
        elif action.action_type == "PROVIDE_HELP":
            reward = self._handle_provide_help(action)
        elif action.action_type == "ESCALATE":
            reward = self._handle_escalate(action)
        elif action.action_type == "REQUEST_RESOURCE":
            reward = self._handle_request_resource(action)
        elif action.action_type == "REPORT_STATUS":
            reward = 0.0

        done = self._check_done()

        return self._get_observation(reward=reward, done=done)

    def _handle_request_task(self, action: OrgAction) -> float:
        my_team = self._get_team_for_agent(self._current_agent)

        available = [
            t
            for t in self._tasks.values()
            if t["team"] == my_team
            and t["status"] == "pending"
            and t.get("assignee") is None
        ]

        if available:
            task = available[0]
            task["status"] = "in_progress"
            task["assignee"] = self._current_agent
            return 0.1

        return -0.05  # No task available — slight penalty to discourage looping

    def _handle_accept_task(self, action: OrgAction) -> float:
        task_id = action.payload.get("task_id")
        if not task_id or task_id not in self._tasks:
            return -0.1

        task = self._tasks[task_id]
        if task["status"] != "pending":
            return -0.1

        task["status"] = "in_progress"
        task["assignee"] = self._current_agent
        return 0.1

    def _handle_complete_task(self, action: OrgAction) -> float:
        task_id = action.payload.get("task_id")
        if not task_id or task_id not in self._tasks:
            return -0.2

        task = self._tasks[task_id]
        if task.get("assignee") != self._current_agent:
            return -0.1

        if task["status"] != "in_progress":
            return -0.1

        # Block completion of resource-gated tasks if resource was never locked
        required = task.get("requires_resource")
        if required and self._episode_resources.get(required, True):
            # Resource still available (never locked) — penalize and don't complete
            return -0.2

        deadline = task.get("deadline", 20)
        time_taken = self._state.step_count

        if time_taken <= deadline:
            base_reward = 1.0
            time_bonus = max(0.0, (deadline - time_taken) / deadline) * 0.3
            priority_scores = {"critical": 0.3, "high": 0.2, "medium": 0.1, "low": 0.0}
            priority_bonus = priority_scores.get(task.get("priority", "medium"), 0.1)
            task["status"] = "completed"
            task["completed_at"] = self._state.step_count
            return base_reward + time_bonus + priority_bonus
        else:
            task["status"] = "failed"
            return -0.3

    def _handle_request_help(self, action: OrgAction) -> float:
        task_id = action.payload.get("task_id")
        if not task_id or task_id not in self._tasks:
            return -0.1

        task = self._tasks[task_id]
        if task.get("assignee") != self._current_agent:
            return -0.1

        task["progress"] = min(1.0, task.get("progress", 0.0) + 0.2)

        lead = self._get_team_lead(self._get_team_for_agent(self._current_agent))
        self._messages.append({
            "from": self._current_agent,
            "to": lead,
            "type": "help_request",
            "task_id": task_id,
        })

        # Scripted reply gives agent real information to act on
        self._messages.append({
            "from": lead,
            "to": self._current_agent,
            "type": "help_response",
            "task_id": task_id,
            "content": f"Acknowledged. Assigned eng_swe2 to assist with {task_id}. Check internal docs for context.",
        })

        return 0.1

    def _handle_provide_help(self, action: OrgAction) -> float:
        task_id = action.payload.get("task_id")
        if not task_id or task_id not in self._tasks:
            return -0.1

        task = self._tasks[task_id]
        task["progress"] = min(1.0, task.get("progress", 0.0) + 0.3)
        return 0.2

    def _handle_escalate(self, action: OrgAction) -> float:
        task_id = action.payload.get("task_id")
        if not task_id or task_id not in self._tasks:
            return -0.2

        task = self._tasks[task_id]

        if task["status"] in ("completed", "failed", "escalated"):
            return -0.1  # Already terminal

        # Appropriate: cross-team task OR difficulty >= 2 when stuck (progress < 0.5)
        my_team = self._get_team_for_agent(self._current_agent)
        is_cross_team = task["team"] != my_team
        is_hard_stuck = task.get("difficulty", 1) >= 2 and task.get("progress", 0.0) < 0.5

        if is_cross_team or is_hard_stuck:
            task["status"] = "escalated"
            task["escalated_at"] = self._state.step_count
            return 0.3
        else:
            return -0.2  # Premature / unnecessary escalation

    def _handle_request_resource(self, action: OrgAction) -> float:
        resource_id = action.payload.get("resource_id")

        if not resource_id or resource_id not in self._episode_resources:
            return -0.2

        if self._episode_resources[resource_id]:
            self._episode_resources[resource_id] = False  # Lock it
            return 0.2

        return -0.1  # Already locked

    def _check_done(self) -> bool:
        """Episode ends when all tasks are in a terminal state OR max_steps reached.

        Checks for terminal states (completed/failed/escalated), NOT just non-pending.
        A task in 'in_progress' is NOT done — episode must continue.
        """
        if self._state.step_count >= self.max_steps:
            return True

        active = sum(
            1 for t in self._tasks.values()
            if t["status"] in ("pending", "in_progress")
        )
        return active == 0

    def grade_episode(self) -> float:
        """Grade the completed episode. Returns 0.0–1.0.

        Called after done=True. Score varies with timing, escalation quality,
        and resource management — never identical across different trajectories.
        """
        if self._current_task_id == "solo_bug_fix":
            return self._grade_solo_bug_fix()
        elif self._current_task_id == "cross_team_launch":
            return self._grade_cross_team_launch()
        elif self._current_task_id == "startup_crisis":
            return self._grade_startup_crisis()
        return 0.0

    def _grade_solo_bug_fix(self) -> float:
        task = self._tasks.get("task_bug", {})
        if task.get("status") == "completed":
            completed_at = task.get("completed_at", self._state.step_count)
            deadline = task.get("deadline", 10)
            time_bonus = max(0.0, (deadline - completed_at) / deadline) * 0.30
            return min(1.0, 0.50 + time_bonus + 0.15)  # base + time_bonus + priority
        elif task.get("status") == "failed":
            return 0.15  # partial credit for attempting
        return 0.0

    def _grade_cross_team_launch(self) -> float:
        score = 0.0
        eng = self._tasks.get("task_eng", {})
        sales = self._tasks.get("task_sales", {})

        if eng.get("status") == "completed":
            completed_at = eng.get("completed_at", self._state.step_count)
            deadline = eng.get("deadline", 15)
            time_bonus = max(0.0, (deadline - completed_at) / deadline) * 0.20
            score += 0.40 + time_bonus  # up to 0.60

        if sales.get("status") == "escalated":
            score += 0.25  # correct action for cross-team task
        elif sales.get("status") == "completed":
            score += 0.20  # completed somehow — less expected but fine

        return min(1.0, score)

    def _grade_startup_crisis(self) -> float:
        score = 0.0
        incident = self._tasks.get("task_incident", {})
        feature = self._tasks.get("task_feature", {})
        client = self._tasks.get("task_client", {})

        # Incident: must complete by step 5
        if incident.get("status") == "completed":
            completed_at = incident.get("completed_at", self._state.step_count)
            if completed_at <= 5:
                score += 0.40
            else:
                score += 0.10  # completed but late

        # Feature: resource must have been locked (False = locked)
        resource_locked = not self._episode_resources.get("senior_engineer", True)
        if feature.get("status") == "completed" and resource_locked:
            score += 0.30
        elif feature.get("status") == "completed":
            score += 0.10  # completed without resource

        # Client: correct cross-team escalation
        if client.get("status") == "escalated":
            score += 0.20
        elif client.get("status") == "completed":
            score += 0.10

        # Resource management bonus
        if resource_locked:
            score += 0.10

        return min(1.0, score)

    def _get_observation(self, reward: float = 0.0, done: bool = False) -> OrgObservation:
        my_team = self._get_team_for_agent(self._current_agent)

        available = [
            t for t in self._tasks.values()
            if t["team"] == my_team
            and t["status"] == "pending"
            and t.get("assignee") is None
        ]

        active = next(
            (t for t in self._tasks.values()
             if t.get("assignee") == self._current_agent and t["status"] == "in_progress"),
            None,
        )

        team_status = {
            team_name: {
                "busy": sum(1 for t in self._tasks.values() if t["team"] == team_name and t.get("assignee")),
                "total": len(team_data["members"]),
            }
            for team_name, team_data in self.teams.items()
        }

        return OrgObservation(
            my_agent_id=self._current_agent,
            my_team=my_team,
            my_role="lead" if self._current_agent == self.teams[my_team]["lead"] else "member",
            available_tasks=available,
            active_task=active,
            inbox=[m for m in self._messages if m.get("to") == self._current_agent],
            team_status=team_status,
            resources=self._episode_resources.copy(),
            metrics={
                "tasks_completed": sum(1 for t in self._tasks.values() if t["status"] == "completed"),
                "tasks_failed": sum(1 for t in self._tasks.values() if t["status"] == "failed"),
                "tasks_escalated": sum(1 for t in self._tasks.values() if t["status"] == "escalated"),
                "current_task_id": self._current_task_id,
                "step_count": self._state.step_count,
            },
            done=done,
            reward=reward,
            metadata={"step_count": self._state.step_count, "task_id": self._current_task_id},
        )

    def _get_team_for_agent(self, agent_id: str) -> str:
        """Get team for agent — checks both lead and members."""
        for team_name, team_data in self.teams.items():
            if agent_id == team_data["lead"] or agent_id in team_data["members"]:
                return team_name
        return "engineering"

    def _get_team_lead(self, team_name: str) -> str:
        return self.teams.get(team_name, {}).get("lead", "")

    @property
    def state(self) -> OrgState:
        self._state.tasks = self._tasks.copy()
        self._state.team_members = {k: v["members"] for k, v in self.teams.items()}
        self._state.resource_status = self._episode_resources.copy()
        return self._state