synth-ai 0.2.9.dev4__py3-none-any.whl → 0.2.9.dev6__py3-none-any.whl

examples/warming_up_to_rl/task_app/synth_envs_hosted/environment_routes.py

@@ -1,14 +1,14 @@
 from __future__ import annotations
 
-import logging
-from typing import Any, Dict, List, Optional
+import contextlib
 import json
+import logging
+from typing import Any
+from uuid import uuid4
 
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
 
-from uuid import uuid4
-
 # Import the actual classes from synth-ai
 from synth_ai.environments.examples.crafter_classic.environment import (
     CrafterClassicEnvironment,
@@ -58,9 +58,7 @@ async def validate_environment_observation(observation: Any, context: str) -> No
         "terminated",
     }
     if wordle_keys.issubset(set(observation.keys())):
-        logger.info(
-            f"🔍 ENV_ROUTES: Validating Wordle observation structure in {context}"
-        )
+        logger.info(f"🔍 ENV_ROUTES: Validating Wordle observation structure in {context}")
         logger.info(f"🔍 ENV_ROUTES: Observation keys: {list(observation.keys())}")
 
         missing_keys = wordle_keys - set(observation.keys())
@@ -100,41 +98,41 @@ async def validate_environment_observation(observation: Any, context: str) -> No
 
 class EnvCreateRequest(BaseModel):
     env_name: str
-    config: Dict[str, Any] = {}
-    seed: Optional[int] = None
-    parent_env_id: Optional[str] = None
+    config: dict[str, Any] = {}
+    seed: int | None = None
+    parent_env_id: str | None = None
     rl_run_id: str
 
 
 class EnvCreateResponse(BaseModel):
     env_id: str
-    observation: Dict[str, Any]
-    info: Optional[Dict[str, Any]] = None
+    observation: dict[str, Any]
+    info: dict[str, Any] | None = None
     step_idx: int
 
 
 class EnvResetRequest(BaseModel):
     env_id: str
-    seed: Optional[int] = None
+    seed: int | None = None
 
 
 class EnvResetResponse(BaseModel):
-    observation: Dict[str, Any]
-    info: Optional[Dict[str, Any]] = None
+    observation: dict[str, Any]
+    info: dict[str, Any] | None = None
     step_idx: int
 
 
 class EnvStepRequest(BaseModel):
     env_id: str
-    tool_calls: List[Dict[str, Any]]
+    tool_calls: list[dict[str, Any]]
 
 
 class EnvStepResponse(BaseModel):
-    observation: Dict[str, Any]
+    observation: dict[str, Any]
     done: bool
-    info: Optional[Dict[str, Any]] = None
-    reward: Optional[float] = None
-    truncated: Optional[bool] = None
+    info: dict[str, Any] | None = None
+    reward: float | None = None
+    truncated: bool | None = None
     step_idx: int
 
 
@@ -155,8 +153,8 @@ class EnvRestoreRequest(BaseModel):
 
 class EnvRestoreResponse(BaseModel):
     env_id: str
-    observation: Dict[str, Any]
-    info: Optional[Dict[str, Any]] = None
+    observation: dict[str, Any]
+    info: dict[str, Any] | None = None
     step_idx: int
 
 
@@ -215,7 +213,8 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
             # Log a world signature for sanity: seed + starting public state hash
             try:
                 pub_state = base_env.engine._get_public_state_from_env()  # type: ignore[attr-defined]
-                import hashlib, json as _json
+                import hashlib
+                import json as _json
 
                 sig_src = {
                     "player_position": list(pub_state.player_position),
@@ -270,27 +269,27 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
         elif env_name_lower == "wordle":
             # Defer imports to avoid hard dependency when not used
             try:
+                from synth_ai.environments.examples.wordle.environment import (
+                    WordleEnvironment,
+                )
                 from synth_ai.environments.examples.wordle.taskset import (
                     WordleTaskInstance,
                     WordleTaskInstanceMetadata,
                 )
-                from synth_ai.environments.examples.wordle.environment import (
-                    WordleEnvironment,
-                )
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Wordle modules unavailable: {e}"
-                )
+                ) from e
 
             # Lazy import of wrapper within branch
             try:
-                from .envs.wordle.environment import (
-                    WordleEnvironmentWrapper as _WordleWrapper,
-                )
+                from .envs.wordle.environment import WordleEnvironmentWrapper
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Wordle wrapper unavailable: {e}"
-                )
+                ) from e
+            else:
+                wordle_wrapper_cls = WordleEnvironmentWrapper
 
             cfg = request.config or {}
             word_length = int(cfg.get("word_length", 5))
