made minor type check changes, open env has strict type checking

client.py

@@ -1,6 +1,6 @@
 """OpenEnv client wrapper for PM-Ops environment."""
 import os
-from openenv.core.client import EnvClient
+from openenv.core.env_client import EnvClient
 from models import PMOpsAction, PMOpsObservation
 
 
@@ -8,7 +8,6 @@ class PMOpsEnv(EnvClient):
     action_class = PMOpsAction
     observation_class = PMOpsObservation
 
-    def __init__(self, base_url: str = None, token: str = None):
+    def __init__(self, base_url: str | None = None, token: str | None= None):
         url = base_url or os.getenv("API_BASE_URL", "https://adityaguntur-pm-ops.hf.space")
-        tok = token or os.getenv("HF_TOKEN")
-        super().__init__(base_url=url, token=tok)
+        super().__init__(base_url=url)

models.py

@@ -46,4 +46,4 @@ class PMOpsState(State):
     codebase: Dict[str, Any] = Field(default_factory=dict)
     step_count: int = Field(default=0)
     finished: bool = Field(default=False)
-    episode_id: str = Field(default="")
+    episode_id: str | None = Field(default="")

server/pm_ops_environment.py

@@ -54,12 +54,18 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         difficulty = rng.choice(_DIFFICULTY_POOL)
         task_type = rng.choice(_TASK_TYPES)
 
+        org: Optional[Dict[str, Any]] = None
+        scenario: Optional[Dict[str, Any]] = None
+
         for attempt in range(10):
             org = generate_org_config(seed + attempt, difficulty)
             scenario = generate_scenario(task_type, org, seed + attempt)
             if _oracle_check(scenario):
                 break
 
+        if org is None or scenario is None:
+            raise RuntimeError("Failed to initialize episode state")
+
         channels = list(org["oncall_channels"].values())
         noise = org.get("noise_channels", [])
 
@@ -87,14 +93,16 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         )
 
     def step(self, action: PMOpsAction, timeout_s: Optional[float] = None, **kwargs) -> PMOpsObservation:
-        if self._ticketing is None:
+        if self._ticketing is None or self._codebase is None or self._chat is None or self._scenario is None:
             raise RuntimeError("Call reset() before step()")
 
+        scenario = self._scenario
+
         if self._done:
             return PMOpsObservation(
                 step=self._step_count,
                 max_steps=MAX_STEPS,
-                task_brief=self._scenario["brief"],
+                task_brief=scenario["brief"],
                 last_action_result={"ok": False, "error": "Episode already finished"},
                 app_state_deltas={"ticketing": [], "chat": [], "codebase": []},
                 steps_remaining=0,
@@ -122,7 +130,7 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         return PMOpsObservation(
             step=self._step_count,
             max_steps=MAX_STEPS,
-            task_brief=self._scenario["brief"],
+            task_brief=scenario["brief"],
             last_action_result=result,
             app_state_deltas=deltas,
             steps_remaining=max(0, MAX_STEPS - self._step_count),
@@ -147,6 +155,12 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         at = action.action_type
         args = action.args or {}
 
+        ticketing = self._ticketing
+        codebase = self._codebase
+        chat = self._chat
+        if ticketing is None or codebase is None or chat is None:
+            return {"ok": False, "error": "Environment not initialized. Call reset() before step()."}
+
         if at == "meta.noop":
             return {"ok": True, "data": "No operation."}
 
@@ -170,13 +184,13 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         if at.startswith("ticketing."):
             op = at.split(".", 1)[1]
             handlers = {
-                "create_ticket":     self._ticketing.create_ticket,
-                "update_ticket":     self._ticketing.update_ticket,
-                "get_ticket":        self._ticketing.get_ticket,
-                "list_tickets":      self._ticketing.list_tickets,
-                "assign_ticket":     self._ticketing.assign_ticket,
-                "comment_ticket":    self._ticketing.comment_ticket,
-                "transition_ticket": self._ticketing.transition_ticket,
+                "create_ticket":     ticketing.create_ticket,
+                "update_ticket":     ticketing.update_ticket,
+                "get_ticket":        ticketing.get_ticket,
+                "list_tickets":      ticketing.list_tickets,
+                "assign_ticket":     ticketing.assign_ticket,
+                "comment_ticket":    ticketing.comment_ticket,
+                "transition_ticket": ticketing.transition_ticket,
             }
             if op not in handlers:
                 return {"ok": False, "error": f"Unknown ticketing action: {op}"}
@@ -185,9 +199,9 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         if at.startswith("codebase."):
             op = at.split(".", 1)[1]
             handlers = {
-                "list_commits": self._codebase.list_commits,
-                "get_commit":   self._codebase.get_commit,
-                "list_prs":     self._codebase.list_prs,
+                "list_commits": codebase.list_commits,
+                "get_commit":   codebase.get_commit,
+                "list_prs":     codebase.list_prs,
             }
             if op not in handlers:
                 return {"ok": False, "error": f"Unknown codebase action: {op}"}
@@ -196,10 +210,10 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         if at.startswith("chat."):
             op = at.split(".", 1)[1]
             handlers = {
-                "post_message":  self._chat.post_message,
-                "read_channel":  self._chat.read_channel,
-                "list_channels": self._chat.list_channels,
-                "search":        self._chat.search,
+                "post_message":  chat.post_message,
+                "read_channel":  chat.read_channel,
+                "list_channels": chat.list_channels,
+                "search":        chat.search,
             }
             if op not in handlers:
                 return {"ok": False, "error": f"Unknown chat action: {op}"}
@@ -208,6 +222,9 @@ class PMOpsEnvironment(Environment[PMOpsAction, PMOpsObservation, PMOpsState]):
         return {"ok": False, "error": f"Unknown action_type: {at}"}
 
     def _grade(self) -> float:
+        if self._ticketing is None or self._chat is None or self._org_config is None or self._scenario is None:
+            return 0.0
+
         final_state = {
             "ticketing": self._ticketing.snapshot(),
             "chat": self._chat.snapshot(),