Upload folder using huggingface_hub

supportdesk_env/server/app.py

@@ -3,24 +3,22 @@
 from __future__ import annotations
 
 import os
-import threading
 from typing import Any
 
 import uvicorn
-from fastapi import Body
-from fastapi.routing import APIRoute
 
 try:
-    from openenv.core.env_server.http_server import create_app
-    from openenv.core.env_server.types import ResetRequest, ResetResponse, StepRequest, StepResponse
+    from openenv.core.env_server import http_server as openenv_http_server
 except ImportError:  # pragma: no cover - package name differs across releases
-    from openenv_core.env_server.http_server import create_app
-    from openenv_core.env_server.types import ResetRequest, ResetResponse, StepRequest, StepResponse
+    from openenv_core.env_server import http_server as openenv_http_server
 
 from supportdesk_env.models import SupportDeskAction, SupportDeskObservation, SupportDeskState
 from supportdesk_env.server.supportdesk_environment import SupportDeskEnvironment
 from supportdesk_env.tasks import TASKS
 
+openenv_http_server.State = SupportDeskState
+create_app = openenv_http_server.create_app
+
 app = create_app(
     SupportDeskEnvironment,
     action_cls=SupportDeskAction,
@@ -28,82 +26,6 @@ app = create_app(
     env_name="supportdesk_env",
 )
 
-_http_session_lock = threading.RLock()
-_http_session_env: SupportDeskEnvironment | None = None
-
-
-def _remove_stateless_http_routes() -> None:
-    """Replace OpenEnv's stateless HTTP simulation routes with a persistent session."""
-
-    kept_routes = []
-    for route in app.router.routes:
-        if isinstance(route, APIRoute):
-            methods = route.methods or set()
-            if route.path == "/reset" and "POST" in methods:
-                continue
-            if route.path == "/step" and "POST" in methods:
-                continue
-            if route.path == "/state" and "GET" in methods:
-                continue
-        kept_routes.append(route)
-    app.router.routes = kept_routes
-    app.openapi_schema = None
-
-
-def _get_http_session_env() -> SupportDeskEnvironment:
-    global _http_session_env
-    if _http_session_env is None:
-        _http_session_env = SupportDeskEnvironment()
-    return _http_session_env
-
-
-def _serialize_observation(observation: SupportDeskObservation) -> dict[str, Any]:
-    return observation.model_dump()
-
-
-_remove_stateless_http_routes()
-
-
-@app.post("/reset", response_model=ResetResponse, tags=["Environment Control"])
-def reset_environment(
-    request: ResetRequest = Body(default_factory=ResetRequest),
-) -> ResetResponse:
-    """Keep a persistent HTTP environment so reset/step/state share the same case."""
-
-    with _http_session_lock:
-        env = _get_http_session_env()
-        observation = env.reset(seed=request.seed, episode_id=request.episode_id)
-        reward = float(observation.reward) if observation.reward is not None else None
-        return ResetResponse(
-            observation=_serialize_observation(observation),
-            reward=reward,
-            done=observation.done,
-        )
-
-
-@app.post("/step", response_model=StepResponse, tags=["Environment Control"])
-def step_environment(request: StepRequest) -> StepResponse:
-    """Advance the current HTTP session instead of stepping a fresh env instance."""
-
-    with _http_session_lock:
-        env = _get_http_session_env()
-        action = SupportDeskAction.model_validate(request.action)
-        observation = env.step(action, timeout_s=request.timeout_s)
-        reward = float(observation.reward) if observation.reward is not None else None
-        return StepResponse(
-            observation=_serialize_observation(observation),
-            reward=reward,
-            done=observation.done,
-        )
-
-
-@app.get("/state", response_model=SupportDeskState, tags=["State Management"])
-def get_environment_state() -> SupportDeskState:
-    """Return the real current HTTP session state for grader-style inspection."""
-
-    with _http_session_lock:
-        return _get_http_session_env().state
-
 
 @app.get("/tasks")
 def list_tasks() -> dict[str, Any]:

supportdesk_env/server/supportdesk_environment.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+import threading
 import uuid
 from pathlib import Path
 
@@ -31,6 +32,18 @@ class SupportDeskEnvironment(
 ):
     """A realistic customer support triage environment with dense rewards."""
 
+    _state_lock = threading.RLock()
+    _shared_task_id: str | None = None
+    _shared_step_count = 0
+    _shared_reward_total = 0.0
+    _shared_done = False
+    _shared_last_feedback = ""
+    _shared_history: list[ActionHistoryEntry] = []
+    _shared_case = SupportCaseProgress()
+    _shared_episode_id: str | None = None
+    _shared_score = 0.0
+    _shared_completed_milestones: list[str] = []
+
     def __init__(self, task_id: str | None = None):
         super().__init__()
         requested_task = task_id or os.getenv("SUPPORTDESK_TASK_ID") or list_task_ids()[0]
@@ -46,23 +59,80 @@ class SupportDeskEnvironment(
         initial_grade = grade_case(self.task, self._case)
         self._score = initial_grade.total_score
         self._completed_milestones = list(initial_grade.completed_milestones)
+        self._ensure_shared_state(self.task)
+
+    @classmethod
+    def _initialize_shared_state(
+        cls,
+        task: SupportTaskSpec,
+        *,
+        episode_id: str | None = None,
+    ) -> None:
+        initial_case = SupportCaseProgress()
+        initial_grade = grade_case(task, initial_case)
+        cls._shared_task_id = task.task_id
+        cls._shared_step_count = 0
+        cls._shared_reward_total = 0.0
+        cls._shared_done = False
+        cls._shared_last_feedback = (
+            "New case loaded. Review the ticket and policy snippets before acting."
+        )
+        cls._shared_history = []
+        cls._shared_case = initial_case
+        cls._shared_episode_id = episode_id
+        cls._shared_score = initial_grade.total_score
+        cls._shared_completed_milestones = list(initial_grade.completed_milestones)
+
+    @classmethod
+    def _ensure_shared_state(cls, task: SupportTaskSpec) -> None:
+        with cls._state_lock:
+            if cls._shared_task_id is None:
+                cls._initialize_shared_state(task)
+
+    def _sync_from_shared(self) -> None:
+        task = get_task(self.__class__._shared_task_id or self.task.task_id)
+        self.task = task
+        self._max_steps = task.max_steps
+        self._step_count = self.__class__._shared_step_count
+        self._reward_total = self.__class__._shared_reward_total
+        self._done = self.__class__._shared_done
+        self._last_feedback = self.__class__._shared_last_feedback
+        self._history = [entry.model_copy(deep=True) for entry in self.__class__._shared_history]
+        self._case = self.__class__._shared_case.model_copy(deep=True)
+        self._episode_id = self.__class__._shared_episode_id
+        self._score = self.__class__._shared_score
+        self._completed_milestones = list(self.__class__._shared_completed_milestones)
+
+    def _sync_to_shared(self) -> None:
+        self.__class__._shared_task_id = self.task.task_id
+        self.__class__._shared_step_count = self._step_count
+        self.__class__._shared_reward_total = self._reward_total
+        self.__class__._shared_done = self._done
+        self.__class__._shared_last_feedback = self._last_feedback
+        self.__class__._shared_history = [entry.model_copy(deep=True) for entry in self._history]
+        self.__class__._shared_case = self._case.model_copy(deep=True)
+        self.__class__._shared_episode_id = self._episode_id
+        self.__class__._shared_score = self._score
+        self.__class__._shared_completed_milestones = list(self._completed_milestones)
 
     @property
     def state(self) -> SupportDeskState:
-        return SupportDeskState(
-            episode_id=self._episode_id,
-            task_id=self.task.task_id,
-            difficulty=self.task.difficulty,
-            step_count=self._step_count,
-            reward=round(self._reward_total, 4),
-            done=self._done,
-            current_score=round(self._score, 4),
-            max_steps=self._max_steps,
-            case=self._case.model_copy(deep=True),
-            action_history=[entry.model_copy(deep=True) for entry in self._history],
-            completed_milestones=list(self._completed_milestones),
-            last_feedback=self._last_feedback,
-        )
+        with self.__class__._state_lock:
+            self._sync_from_shared()
+            return SupportDeskState(
+                episode_id=self._episode_id,
+                task_id=self.task.task_id,
+                difficulty=self.task.difficulty,
+                step_count=self._step_count,
+                reward=round(self._reward_total, 4),
+                done=self._done,
+                current_score=round(self._score, 4),
+                max_steps=self._max_steps,
+                case=self._case.model_copy(deep=True),
+                action_history=[entry.model_copy(deep=True) for entry in self._history],
+                completed_milestones=list(self._completed_milestones),
+                last_feedback=self._last_feedback,
+            )
 
     def reset(
         self,
@@ -70,17 +140,13 @@ class SupportDeskEnvironment(
         episode_id: str | None = None,
         **kwargs,
     ) -> SupportDeskObservation:
-        self._episode_id = episode_id or f"{self.task.task_id}-{uuid.uuid4().hex[:8]}"
-        self._step_count = 0
-        self._reward_total = 0.0
-        self._done = False
-        self._last_feedback = "New case loaded. Review the ticket and policy snippets before acting."
-        self._history = []
-        self._case = SupportCaseProgress()
-        initial_grade = grade_case(self.task, self._case)
-        self._score = initial_grade.total_score
-        self._completed_milestones = list(initial_grade.completed_milestones)
-        return self._build_observation(reward=0.0, done=False)
+        with self.__class__._state_lock:
+            self.__class__._initialize_shared_state(
+                self.task,
+                episode_id=episode_id or f"{self.task.task_id}-{uuid.uuid4().hex[:8]}",
+            )
+            self._sync_from_shared()
+            return self._build_observation(reward=0.0, done=False)
 
     def step(
         self,
@@ -88,51 +154,59 @@ class SupportDeskEnvironment(
         timeout_s: float | None = None,
         **kwargs,
     ) -> SupportDeskObservation:
-        if self._done:
-            return self._build_observation(
-                reward=-0.05,
-                done=True,
-                feedback="Episode already finished. Call reset() before taking more actions.",
-            )
-
-        previous_grade = grade_case(self.task, self._case)
-        self._apply_action(action)
-        self._step_count += 1
+        with self.__class__._state_lock:
+            self._sync_from_shared()
 
-        current_grade = grade_case(self.task, self._case)
-        reward = current_grade.total_score - previous_grade.total_score
-        reward += self._action_penalty(action, current_grade.total_score, previous_grade.total_score)
-        reward = round(reward, 4)
+            if self._done:
+                return self._build_observation(
+                    reward=-0.05,
+                    done=True,
+                    feedback="Episode already finished. Call reset() before taking more actions.",
+                )
 
-        self._score = current_grade.total_score
-        self._completed_milestones = list(current_grade.completed_milestones)
+            previous_grade = grade_case(self.task, self._case)
+            self._apply_action(action)
+            self._step_count += 1
 
-        if action.operation == "submit":
-            self._done = True
-            self._last_feedback = (
-                "Case submitted. Final deterministic grade is "
-                f"{current_grade.total_score:.2f}."
+            current_grade = grade_case(self.task, self._case)
+            reward = current_grade.total_score - previous_grade.total_score
+            reward += self._action_penalty(
+                action,
+                current_grade.total_score,
+                previous_grade.total_score,
             )
-        elif self._step_count >= self._max_steps:
-            self._done = True
-            self._last_feedback = (
-                f"Reached max steps ({self._max_steps}). Final deterministic grade is "
-                f"{current_grade.total_score:.2f}."
-            )
-        else:
-            self._last_feedback = self._build_feedback(current_grade, reward)
-
-        self._reward_total = round(self._reward_total + reward, 4)
-        self._history.append(
-            ActionHistoryEntry(
-                step=self._step_count,
-                operation=action.operation,
-                summary=self._summarize_action(action),
-                reward_delta=reward,
+            reward = round(reward, 4)
+
+            self._score = current_grade.total_score
+            self._completed_milestones = list(current_grade.completed_milestones)
+
+            if action.operation == "submit":
+                self._done = True
+                self._last_feedback = (
+                    "Case submitted. Final deterministic grade is "
+                    f"{current_grade.total_score:.2f}."
+                )
+            elif self._step_count >= self._max_steps:
+                self._done = True
+                self._last_feedback = (
+                    f"Reached max steps ({self._max_steps}). Final deterministic grade is "
+                    f"{current_grade.total_score:.2f}."
+                )
+            else:
+                self._last_feedback = self._build_feedback(current_grade, reward)
+
+            self._reward_total = round(self._reward_total + reward, 4)
+            self._history.append(
+                ActionHistoryEntry(
+                    step=self._step_count,
+                    operation=action.operation,
+                    summary=self._summarize_action(action),
+                    reward_delta=reward,
+                )
             )
-        )
+            self._sync_to_shared()
 
-        return self._build_observation(reward=reward, done=self._done)
+            return self._build_observation(reward=reward, done=self._done)
 
     def close(self) -> None:
         """No-op close hook for compatibility with local scripts."""