@@ -307,12 +306,8 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
             )
             instance = WordleTaskInstance(
                 id=uuid4(),
-                impetus=Impetus(
-                    instructions="Play Wordle. Submit one 5-letter word per turn."
-                ),
-                intent=Intent(
-                    rubric="guess the word", gold_trajectories=None, gold_state_diff={}
-                ),
+                impetus=Impetus(instructions="Play Wordle. Submit one 5-letter word per turn."),
+                intent=Intent(rubric="guess the word", gold_trajectories=None, gold_state_diff={}),
                 metadata=md,
                 is_reproducible=True,
                 initial_engine_snapshot=None,
@@ -322,7 +317,7 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
             # Try to preserve the exact puzzle snapshot for reproducibility
             init_snap = getattr(instance, "initial_engine_snapshot", None)
 
-            wrapper = _WordleWrapper(
+            wrapper = wordle_wrapper_cls(
                 env=base_env,
                 seed=request.seed,
                 word_length=word_length,
@@ -345,9 +340,7 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
                     if key in observation_for_registry:
                         del observation_for_registry[key]
 
-            await validate_environment_observation(
-                observation_for_registry, "initialize"
-            )
+            await validate_environment_observation(observation_for_registry, "initialize")
 
             env_id = registry.register_env(
                 env=wrapper,
@@ -368,27 +361,25 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
 
         elif env_name_lower == "sokoban":
             try:
+                from synth_ai.environments.examples.sokoban.environment import (
+                    SokobanEnvironment,
+                )
                 from synth_ai.environments.examples.sokoban.taskset import (
                     SokobanTaskInstance,
                     SokobanTaskInstanceMetadata,
                 )
-                from synth_ai.environments.examples.sokoban.environment import (
-                    SokobanEnvironment,
-                )
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Sokoban modules unavailable: {e}"
-                )
+                ) from e
 
             # Lazy import of wrapper within branch
             try:
-                from .envs.sokoban.environment import (
-                    SokobanEnvironmentWrapper as _SokobanWrapper,
-                )
+                from .envs.sokoban.environment import SokobanEnvironmentWrapper
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Sokoban wrapper unavailable: {e}"
-                )
+                ) from e
 
             cfg = request.config or {}
             difficulty = cfg.get("difficulty", "easy")
@@ -411,9 +402,7 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
             )
             base_env = SokobanEnvironment(task_instance=instance)
 
-            wrapper = _SokobanWrapper(
-                env=base_env, seed=request.seed, config=cfg
-            )
+            wrapper = SokobanEnvironmentWrapper(env=base_env, seed=request.seed, config=cfg)
             result = await wrapper.initialize()
 
             # Handle the observation structure consistently for Sokoban
@@ -449,22 +438,22 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
             cfg = request.config or {}
             # Lazy import of wrapper within branch
             try:
-                from .envs.math.environment import (
-                    MathEnvironmentWrapper as _MathWrapper,
-                )
+                from .envs.math.environment import MathEnvironmentWrapper
             except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Math wrapper unavailable: {e}"
-                )
+                raise HTTPException(status_code=500, detail=f"Math wrapper unavailable: {e}") from e
 
