[FEATURE] WebSocket-based Concurrency Architecture

docs/environment-builder.md

@@ -10,7 +10,7 @@ A typical workflow looks like:
 
 1. Scaffold a new environment with `openenv init`.
 2. Customize your models, environment logic, and FastAPI server.
-3. Implement a typed `HTTPEnvClient`.
+3. Implement a typed `EnvClient` (WebSocket-based for persistent sessions).
 4. Configure dependencies and the Dockerfile once.
 5. Use the CLI (`openenv build`, `openenv validate`, `openenv push`) to package and share your work.
 
@@ -58,33 +58,26 @@ my_env/
     └── Dockerfile
 ```
 
-Python classes are generated for the action, observation, and state, and a client is generated for the environment. For example, you will find `MyEnvironment`, `MyAction`, `MyObservation`, and `MyState` in the `my_env` directory based on the name of the environment you provided.
+Python classes are generated for the action, observation, environment, and client. For example, you will find `MyEnvironment`, `MyAction`, `MyObservation`, and `MyEnv` (client) in the `my_env` directory based on the name you provided. The environment uses the core `State` class from `openenv.core.env_server.types`.
 
 ### 2. Define Models
 
-Edit `models.py` to describe your action, observation, and state dataclasses:
+Edit `models.py` to describe your action and observation using Pydantic:
 
 ```python
 # models.py
-from dataclasses import dataclass
-from openenv.core.env_server import Action, Observation, State
+from pydantic import Field
+from openenv.core.env_server.types import Action, Observation
 
-@dataclass
 class MyAction(Action):
     """Your custom action."""
-    command: str
-    parameters: dict
+    command: str = Field(..., description="Command to execute")
+    parameters: dict = Field(default_factory=dict, description="Command parameters")
 
-@dataclass
 class MyObservation(Observation):
     """Your custom observation."""
-    result: str
-    success: bool
-
-@dataclass
-class MyState(State):
-    """Custom state fields."""
-    custom_field: int = 0
+    result: str = Field(..., description="Result of the action")
+    success: bool = Field(..., description="Whether the action succeeded")
 ```
 
 ### 3. Implement Environment Logic
@@ -93,70 +86,104 @@ Customize `server/my_environment.py` by extending `Environment`:
 
 ```python
 # server/my_environment.py
-import uuid
-from openenv.core.env_server import Environment
-from ..models import MyAction, MyObservation, MyState
+from uuid import uuid4
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+from models import MyAction, MyObservation
 
 class MyEnvironment(Environment):
     def __init__(self):
-        super().__init__()
-        self._state = MyState()
+        self._state = State(episode_id=str(uuid4()), step_count=0)
 
     def reset(self) -> MyObservation:
-        self._state = MyState(episode_id=str(uuid.uuid4()))
-        return MyObservation(result="Ready", success=True)
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        return MyObservation(result="Ready", success=True, done=False, reward=0.0)
 
     def step(self, action: MyAction) -> MyObservation:
         # Implement your logic here
         self._state.step_count += 1
         result = self._execute_command(action.command)
-        return MyObservation(result=result, success=True)
+        return MyObservation(result=result, success=True, done=False, reward=1.0)
 
     @property
-    def state(self) -> MyState:
+    def state(self) -> State:
         return self._state
 ```
 
 ### 4. Create the FastAPI Server
 
-`server/app.py` should expose the environment through `create_fastapi_app`:
+`server/app.py` should expose the environment through `create_app`.
+
+**Important:** You must pass a class or factory function (not an instance) to enable WebSocket-based concurrent sessions:
+
+```python
+# server/app.py
+from openenv.core.env_server import create_app
+from ..models import MyAction, MyObservation
+from .my_environment import MyEnvironment
+
+# Pass the class (factory) - each WebSocket session gets its own instance
+app = create_app(MyEnvironment, MyAction, MyObservation, env_name="my_env")
+```
+
+For environments with constructor arguments, create a factory function:
 
 ```python
 # server/app.py
-from openenv.core.env_server import create_fastapi_app
+import os
+from openenv.core.env_server import create_app
 from ..models import MyAction, MyObservation
 from .my_environment import MyEnvironment
 
-env = MyEnvironment()
-app = create_fastapi_app(env, MyAction, MyObservation)
+# Read config from environment variables
+api_key = os.getenv("MY_API_KEY")
+timeout = int(os.getenv("MY_TIMEOUT", "30"))
+
+def create_my_environment():
+    """Factory function that creates MyEnvironment with config."""
+    return MyEnvironment(api_key=api_key, timeout=timeout)
+
+# Pass the factory function
+app = create_app(create_my_environment, MyAction, MyObservation, env_name="my_env")
 ```
 
 ### 5. Implement the Client
 
-`client.py` extends `HTTPEnvClient` so users can interact with your server over HTTP or Docker:
+`client.py` extends `EnvClient` so users can interact with your server via WebSocket for persistent sessions:
 
 ```python
 # client.py
