khushmagrawal/devsecops_env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8DevSecOps Environment Client.9 10HTTP/WebSocket client for connecting to the DevSecOps RL environment server.11Handles action serialization and observation deserialization.12"""13 14from typing import Dict, Any, Optional15 16from openenv.core import EnvClient17from openenv.core.client_types import StepResult18from openenv.core.env_server.types import State19 20from .models import (21    DevsecopsAction,22    DevsecopsObservation,23    PullRequest,24    RepositoryContext,25    Budget,26    ToolUseRecord,27)28 29 30class DevsecopsEnv(31    EnvClient[DevsecopsAction, DevsecopsObservation, State]32):33    """34    Client for the DevSecOps Environment.35 36    This client maintains a persistent WebSocket connection to the environment server,37    enabling efficient multi-step interactions with lower latency.38    Each client instance has its own dedicated environment session on the server.39 40    The client handles:41    - Serialization of DevsecopsAction to JSON42    - Deserialization of JSON responses into DevsecopsObservation43    - Connection management (HTTP/WebSocket)44 45    Example:46        >>> # Connect to a running server47        >>> with DevsecopsEnv(base_url="http://localhost:8000") as client:48        ...     result = client.reset()49        ...     print(f"Task: {result.observation.task_id}")50        ...51        ...     # Inspect the PR changes52        ...     action = DevsecopsAction(tool_name="inspect_diff")53        ...     result = client.step(action)54        ...     print(result.observation.last_tool_output)55        ...56        ...     # Make a decision57        ...     action = DevsecopsAction(58        ...         tool_name="make_decision",59        ...         verdict="MERGE",60        ...         justification="Docs only, no functional changes"61        ...     )62        ...     result = client.step(action)63        ...     print(f"Done: {result.done}, Reward: {result.observation.episode_reward}")64 65    Example with Docker:66        >>> # Automatically start container and connect67        >>> client = DevsecopsEnv.from_docker_image("devsecops_env:latest")68        >>> try:69        ...     result = client.reset()70        ...     # ... interact with environment ...71        ... finally:72        ...     client.close()73    """74 75    def _step_payload(self, action: DevsecopsAction) -> Dict[str, Any]:76        """77        Convert DevsecopsAction to JSON payload for step message.78 79        Serializes all non-None action fields as a dictionary.80 81        Args:82            action: DevsecopsAction instance83 84        Returns:85            Dictionary representation suitable for JSON encoding86        """87        # Serialize action dict, removing None values88        return {k: v for k, v in action.dict().items() if v is not None}89 90    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[DevsecopsObservation]:91        """92        Parse server response into StepResult[DevsecopsObservation].93 94        Reconstructs all nested Pydantic models from JSON data.95 96        Args:97            payload: JSON response data from server98 99        Returns:100            StepResult with DevsecopsObservation101        """102        obs_data = payload.get("observation", {})103        104        # Reconstruct nested Pydantic models105        pr = PullRequest(**obs_data.get("pr", {})) if obs_data.get("pr") else None106        repo = RepositoryContext(**obs_data.get("repo_context", {})) if obs_data.get("repo_context") else None107        budget = Budget(**obs_data.get("budget", {})) if obs_data.get("budget") else None108        109        # Reconstruct tool use records110        history = [111            ToolUseRecord(**record)112            for record in obs_data.get("pipeline_history", [])113        ]114        115        # Create observation116        observation = DevsecopsObservation(117            task_id=obs_data.get("task_id", ""),118            pr=pr or PullRequest(pr_id="", title="", author=""),119            repo_context=repo or RepositoryContext(repo_name=""),120            pipeline_history=history,121            budget=budget or Budget(),122            last_tool_output=obs_data.get("last_tool_output"),123            done=payload.get("done", False),124            reward=payload.get("reward", 0.0),125            episode_reward=obs_data.get("episode_reward", 0.0),126            step_count=obs_data.get("step_count", 0),127            internal_state=obs_data.get("internal_state", {}),128            metadata=obs_data.get("metadata", {}),129        )130 131        return StepResult(132            observation=observation,133            reward=payload.get("reward", 0.0),134            done=payload.get("done", False),135        )136 137    def _parse_state(self, payload: Dict[str, Any]) -> State:138        """139        Parse server response into State object.140 141        Args:142            payload: JSON response from state request143 144        Returns:145            State object with episode_id and step_count146        """147        return State(148            episode_id=payload.get("episode_id", ""),149            step_count=payload.get("step_count", 0),150        )151