-            wrapper = _MathWrapper(
+            wrapper = MathEnvironmentWrapper(
                 seed=request.seed,
                 problem_id=cfg.get("problem_id"),
                 problem_text=cfg.get("problem_text"),
             )
             result = await wrapper.initialize()
 
-            observation_for_registry = result["observation"].copy() if isinstance(result, dict) and "observation" in result else result.copy()
+            observation_for_registry = (
+                result["observation"].copy()
+                if isinstance(result, dict) and "observation" in result
+                else result.copy()
+            )
             for key in ["step_idx", "info"]:
                 if key in observation_for_registry:
                     del observation_for_registry[key]
@@ -493,7 +482,7 @@ async def create_environment(request: EnvCreateRequest) -> EnvCreateResponse:
 
     except Exception as e:
         logger.error(f"Failed to create environment: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e)) from e
 
 
 # --- Compatibility routes for existing eval scripts that expect CrafterClassic paths ---
@@ -509,7 +498,9 @@ async def compat_initialize(payload: dict) -> EnvCreateResponse:
         difficulty = str(wc.get("difficulty"))
     elif isinstance(cfg, dict) and cfg.get("difficulty"):
         difficulty = str(cfg.get("difficulty"))
-    req = EnvCreateRequest(env_name="crafter", config={"difficulty": difficulty}, seed=seed, rl_run_id="eval")
+    req = EnvCreateRequest(
+        env_name="crafter", config={"difficulty": difficulty}, seed=seed, rl_run_id="eval"
+    )
     return await create_environment(req)
 
 
@@ -525,10 +516,12 @@ async def compat_step(payload: dict) -> EnvStepResponse:
         actions_list = action.get("actions") if isinstance(action, dict) else None
         if isinstance(actions_list, list) and actions_list:
             for a in actions_list:
-                tool_calls.append({
-                    "tool": "interact",
-                    "args": {"action": a},
-                })
+                tool_calls.append(
+                    {
+                        "tool": "interact",
+                        "args": {"action": a},
+                    }
+                )
     req = EnvStepRequest(env_id=env_id, tool_calls=tool_calls)
     return await step_environment(req)
 
@@ -545,9 +538,7 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
     """Reset an environment to its initial state."""
     handle = registry.get_env(request.env_id)
     if not handle:
-        raise HTTPException(
-            status_code=404, detail=f"Environment {request.env_id} not found"
-        )
+        raise HTTPException(status_code=404, detail=f"Environment {request.env_id} not found")
 
     try:
         # Determine wrapper type and rebuild base env if a new seed is provided
@@ -586,29 +577,28 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
 
         elif True:
             # Try to dynamically import Wordle wrapper and check instance safely
-            try:
-                from .envs.wordle.environment import (
-                    WordleEnvironmentWrapper as _WordleWrapper,
-                )
-            except Exception:
-                _WordleWrapper = None  # type: ignore
+            wordle_wrapper_cls = None
+            with contextlib.suppress(Exception):
+                from .envs.wordle.environment import WordleEnvironmentWrapper
+
+                wordle_wrapper_cls = WordleEnvironmentWrapper  # type: ignore[assignment]
 
-            if _WordleWrapper is not None and isinstance(wrapper, _WordleWrapper):
+            if wordle_wrapper_cls is not None and isinstance(wrapper, wordle_wrapper_cls):
                 # Rebuild Wordle env with the same configuration; if we have a preserved
                 # initial_engine_snapshot, prefer constructing the instance directly.
                 try:
+                    from synth_ai.environments.examples.wordle.environment import (
+                        WordleEnvironment,
+                    )
                     from synth_ai.environments.examples.wordle.taskset import (
-                        create_wordle_taskset,
                         WordleTaskInstance,
                         WordleTaskInstanceMetadata,
-                    )
-                    from synth_ai.environments.examples.wordle.environment import (
-                        WordleEnvironment,
+                        create_wordle_taskset,
                     )
                 except Exception as e:
                     raise HTTPException(
                         status_code=500, detail=f"Wordle modules unavailable: {e}"
-                    )
+                    ) from e
 
                 init_snap = getattr(wrapper, "initial_engine_snapshot", None)
                 if init_snap is not None:
@@ -646,18 +636,18 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
             # Rebuild Wordle env with the same configuration; if we have a preserved
             # initial_engine_snapshot, prefer constructing the instance directly.
             try:
+                from synth_ai.environments.examples.wordle.environment import (
+                    WordleEnvironment,
+                )
                 from synth_ai.environments.examples.wordle.taskset import (
-                    create_wordle_taskset,
                     WordleTaskInstance,
                     WordleTaskInstanceMetadata,
-                )
-                from synth_ai.environments.examples.wordle.environment import (
-                    WordleEnvironment,
+                    create_wordle_taskset,
                 )
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Wordle modules unavailable: {e}"
-                )
+                ) from e
 
             init_snap = getattr(wrapper, "initial_engine_snapshot", None)
             if init_snap is not None:
@@ -693,27 +683,26 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
 
         elif True:
             # Try to dynamically import Sokoban wrapper and check instance safely
-            try:
-                from .envs.sokoban.environment import (
-                    SokobanEnvironmentWrapper as _SokobanWrapper,
-                )
-            except Exception:
-                _SokobanWrapper = None  # type: ignore
+            sokoban_wrapper_cls = None
+            with contextlib.suppress(Exception):
+                from .envs.sokoban.environment import SokobanEnvironmentWrapper
 