-from openenv.core.http_env_client import HTTPEnvClient
-from openenv.core.types import StepResult
+from openenv.core.env_client import EnvClient
+from openenv.core.client_types import StepResult
 from .models import MyAction, MyObservation, MyState
 
-class MyEnv(HTTPEnvClient[MyAction, MyObservation]):
+class MyEnv(EnvClient[MyAction, MyObservation, MyState]):
     def _step_payload(self, action: MyAction) -> dict:
         return {"command": action.command, "parameters": action.parameters}
 
     def _parse_result(self, payload: dict) -> StepResult[MyObservation]:
-        obs = MyObservation(**payload["observation"])
+        obs_data = payload.get("observation", {})
+        obs = MyObservation(
+            result=obs_data.get("result", ""),
+            success=obs_data.get("success", False),
+            done=payload.get("done", False),
+            reward=payload.get("reward"),
+        )
         return StepResult(
             observation=obs,
             reward=payload.get("reward"),
             done=payload.get("done", False),
         )
 
-    def _parse_state(self, payload: dict) -> MyState:
-        return MyState(**payload)
+    def _parse_state(self, payload: dict) -> State:
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )
 ```
 
+The `EnvClient` maintains a persistent WebSocket connection to the server, enabling efficient multi-step interactions with lower latency compared to HTTP. Each client instance gets its own dedicated environment session on the server.
+
 ### 6. Configure Dependencies & Dockerfile
 
 The CLI template ships with `pyproject.toml` and `server/Dockerfile`. You should manage your python dependencies with `uv` or `pip` in the `pyproject.toml` file. Other dependencies should be installed in the Dockerfile.
@@ -322,22 +349,29 @@ client = MyEnv.from_hub("my-org/my-env")
 # Or, connect to the local server
 client = MyEnv(base_url="http://localhost:8000")
 
-# Reset
-result = client.reset()
-print(result.observation.result)  # "Ready"
-
-# Execute actions
-result = client.step(MyAction(command="test", parameters={}))
-print(result.observation.result)
-print(result.observation.success)
-
-# Get state
-state = client.state()
-print(state.episode_id)
-print(state.step_count)
-
-# Cleanup
-client.close()
+# Use context manager for automatic cleanup (recommended)
+with client:
+    # Reset
+    result = client.reset()
+    print(result.observation.result)  # "Ready"
+
+    # Execute actions
+    result = client.step(MyAction(command="test", parameters={}))
+    print(result.observation.result)
+    print(result.observation.success)
+
+    # Get state
+    state = client.state()
+    print(state.episode_id)
+    print(state.step_count)
+
+# Or manually manage the connection
+try:
+    client = MyEnv(base_url="http://localhost:8000")
+    result = client.reset()
+    result = client.step(MyAction(command="test", parameters={}))
+finally:
+    client.close()
 ```
 
 ## Nice work! You've now built and used your own OpenEnv environment.

envs/atari_env/client.py

@@ -5,10 +5,10 @@
 # LICENSE file in the root directory of this source tree.
 
 """
-Atari Environment HTTP Client.
+Atari Environment Client.
 
 This module provides the client for connecting to an Atari Environment server
-over HTTP.
+via WebSocket for persistent sessions.
 """
 
 from __future__ import annotations
@@ -17,36 +17,38 @@
 
 from openenv.core.client_types import StepResult
 
-from openenv.core.http_env_client import HTTPEnvClient
+from openenv.core.env_client import EnvClient
 
 from .models import AtariAction, AtariObservation, AtariState
 
 if TYPE_CHECKING:
     from openenv.core.containers.runtime import ContainerProvider
 
 
