# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. """FixOS environment implementation.""" from dataclasses import dataclass from uuid import uuid4 from typing import Any, Callable, Dict, List from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import State try: from ..models import FileInfo, FixOSAction, FixOSObservation, LogEntry, ProcessInfo, ServiceInfo except ImportError: from models import FileInfo, FixOSAction, FixOSObservation, LogEntry, ProcessInfo, ServiceInfo # --------------------------- # SAFE SCORE FUNCTION # --------------------------- def fix_score(score: float) -> float: # strictly (0,1) return max(0.01, min(0.99, float(score))) # --------------------------- # TASK DEFINITION # --------------------------- @dataclass(frozen=True) class TaskDefinition: task_id: str difficulty: str failures: Dict[str, Any] optimal_steps: int success_check: Callable[[Dict[str, Any]], bool] # --------------------------- # TASK ORDER # --------------------------- TASK_ORDER: List[str] = [ "easy_1", "easy_2", "medium_1", "medium_2", "hard_1", "hard_2", "hard_3", ] # --------------------------- # TASKS (REQUIRED — NO {...}) # --------------------------- TASKS = { "easy_1": TaskDefinition( task_id="easy_1", difficulty="easy", failures={"nginx_status": "stopped"}, optimal_steps=3, success_check=lambda s: s["nginx_status"] == "running", ), "easy_2": TaskDefinition( task_id="easy_2", difficulty="easy", failures={"mysql_status": "stopped"}, optimal_steps=3, success_check=lambda s: s["mysql_status"] == "running", ), "medium_1": TaskDefinition( task_id="medium_1", difficulty="medium", failures={"nginx_config_valid": False}, optimal_steps=5, success_check=lambda s: s["nginx_config_valid"] and s["nginx_status"] == "running", ), "medium_2": TaskDefinition( task_id="medium_2", difficulty="medium", failures={"mysql_config_valid": False}, optimal_steps=5, success_check=lambda s: s["mysql_config_valid"] and s["mysql_status"] == "running", ), "hard_1": TaskDefinition( task_id="hard_1", difficulty="hard", failures={"disk_percent": 98.0}, optimal_steps=7, success_check=lambda s: s["disk_percent"] < 90, ), "hard_2": TaskDefinition( task_id="hard_2", difficulty="hard", failures={"nginx_status": "stopped", "nginx_config_valid": False}, optimal_steps=8, success_check=lambda s: s["nginx_status"] == "running" and s["nginx_config_valid"], ), "hard_3": TaskDefinition( task_id="hard_3", difficulty="hard", failures={"mysql_status": "stopped", "mysql_config_valid": False}, optimal_steps=8, success_check=lambda s: s["mysql_status"] == "running" and s["mysql_config_valid"], ), } # --------------------------- # BASE SCORE FUNCTION # --------------------------- def _task_score_base(state: Dict[str, Any]) -> float: score = 0.0 if state["disk_percent"] <= 95: score += 0.2 if state["nginx_status"] == "running": score += 0.2 if state["mysql_status"] == "running": score += 0.2 if state["nginx_config_valid"]: score += 0.2 if state["mysql_config_valid"]: score += 0.2 return fix_score(score) # --------------------------- # ENVIRONMENT CLASS # --------------------------- class FixOSEnvironment(Environment): SUPPORTS_CONCURRENT_SESSIONS: bool = True MAX_STEPS_PER_EPISODE: int = 50 def __init__(self): self._state = State(episode_id=str(uuid4()), step_count=0) self._current_task = None self._task_pointer = -1 self._env_state: Dict[str, Any] = {} self._previous_score = 0.0 # --------------------------- # RESET # --------------------------- def reset(self): self._state = State(episode_id=str(uuid4()), step_count=0) self._task_pointer = (self._task_pointer + 1) % len(TASK_ORDER) self._current_task = TASKS[TASK_ORDER[self._task_pointer]] self._env_state = { "task_id": self._current_task.task_id, "difficulty": self._current_task.difficulty, "nginx_status": "running", "mysql_status": "running", "nginx_config_valid": True, "mysql_config_valid": True, "disk_percent": 45.0, "processes": [], "files": {}, "logs": [], "history": [], "logs_checked": False, "config_checked": {"nginx": False, "mysql": False}, } self._apply_failures(self._current_task.failures) self._refresh_resources() # FIXED self._previous_score = fix_score(self._compute_task_score()) return self._build_observation("[INFO] Environment reset") # --------------------------- # STEP # --------------------------- def step(self, action): self._state.step_count += 1 reward = 0.0 current_score = fix_score(self._compute_task_score()) delta = max(0.0, current_score - self._previous_score) if delta > 0: reward += 0.4 * delta self._previous_score = current_score is_success = self._current_task.success_check(self._env_state) if is_success: reward += 0.5 done = is_success or self._state.step_count >= self.MAX_STEPS_PER_EPISODE reward = fix_score(reward) return self._build_observation( command_output="step executed", reward=reward, done=done, is_success_step=is_success, task_score=current_score, # already fixed ) # --------------------------- # SCORE COMPUTATION # --------------------------- def _compute_task_score(self) -> float: s = self._env_state task_id = self._current_task.task_id if task_id.startswith("easy"): raw = ( (0.7 if ((s["nginx_status"] == "running" and task_id == "easy_1") or (s["mysql_status"] == "running" and task_id == "easy_2")) else 0.0) + (0.3 if s["disk_percent"] < 95 else 0.0) ) return fix_score(raw) score = 0.0 if task_id == "medium_1": score += 0.2 if s["logs_checked"] else 0.0 score += 0.3 if s["nginx_config_valid"] else 0.0 score += 0.5 if s["nginx_status"] == "running" else 0.0 elif task_id == "medium_2": score += 0.2 if s["logs_checked"] else 0.0 score += 0.3 if s["mysql_config_valid"] else 0.0 score += 0.5 if s["mysql_status"] == "running" else 0.0 else: return _task_score_base(s) return fix_score(score) # --------------------------- # HELPERS # --------------------------- def _apply_failures(self, failures): for k, v in failures.items(): self._env_state[k] = v def _refresh_resources(self): self._env_state["cpu_percent"] = 50.0 self._env_state["memory_mb"] = 500.0 def _build_observation( self, command_output, reward=0.0, done=False, is_success_step=False, task_score=None, ): return FixOSObservation( command_output=command_output, processes=[], services=[], filesystem=[], resources={}, logs=[], history=[], task_id=self._env_state["task_id"], task_difficulty=self._env_state["difficulty"], task_score=fix_score( self._compute_task_score() if task_score is None else task_score ), is_success_step=is_success_step, remaining_steps=max(0, self.MAX_STEPS_PER_EPISODE - self._state.step_count), done=done, reward=fix_score(reward), ) @property def state(self): return self._state