-            if _SokobanWrapper is not None and isinstance(wrapper, _SokobanWrapper):
+                sokoban_wrapper_cls = SokobanEnvironmentWrapper  # type: ignore[assignment]
+
+            if sokoban_wrapper_cls is not None and isinstance(wrapper, sokoban_wrapper_cls):
                 # Rebuild Sokoban env using stored config snapshot
                 try:
+                    from synth_ai.environments.examples.sokoban.environment import (
+                        SokobanEnvironment,
+                    )
                     from synth_ai.environments.examples.sokoban.taskset import (
                         SokobanTaskInstance,
                         SokobanTaskInstanceMetadata,
                     )
-                    from synth_ai.environments.examples.sokoban.environment import (
-                        SokobanEnvironment,
-                    )
                 except Exception as e:
                     raise HTTPException(
                         status_code=500, detail=f"Sokoban modules unavailable: {e}"
-                    )
+                    ) from e
 
                 cfg = dict(wrapper.config or {})
                 metadata = SokobanTaskInstanceMetadata(
@@ -738,17 +727,17 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
                 pass
             # Rebuild Sokoban env using stored config snapshot
             try:
+                from synth_ai.environments.examples.sokoban.environment import (
+                    SokobanEnvironment,
+                )
                 from synth_ai.environments.examples.sokoban.taskset import (
                     SokobanTaskInstance,
                     SokobanTaskInstanceMetadata,
                 )
-                from synth_ai.environments.examples.sokoban.environment import (
-                    SokobanEnvironment,
-                )
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Sokoban modules unavailable: {e}"
-                )
+                ) from e
 
             cfg = dict(wrapper.config or {})
             metadata = SokobanTaskInstanceMetadata(
@@ -757,9 +746,7 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
             instance = SokobanTaskInstance(
                 id=uuid4(),
                 impetus=Impetus(instructions="Reset"),
-                intent=Intent(
-                    rubric={"goal": "Reset"}, gold_trajectories=None, gold_state_diff={}
-                ),
+                intent=Intent(rubric={"goal": "Reset"}, gold_trajectories=None, gold_state_diff={}),
                 metadata=metadata,
                 is_reproducible=True,
                 initial_engine_snapshot=cfg.get("initial_state"),
@@ -777,7 +764,8 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
         try:
             base_env = handle.env.env  # type: ignore[attr-defined]
             pub_state = base_env.engine._get_public_state_from_env()  # type: ignore[attr-defined]
-            import hashlib, json as _json
+            import hashlib
+            import json as _json
 
             sig_src = {
                 "player_position": list(pub_state.player_position),
@@ -810,7 +798,7 @@ async def reset_environment(request: EnvResetRequest) -> EnvResetResponse:
 
     except Exception as e:
         logger.error(f"Failed to reset environment {request.env_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e)) from e
 
 
 @router.post("/step", response_model=EnvStepResponse)
@@ -818,30 +806,22 @@ async def step_environment(request: EnvStepRequest) -> EnvStepResponse:
     """Execute a step in the environment."""
     handle = registry.get_env(request.env_id)
     if not handle:
-        raise HTTPException(
-            status_code=404, detail=f"Environment {request.env_id} not found"
-        )
+        raise HTTPException(status_code=404, detail=f"Environment {request.env_id} not found")
 
     try:
         # Execute the step, pre-normalizing invalid Wordle guesses to avoid hard failures
         wrapper = handle.env
-        try:
-            from .envs.wordle.environment import (
-                WordleEnvironmentWrapper as _WordleWrapper,
-            )
-        except Exception:
-            _WordleWrapper = None  # type: ignore
+        wordle_wrapper_cls = None
+        with contextlib.suppress(Exception):
+            from .envs.wordle.environment import WordleEnvironmentWrapper
 
-        if _WordleWrapper is not None and isinstance(wrapper, _WordleWrapper):
+            wordle_wrapper_cls = WordleEnvironmentWrapper  # type: ignore[assignment]
+
+        if wordle_wrapper_cls is not None and isinstance(wrapper, wordle_wrapper_cls):
             expected_len = int(getattr(wrapper, "word_length", 5))
-            normalized: List[Dict[str, Any]] = []
+            normalized: list[dict[str, Any]] = []
             for tc in request.tool_calls or []:
-                tool = (
-                    tc.get("tool")
-                    or tc.get("tool_name")
-                    or tc.get("name")
-                    or "interact"
-                )
+                tool = tc.get("tool") or tc.get("tool_name") or tc.get("name") or "interact"
                 args = tc.get("arguments") or tc.get("args") or {}
                 if isinstance(args, str):
                     try:
@@ -861,9 +841,7 @@ async def step_environment(request: EnvStepRequest) -> EnvStepResponse:
                         # Preserve the original tool name (interact or submit) for the environment to handle
                         normalized.append({"tool": tool, "args": {"guess": g}})
                 else:
-                    normalized.append(
-                        {"tool": "invalid_guess", "args": {"original_guess": guess}}
-                    )
+                    normalized.append({"tool": "invalid_guess", "args": {"original_guess": guess}})
             result = await wrapper.step(normalized)
         else:
             result = await handle.env.step(request.tool_calls)
@@ -913,21 +891,18 @@ async def step_environment(request: EnvStepRequest) -> EnvStepResponse:
         logger.error(f"Failed to step environment {request.env_id}: {e}")
         # Fallback for Wordle: convert invalid guesses into 'invalid_guess' tool calls and retry once
         try:
-            from .envs.wordle.environment import (
-                WordleEnvironmentWrapper as _WordleWrapper,
-            )
+            wordle_wrapper_cls = None
+            with contextlib.suppress(Exception):
+                from .envs.wordle.environment import WordleEnvironmentWrapper
+
+                wordle_wrapper_cls = WordleEnvironmentWrapper  # type: ignore[assignment]
 
             wrapper = handle.env
-            if isinstance(wrapper, _WordleWrapper):
+            if wordle_wrapper_cls is not None and isinstance(wrapper, wordle_wrapper_cls):
                 expected_len = int(getattr(wrapper, "word_length", 5))
-                normalized: List[Dict[str, Any]] = []
+                normalized: list[dict[str, Any]] = []
                 for tc in request.tool_calls or []:
-                    tool = (
-                        tc.get("tool")
-                        or tc.get("tool_name")
-                        or tc.get("name")
-                        or "interact"
-                    )
+                    tool = tc.get("tool") or tc.get("tool_name") or tc.get("name") or "interact"
                     args = tc.get("arguments") or tc.get("args") or {}
                     if isinstance(args, str):
                         try:
@@ -947,9 +922,7 @@ async def step_environment(request: EnvStepRequest) -> EnvStepResponse:
                                 }
                             )
                         else:
-                            normalized.append(
-                                {"tool": "interact", "args": {"guess": g}}
-                            )
+                            normalized.append({"tool": "interact", "args": {"guess": g}})
                     else:
                         normalized.append(
                             {"tool": "invalid_guess", "args": {"original_guess": guess}}
@@ -981,7 +954,7 @@ async def step_environment(request: EnvStepRequest) -> EnvStepResponse:
             # Ignore fallback errors; fall through to generic error
             pass
 
-        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}")
+        raise HTTPException(status_code=500, detail=f"{type(e).__name__}: {e}") from e
 
 
 @router.post("/snapshot", response_model=EnvSnapshotResponse)
@@ -989,9 +962,7 @@ async def snapshot_environment(request: EnvSnapshotRequest) -> EnvSnapshotRespon
     """Create a snapshot of the environment state."""
     handle = registry.get_env(request.env_id)
     if not handle:
-        raise HTTPException(
-            status_code=404, detail=f"Environment {request.env_id} not found"
-        )
+        raise HTTPException(status_code=404, detail=f"Environment {request.env_id} not found")
 
     try:
         # Serialize environment state
@@ -1022,7 +993,7 @@ async def snapshot_environment(request: EnvSnapshotRequest) -> EnvSnapshotRespon
 
     except Exception as e:
         logger.error(f"Failed to snapshot environment {request.env_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e)) from e
 
 
 @router.post("/restore", response_model=EnvRestoreResponse)
@@ -1030,9 +1001,7 @@ async def restore_environment(request: EnvRestoreRequest) -> EnvRestoreResponse:
     """Restore an environment from a snapshot."""
     snapshot = registry.get_snapshot(request.snapshot_id)
     if not snapshot:
-        raise HTTPException(
-            status_code=404, detail=f"Snapshot {request.snapshot_id} not found"
-        )
+        raise HTTPException(status_code=404, detail=f"Snapshot {request.snapshot_id} not found")
 
     if snapshot.kind != "env":
         raise HTTPException(
@@ -1104,18 +1073,18 @@ async def restore_environment(request: EnvRestoreRequest) -> EnvRestoreResponse:
             )
         elif name_lower == "wordle":
             try:
+                from synth_ai.environments.examples.wordle.environment import (
+                    WordleEnvironment,
+                )
                 from synth_ai.environments.examples.wordle.taskset import (
-                    create_wordle_taskset,
                     WordleTaskInstance,
                     WordleTaskInstanceMetadata,
-                )
-                from synth_ai.environments.examples.wordle.environment import (
-                    WordleEnvironment,
+                    create_wordle_taskset,
                 )
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Wordle modules unavailable: {e}"
-                )
+                ) from e
 
             cfg = state_dict.get("config", {}) or {}
             word_length = int(cfg.get("word_length", 5))
@@ -1146,16 +1115,12 @@ async def restore_environment(request: EnvRestoreRequest) -> EnvRestoreResponse:
                 base_env = WordleEnvironment(task_instance=instance)
             # Lazy import of wrapper only when needed
             try:
-                from .envs.wordle.environment import (
-                    WordleEnvironmentWrapper as _WordleWrapper,
-                )
+                from .envs.wordle.environment import WordleEnvironmentWrapper
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Wordle wrapper unavailable: {e}"
-                )
-            wrapper = await _WordleWrapper.deserialize(
-                payload=state_dict, env=base_env
-            )
+                ) from e
+            wrapper = await WordleEnvironmentWrapper.deserialize(payload=state_dict, env=base_env)
 
             env_id = registry.register_env(
                 env=wrapper,
@@ -1176,22 +1141,20 @@ async def restore_environment(request: EnvRestoreRequest) -> EnvRestoreResponse:
 
         elif name_lower == "sokoban":
             try:
+                from synth_ai.environments.examples.sokoban.environment import (
+                    SokobanEnvironment,
+                )
                 from synth_ai.environments.examples.sokoban.taskset import (
                     SokobanTaskInstance,
                     SokobanTaskInstanceMetadata,
                 )
-                from synth_ai.environments.examples.sokoban.environment import (
-                    SokobanEnvironment,
-                )
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Sokoban modules unavailable: {e}"
-                )
+                ) from e
 
             cfg = state_dict.get("config", {}) or {}