-class AtariEnv(HTTPEnvClient[AtariAction, AtariObservation]):
+class AtariEnv(EnvClient[AtariAction, AtariObservation, AtariState]):
     """
-    HTTP client for Atari Environment.
+    Client for Atari Environment.
 
-    This client connects to an AtariEnvironment HTTP server and provides
-    methods to interact with it: reset(), step(), and state access.
+    This client maintains a persistent WebSocket connection to the environment
+    server, enabling efficient multi-step interactions with lower latency.
 
     Example:
         >>> # Connect to a running server
-        >>> client = AtariEnv(base_url="http://localhost:8000")
-        >>> result = client.reset()
-        >>> print(result.observation.screen_shape)
-        >>>
-        >>> # Take an action
-        >>> result = client.step(AtariAction(action_id=2))  # UP
-        >>> print(result.reward, result.done)
+        >>> with AtariEnv(base_url="http://localhost:8000") as client:
+        ...     result = client.reset()
+        ...     print(result.observation.screen_shape)
+        ...
+        ...     result = client.step(AtariAction(action_id=2))  # UP
+        ...     print(result.reward, result.done)
 
     Example with Docker:
         >>> # Automatically start container and connect
         >>> client = AtariEnv.from_docker_image("atari-env:latest")
-        >>> result = client.reset()
-        >>> result = client.step(AtariAction(action_id=0))  # NOOP
+        >>> try:
+        ...     result = client.reset()
+        ...     result = client.step(AtariAction(action_id=0))  # NOOP
+        ... finally:
+        ...     client.close()
     """
 
     def _step_payload(self, action: AtariAction) -> Dict[str, Any]:

envs/atari_env/server/app.py

@@ -8,7 +8,7 @@
 FastAPI application for the Atari Environment.
 
 This module creates an HTTP server that exposes Atari games
-over HTTP endpoints, making them compatible with HTTPEnvClient.
+over HTTP and WebSocket endpoints, compatible with EnvClient.
 
 Usage:
     # Development (with auto-reload):
@@ -52,19 +52,24 @@
 mode = int(mode) if mode is not None else None
 difficulty = int(difficulty) if difficulty is not None else None
 
-# Create the environment instance
-env = AtariEnvironment(
-    game_name=game_name,
-    obs_type=obs_type,
-    full_action_space=full_action_space,
-    mode=mode,
-    difficulty=difficulty,
-    repeat_action_probability=repeat_action_prob,
-    frameskip=frameskip,
-)
+
+# Factory function to create AtariEnvironment instances
+def create_atari_environment():
+    """Factory function that creates AtariEnvironment with config."""
+    return AtariEnvironment(
+        game_name=game_name,
+        obs_type=obs_type,
+        full_action_space=full_action_space,
+        mode=mode,
+        difficulty=difficulty,
+        repeat_action_probability=repeat_action_prob,
+        frameskip=frameskip,
+    )
+
 
 # Create the FastAPI app with web interface and README integration
-app = create_app(env, AtariAction, AtariObservation, env_name="atari_env")
+# Pass the factory function instead of an instance for WebSocket session support
+app = create_app(create_atari_environment, AtariAction, AtariObservation, env_name="atari_env")
 
 
 if __name__ == "__main__":

envs/browsergym_env/client.py

@@ -1,17 +1,18 @@
-"""HTTP client for the BrowserGym environment."""
+"""Client for the BrowserGym environment."""
 
 from typing import Any, Dict
 
-from openenv.core.http_env_client import HTTPEnvClient, StepResult
+from openenv.core.client_types import StepResult
+from openenv.core.env_client import EnvClient
 from .models import (
     BrowserGymAction,
     BrowserGymObservation,
     BrowserGymState,
 )
 
 
-class BrowserGymEnv(HTTPEnvClient[BrowserGymAction, BrowserGymObservation]):
-    """Client for interacting with the BrowserGym environment over HTTP.
+class BrowserGymEnv(EnvClient[BrowserGymAction, BrowserGymObservation, BrowserGymState]):
+    """Client for interacting with the BrowserGym environment.
 
     BrowserGym provides unified access to multiple web navigation benchmarks:
     - MiniWoB++: 100+ training tasks (no external infrastructure needed!)

envs/browsergym_env/server/app.py

@@ -15,19 +15,24 @@
 timeout = float(os.environ.get("BROWSERGYM_TIMEOUT", "10000"))
 port = int(os.environ.get("BROWSERGYM_PORT", "8000"))
 
-# Create the environment instance
-env = BrowserGymEnvironment(
-    benchmark=benchmark,
-    task_name=task_name,
-    headless=headless,
-    viewport_width=viewport_width,
-    viewport_height=viewport_height,
-    timeout=timeout,
-)
+
+# Factory function to create BrowserGymEnvironment instances
+def create_browsergym_environment():
+    """Factory function that creates BrowserGymEnvironment with config."""
+    return BrowserGymEnvironment(
+        benchmark=benchmark,
+        task_name=task_name,
+        headless=headless,
+        viewport_width=viewport_width,
+        viewport_height=viewport_height,
+        timeout=timeout,
+    )
+
 
 # Create the FastAPI app
+# Pass the factory function instead of an instance for WebSocket session support
 app = create_app(
-    env,
+    create_browsergym_environment,
     BrowserGymAction,
     BrowserGymObservation,
     env_name="browsergym_env",