server/toolforge_env_environment.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Toolforge Env Environment Implementation.

A simple test environment that echoes back messages sent to it.
Perfect for testing HTTP server infrastructure.
"""

import logging
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import EnvironmentMetadata, State
from typing import Any, Dict, List, Optional

try:
    from .inputs.base import InputProvider
    from .inputs.simulated.task_selector import TaskSelector
    from .inputs.factory import create_input_provider
except ImportError:
    from server.inputs.base import InputProvider
    from server.inputs.simulated.task_selector import TaskSelector
    from server.inputs.factory import create_input_provider

try:
    from ..models import (
        Tool,
        ToolCall,
        ToolforgeAction,
        ToolforgeObservation,
        ToolForgeState,
        Task,
        EpisodeGradingState,
    )
except ImportError:
    from models import (
        Tool,
        ToolCall,
        ToolforgeAction,
        ToolforgeObservation,
        ToolForgeState,
        Task,
        EpisodeGradingState,
    )

try:
    from server.tools import build_atomic_tools
    from server.evaluation_pipeline import run_evaluation_pipeline
    from server.plan_evaluator import update_sequence_counts
except ImportError:
    from .tools import build_atomic_tools
    from .evaluation_pipeline import run_evaluation_pipeline
    from .plan_evaluator import update_sequence_counts


logger = logging.getLogger(__name__)


class ToolforgeEnvironment(Environment):
    """
    A simple echo environment that echoes back messages.

    This environment is designed for testing the HTTP server infrastructure.
    It maintains minimal state and simply echoes back whatever message it receives.

    Example:
        >>> env = ToolforgeEnvironment()
        >>> obs = env.reset()
        >>> print(obs.echoed_message)  # "Toolforge Env environment ready!"
        >>>
        >>> obs = env.step(ToolforgeAction(message="Hello"))
        >>> print(obs.echoed_message)  # "Hello"
        >>> print(obs.message_length)  # 5
    """

    # Enable concurrent WebSocket sessions.
    # Set to True if your environment isolates state between instances.
    # When True, multiple WebSocket clients can connect simultaneously, each
    # getting their own environment instance (when using factory mode in app.py).
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    # Hard guardrail to ensure episodes terminate deterministically.
    MAX_EPISODE_STEPS: int = 100

    def __init__(self):
        """Initialize the toolforge_env environment."""
        super().__init__(transform=None, rubric=None)
        self._state = self._create_default_state()
        self._reset_count = 0

        self._task_selector = TaskSelector()
        self._input_provider_factory = create_input_provider
        self._input_provider: Optional[InputProvider] = None
        self._last_approval: Optional[bool] = None

        # persistent config
        self.mode = None
        self.difficulty: str = "easy"
        self.initialized = False

    def _create_default_state(self) -> ToolForgeState:
        """
        Create a default ToolForgeState with basic parameters.

        Returns:
            ToolForgeState with default values
        """
        default_task = Task(
            id="default-task",
            prompt="Default task",
            difficulty="easy",
            required_slots=[],
            baseline_token_cost=0,
        )

        return ToolForgeState(
            episode_id=str(uuid4()),
            step_count=0,
            current_task=default_task,
            task_queue=[],
            completed_tasks=[],
            available_tools=[],
            accepted_macros=[],
            rejected_macro_count=0,
            call_history=[],
            tokens_used=0,
            done=False,
        )

    def _create_default_observation(self) -> ToolforgeObservation:
        """
        Create a default ToolforgeObservation with basic parameters.

        Returns:
            ToolforgeObservation with default values
        """
        return ToolforgeObservation(
            current_task=self._state.current_task,
            available_tools=self._available_tools_to_prompt_specs(self._state.available_tools),
        )

    def _tool_to_prompt_spec(self, tool: Tool) -> Dict[str, Any]:
        """Convert a Tool model into a plain prompt-friendly dictionary."""

        return {
            "name": tool.name,
            "description": tool.description,
            "is_macro": tool.is_macro,
            "steps": tool.steps or [],
        }

    def _available_tools_to_prompt_specs(self, tools: List[Tool]) -> List[Dict[str, Any]]:
        """Convert available tools to plain dictionaries ready for prompt injection.
        Macros are listed first, followed by atomic tools.
        """
        sorted_tools = sorted(tools, key=lambda t: not t.is_macro)
        return [self._tool_to_prompt_spec(tool) for tool in sorted_tools]

    def reset(        
            self, 
            seed: Optional[int] = None, 
            episode_id: Optional[str] = None, 
            task_id: Optional[str] = None, 
            **kwargs
        ) -> ToolforgeObservation:
        """
        Reset the environment.

        Returns:
            ToolforgeObservation with a ready message
        """
        ep_id = episode_id if episode_id is not None else str(uuid4())

        # old state preservation
        prev_accepted_macros = self._state.accepted_macros if hasattr(self, '_state') else []
        prev_macro_defs = self._state.macro_definitions if hasattr(self, '_state') else {}
        prev_rejected_count = self._state.rejected_macro_count if hasattr(self, '_state') else 0

        self._state = self._create_default_state()
        self._state.episode_id = ep_id
        self._reset_count += 1
        resolved_task_id = task_id or kwargs.get("task_id", "easy")

        # subsequent resets ignore incoming values
        task_list = self._task_selector.next_task_list(resolved_task_id)

        self._input_provider = self._input_provider_factory(task_list)
        first_task = self._get_next_task_from_generator()
        self._state.current_task = first_task
        self._sync_task_queue_from_generator()
        self._state.completed_tasks = []
        self._state.available_tools = build_atomic_tools() + prev_accepted_macros
        self._state.accepted_macros = prev_accepted_macros
        self._state.rejected_macro_count = prev_rejected_count
        self._state.call_history = []
        self._state.tokens_used = 0
        self._state.done = False
        # Reset episode-level macro tracking state
        self._state.sequence_counts = {}
        self._state.macro_usage_counts = {}
        self._state.macro_definitions = prev_macro_defs
        self._last_approval = None

        # Fresh episode-level grading accumulator
        self._state.grading = EpisodeGradingState()


        obs = self._create_default_observation()
        obs.done = False
        obs.reward = 0.0

        return obs

    def step(self, action: ToolforgeAction) -> ToolforgeObservation:  # type: ignore[override]
        """
        Execute a step in the environment by echoing the message.

        Args:
            action: ToolforgeAction containing the message to echo

        Returns:
            ToolforgeObservation with the echoed message and its length
        """
        if self._is_done():
            obs_terminal = self._create_default_observation()
            obs_terminal.done = True
            obs_terminal.reward = 0.0
            obs_terminal.metadata = {
                "summary": "Episode already completed.",
                "terminal": True,
            }
            return obs_terminal

        if not isinstance(action, ToolforgeAction):
            logger.warning(
                "Malformed action rejected. Expected ToolforgeAction, got %s",
                type(action).__name__,
            )
            self._state.step_count += 1
            self._last_approval = False

            # Track malformed actions in grading state
            self._state.grading.episode_steps += 1
            self._state.grading.validation_failures += 1

            progression = "advanced_to_next_task_malformed_action"
            if self._state.step_count >= self.MAX_EPISODE_STEPS:
                self._state.done = True
                progression = "episode_terminated_max_steps"
            else:
                advanced = self._advance_to_next_task()
                progression = "advanced_to_next_task" if advanced else "episode_completed"

            obs_malformed = self._create_default_observation()
            obs_malformed.done = self._is_done()
            obs_malformed.reward = 0.0
            obs_malformed.metadata = {
                "summary": "Malformed action treated as failed attempt.",
                "malformed_action": True,
                "plan_accepted": False,
                "task_prompt": self._state.current_task.prompt,
                "task_level": self._state.current_task.difficulty,
                "progression": progression,
            }
            return obs_malformed

        self._state.step_count += 1

        plan_accounting = self._analyze_plan(action.plan)
        step_call_count = plan_accounting["step_call_count"]
        unknown_tool_calls = plan_accounting["unknown_tool_calls"]

        self._state.call_history.extend(action.plan)
        self._state.tokens_used += step_call_count

        available_tools_by_name = {
            tool.name: tool for tool in self._state.available_tools
        }
        pipeline_result = run_evaluation_pipeline(
            plan=action.plan,
            task=self._state.current_task,
            available_tools=available_tools_by_name,
            accepted_macros=self._state.accepted_macros,
            baseline_token_cost=self._state.current_task.baseline_token_cost,
            sequence_counts=self._state.sequence_counts,
            macro_definitions=self._state.macro_definitions,
            macro_proposal=action.macro_proposal,
        )
        self._last_approval = bool(pipeline_result.passed_validation)

        # Update sequence counts AFTER pipeline evaluation so recognition uses prior counts only
        update_sequence_counts(action.plan, self._state.sequence_counts)

        # Update macro usage counts for any macro tools used in this plan
        macro_names = {m.name for m in self._state.accepted_macros}
        for call in action.plan:
            if call.tool_name in macro_names:
                self._state.macro_usage_counts[call.tool_name] = (
                    self._state.macro_usage_counts.get(call.tool_name, 0) + 1
                )

        progression = "advanced_to_next_task"
        if self._state.step_count >= self.MAX_EPISODE_STEPS:
            self._state.done = True
            progression = "episode_terminated_max_steps"
        else:
            advanced = self._advance_to_next_task()
            progression = "advanced_to_next_task" if advanced else "episode_completed"

        macro_result = self._process_macro_proposal(
            action=action,
            can_accept=bool(pipeline_result.passed_validation),
            reject_reason="plan_not_accepted",
        )

        if macro_result["decision"] == "approved" and action.macro_proposal is not None:
            self._state.available_tools.append(action.macro_proposal)

        # Accumulate grading state from this step
        self._update_grading_state(pipeline_result, macro_result, action)

        # Fetch the scalar reward from the evaluation pipeline
        reward = float(pipeline_result.reward)
        print(self._state.current_task.prompt, reward, progression, macro_result)
        return ToolforgeObservation(
            current_task=self._state.current_task,
            available_tools=self._available_tools_to_prompt_specs(self._state.available_tools),
            done=self._is_done(),
            reward=reward
        )

    def _advance_to_next_task(self) -> bool:
        """Advance from current task to the next queued task."""

        completed_task_id = self._state.current_task.id
        self._state.completed_tasks.append(self._state.current_task)

        if self._input_provider is None or self._input_provider.is_done():
            self._state.done = True
            self._state.task_queue = []
            logger.info(
                "Episode complete. Final task '%s' finished; no tasks remain.",
                completed_task_id,
            )
            return False

        next_task = self._get_next_task_from_generator()
        self._state.current_task = next_task
        print("Prompt change")
        self._sync_task_queue_from_generator()
        logger.info(
            "Task advanced from '%s' to '%s'. Remaining tasks=%d",
            completed_task_id,
            next_task.id,
            len(self._state.task_queue),
        )
        logger.debug("Next task prompt: %s", next_task.prompt)
        return True

    def _update_grading_state(
        self,
        pipeline_result,
        macro_result: Dict[str, Any],
        action: ToolforgeAction,
    ) -> None:
        """Accumulate episode-level grading counters from a single step."""
        g = self._state.grading
        g.episode_steps += 1

        # Validation failures
        if not pipeline_result.passed_validation:
            g.validation_failures += 1

        # Harmful calls
        if pipeline_result.step_harmful:
            g.harmful_plan_count += 1

        # Correct plans (full slot fill + valid)
        if pipeline_result.step_task_complete and pipeline_result.passed_validation:
            g.correct_plan_count += 1

        # Efficiency tracking (only when slot_ratio == 1.0)
        sr = pipeline_result.step_slot_ratio
        if sr is not None and sr >= 1.0:
            g.fully_correct_efficiency_opportunities += 1
            g.sum_efficiency_score += (pipeline_result.step_efficiency_score or 0.0)

        # Macro creation tracking
        if macro_result["attempted"]:
            g.macro_creation_attempts += 1
            if macro_result["decision"] == "approved":
                g.macro_creation_approved += 1
                if pipeline_result.passed_validation:
                    g.macro_creation_correct += 1
                g.macro_creation_bonus_total += (pipeline_result.step_macro_creation_bonus or 0.0)
            elif macro_result["decision"] == "rejected":
                g.macro_rejected_count += 1

        # Macro usage tracking
        macro_names = {m.name for m in self._state.accepted_macros}
        if any(c.tool_name in macro_names for c in action.plan):
            g.macro_usage_attempts += 1
            if sr is not None and sr >= 0.65:
                g.macro_usage_correct += 1

        # Keep completed task count current
        g.final_completed_tasks = len(self._state.completed_tasks)


    def _analyze_plan(self, plan: List[ToolCall]) -> Dict[str, Any]:
        """Compute deterministic call accounting from server-side tool registry."""

        available_tools_by_name = {
            tool.name: tool for tool in self._state.available_tools
        }
        unknown_tool_calls: List[str] = []

        for call in plan:
            tool_def = available_tools_by_name.get(call.tool_name)
            if tool_def is None:
                unknown_tool_calls.append(call.tool_name)

        return {
            "step_call_count": len(plan),
            "unknown_tool_calls": unknown_tool_calls,
        }

    def _process_macro_proposal(
        self,
        action: ToolforgeAction,
        can_accept: bool,
        reject_reason: str,
    ) -> Dict[str, Any]:
        """Evaluate and apply macro proposal lifecycle for this step."""

        result: Dict[str, Any] = {
            "attempted": False,
            "decision": "none",
            "name": None,
            "reason": "no_macro_proposal",
        }

        has_macro_intent = (
            action.action_type == "propose_plan_with_macro"
            or action.macro_proposal is not None
        )
        if not has_macro_intent:
            return result

        result["attempted"] = True

        if action.action_type != "propose_plan_with_macro":
            return self._reject_macro(
                result=result,
                name=action.macro_proposal.name if action.macro_proposal else None,
                reason="macro_proposal_requires_propose_plan_with_macro_action_type",
            )

        if action.macro_proposal is None:
            return self._reject_macro(
                result=result,
                name=None,
                reason="missing_macro_proposal_payload",
            )

        proposal = action.macro_proposal
        macro_name = proposal.name.strip()

        if not can_accept:
            return self._reject_macro(
                result=result,
                name=macro_name,
                reason=reject_reason,
            )

        if not macro_name:
            return self._reject_macro(
                result=result,
                name=None,
                reason="macro_name_cannot_be_empty",
            )

        if not proposal.steps:
            return self._reject_macro(
                result=result,
                name=None,
                reason="macro_steps_cannot_be_empty",
            )


        existing_names = {tool.name for tool in self._state.available_tools}
        if macro_name in existing_names:
            return self._reject_macro(
                result=result,
                name=macro_name,
                reason="macro_name_already_exists",
            )

        if len(proposal.steps) < 2:
            return self._reject_macro(
                result=result,
                name=macro_name,
                reason="macro_requires_at_least_two_steps",
            )

        available_tools_by_name = {
            tool.name: tool for tool in self._state.available_tools
        }

        missing_steps = [
            call.tool_name
            for call in proposal.steps
            if call.tool_name not in available_tools_by_name
        ]
        if missing_steps:
            return self._reject_macro(
                result=result,
                name=macro_name,
                reason=f"macro_contains_unknown_tools:{','.join(missing_steps)}",
            )

        if any(
            available_tools_by_name[call.tool_name].is_macro
            for call in proposal.steps
        ):
            return self._reject_macro(
                result=result,
                name=macro_name,
                reason="nested_macro_steps_not_supported",
            )

        composed_of: List[str] = [call.tool_name for call in proposal.steps]

        macro_tool = Tool(
            name=macro_name,
            description=proposal.description.strip() or f"Macro: {' -> '.join(composed_of)}",
            is_macro=True,
            steps=proposal.steps,
        )

        self._state.accepted_macros.append(macro_tool)
        self._state.available_tools.append(macro_tool)

        # Store macro definition for sequence-based recognition tracking
        self._state.macro_definitions[macro_name] = composed_of

        result["decision"] = "approved"
        result["name"] = macro_name
        result["reason"] = "macro_registered"
        logger.info(
            "Macro approved: name='%s', steps=%s",
            macro_name,
            composed_of,
        )
        return result

    def _reject_macro(
        self,
        result: Dict[str, Any],
        name: Optional[str],
        reason: str,
    ) -> Dict[str, Any]:
        """Record macro rejection and return standardized metadata payload."""

        self._state.rejected_macro_count += 1
        result["decision"] = "rejected"
        result["name"] = name
        result["reason"] = reason
        logger.info("Macro rejected: name='%s', reason='%s'", name, reason)
        return result

    def _get_next_task_from_generator(self) -> Task:
        """Return next task from generator supporting get() or get_input()."""
        print("Prompt change")
        if self._input_provider is None:
            raise RuntimeError("Task data generator not initialized. Call reset() first.")

        return self._input_provider.get_input()

    def _sync_task_queue_from_generator(self) -> None:
        """Best-effort sync of remaining tasks for state visibility."""

        if self._input_provider is None:
            self._state.task_queue = []
            return

        if hasattr(self._input_provider, "data") and hasattr(self._input_provider, "idx"):
            data = getattr(self._input_provider, "data")
            idx = getattr(self._input_provider, "idx")
            self._state.task_queue = list(data[idx:])
            return

        # For generic providers without index access, keep queue opaque.
        self._state.task_queue = []

    def _is_done(self) -> bool:
        """Return True when the current episode should terminate."""

        return bool(self._state.done) or self._state.step_count >= self.MAX_EPISODE_STEPS

    @property
    def state(self) -> State:
        """
        Get the current environment state.

        Returns:
            Current State with episode_id and step_count
        """
        return self._state

    def get_metadata(self) -> EnvironmentMetadata:
        """Return descriptive metadata about this environment."""

        return EnvironmentMetadata(
            name="toolforge_env",
            description=(
                "A DevOps benchmark where an LLM agent learns to identify "
                "recurring tool-call patterns and compose them into reusable "
                "macro tools to minimise token consumption."
            ),
            version="0.1.0",
        )