-            metadata = SokobanTaskInstanceMetadata(
-                difficulty=cfg.get("difficulty", "easy")
-            )
+            metadata = SokobanTaskInstanceMetadata(difficulty=cfg.get("difficulty", "easy"))
             instance = SokobanTaskInstance(
                 id=uuid4(),
                 impetus=Impetus(instructions="Restore"),
@@ -1207,16 +1170,12 @@ async def restore_environment(request: EnvRestoreRequest) -> EnvRestoreResponse:
             base_env = SokobanEnvironment(task_instance=instance)
             # Lazy import of wrapper only when needed
             try:
-                from .envs.sokoban.environment import (
-                    SokobanEnvironmentWrapper as _SokobanWrapper,
-                )
+                from .envs.sokoban.environment import SokobanEnvironmentWrapper
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Sokoban wrapper unavailable: {e}"
-                )
-            wrapper = await _SokobanWrapper.deserialize(
-                payload=state_dict, env=base_env
-            )
+                ) from e
+            wrapper = await SokobanEnvironmentWrapper.deserialize(payload=state_dict, env=base_env)
 
             env_id = registry.register_env(
                 env=wrapper,
@@ -1242,10 +1201,8 @@ async def restore_environment(request: EnvRestoreRequest) -> EnvRestoreResponse:
             )
 
     except Exception as e:
-        logger.error(
-            f"Failed to restore environment from snapshot {request.snapshot_id}: {e}"
-        )
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Failed to restore environment from snapshot {request.snapshot_id}: {e}")
+        raise HTTPException(status_code=500, detail=str(e)) from e
 
 
 @router.post("/terminate", response_model=EnvTerminateResponse)
@@ -1253,9 +1210,7 @@ async def terminate_environment(request: EnvTerminateRequest) -> EnvTerminateRes
     """Terminate an environment and clean up resources."""
     handle = registry.get_env(request.env_id)
     if not handle:
-        raise HTTPException(
-            status_code=404, detail=f"Environment {request.env_id} not found"
-        )
+        raise HTTPException(status_code=404, detail=f"Environment {request.env_id} not found")
 
     try:
         # Call terminate on the environment
@@ -1268,4 +1223,4 @@ async def terminate_environment(request: EnvTerminateRequest) -> EnvTerminateRes
 
     except Exception as e:
         logger.error(f"Failed to terminate environment {request.env_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e)) from e