[PATCH] Deprecate HTTP in websockets

pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
     "huggingface_hub>=0.20.0",
     "openai>=2.7.2",
     "tomli>=2.3.0",
-    "tomli-w>=1.2.0"
+    "tomli-w>=1.2.0",
+    "websockets>=15.0.1",
 ]
 
 [project.optional-dependencies]

src/openenv/cli/templates/openenv_env/README.md

@@ -114,6 +114,7 @@ The deployed space includes:
 - **Web Interface** at `/web` - Interactive UI for exploring the environment
 - **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
 - **Health Check** at `/health` - Container health monitoring
+- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
 
 ## Environment Details
 
@@ -154,6 +155,61 @@ result = __ENV_NAME__env.step(__ENV_CLASS_NAME__Action(message="Hello!"))
 
 Note: When connecting to an existing server, `__ENV_NAME__env.close()` will NOT stop the server.
 
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env
+
+# Connect with context manager (auto-connects and closes)
+with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as env:
+    result = env.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+    # Multiple steps with low latency
+    for msg in ["Hello", "World", "!"]:
+        result = env.step(__ENV_CLASS_NAME__Action(message=msg))
+        print(f"Echoed: {result.observation.echoed_message}")
+```
+
+The client uses WebSocket connections for:
+- **Lower latency**: No HTTP connection overhead per request
+- **Persistent session**: Server maintains your environment state
+- **Efficient for episodes**: Better for many sequential steps
+
+### Concurrent WebSocket Sessions
+
+The server supports multiple concurrent WebSocket connections. To enable this,
+modify `server/app.py` to use factory mode:
+
+```python
+# In server/app.py - use factory mode for concurrent sessions
+app = create_app(
+    __ENV_CLASS_NAME__Environment,  # Pass class, not instance
+    __ENV_CLASS_NAME__Action,
+    __ENV_CLASS_NAME__Observation,
+    max_concurrent_envs=4,  # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env
+from concurrent.futures import ThreadPoolExecutor
+
+def run_episode(client_id: int):
+    with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as env:
+        result = env.reset()
+        for i in range(10):
+            result = env.step(__ENV_CLASS_NAME__Action(message=f"Client {client_id}, step {i}"))
+        return client_id, result.observation.message_length
+
+# Run 4 episodes concurrently
+with ThreadPoolExecutor(max_workers=4) as executor:
+    results = list(executor.map(run_episode, range(4)))
+```
+
 ## Development & Testing
 
 ### Direct Environment Testing
@@ -189,11 +245,11 @@ __ENV_NAME__/
 ├── openenv.yaml           # OpenEnv manifest
 ├── pyproject.toml         # Project metadata and dependencies
 ├── uv.lock                # Locked dependencies (generated)
-├── client.py              # __ENV_CLASS_NAME__Env client implementation
+├── client.py              # __ENV_CLASS_NAME__Env client
 ├── models.py              # Action and Observation models
 └── server/
     ├── __init__.py        # Server module exports
     ├── __ENV_NAME___environment.py  # Core environment logic
-    ├── app.py             # FastAPI application
+    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
     └── Dockerfile         # Container image definition
 ```

src/openenv/cli/templates/openenv_env/__init__.py

@@ -4,10 +4,13 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""__ENV_TITLE_NAME__ Environment - A simple test environment for HTTP server."""
+"""__ENV_TITLE_NAME__ Environment."""
 
 from .client import __ENV_CLASS_NAME__Env
 from .models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
 
-__all__ = ["__ENV_CLASS_NAME__Action", "__ENV_CLASS_NAME__Observation", "__ENV_CLASS_NAME__Env"]
-
+__all__ = [
+    "__ENV_CLASS_NAME__Action",
+    "__ENV_CLASS_NAME__Observation",
+    "__ENV_CLASS_NAME__Env",
+]

src/openenv/cli/templates/openenv_env/client.py

@@ -4,50 +4,47 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-"""
-__ENV_TITLE_NAME__ Environment HTTP Client.
+"""__ENV_TITLE_NAME__ Environment Client."""
 
-This module provides the client for connecting to a __ENV_TITLE_NAME__ Environment server
-over HTTP.
-"""
-
-from typing import Any, Dict
+from typing import Dict
 
 from openenv.core.client_types import StepResult
 from openenv.core.env_server.types import State
-from openenv.core.http_env_client import HTTPEnvClient
+from openenv.core import EnvClient
 
 from .models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
 
 
-class __ENV_CLASS_NAME__Env(HTTPEnvClient[__ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation]):
+class __ENV_CLASS_NAME__Env(EnvClient[__ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation]):
     """
-    HTTP client for the __ENV_TITLE_NAME__ Environment.
+    Client for the __ENV_TITLE_NAME__ Environment.
 
-    This client connects to a __ENV_CLASS_NAME__Environment 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.
+    Each client instance has its own dedicated environment session on the server.
 
     Example:
         >>> # Connect to a running server
-        >>> client = __ENV_CLASS_NAME__Env(base_url="http://localhost:8000")
-        >>> result = client.reset()
-        >>> print(result.observation.echoed_message)
-        >>>
-        >>> # Send a message
-        >>> result = client.step(__ENV_CLASS_NAME__Action(message="Hello!"))
-        >>> print(result.observation.echoed_message)
-        >>> print(result.reward)
+        >>> with __ENV_CLASS_NAME__Env(base_url="http://localhost:8000") as client:
+        ...     result = client.reset()
+        ...     print(result.observation.echoed_message)
+        ...
+        ...     result = client.step(__ENV_CLASS_NAME__Action(message="Hello!"))
+        ...     print(result.observation.echoed_message)
 
     Example with Docker:
         >>> # Automatically start container and connect
         >>> client = __ENV_CLASS_NAME__Env.from_docker_image("__ENV_NAME__-env:latest")
-        >>> result = client.reset()
-        >>> result = client.step(__ENV_CLASS_NAME__Action(message="Test"))
+        >>> try:
+        ...     result = client.reset()
+        ...     result = client.step(__ENV_CLASS_NAME__Action(message="Test"))
+        ... finally:
+        ...     client.close()
     """
 
     def _step_payload(self, action: __ENV_CLASS_NAME__Action) -> Dict:
         """
-        Convert __ENV_CLASS_NAME__Action to JSON payload for step request.
+        Convert __ENV_CLASS_NAME__Action to JSON payload for step message.
 
         Args:
             action: __ENV_CLASS_NAME__Action instance
@@ -64,7 +61,7 @@ def _parse_result(self, payload: Dict) -> StepResult[__ENV_CLASS_NAME__Observati
         Parse server response into StepResult[__ENV_CLASS_NAME__Observation].
 
         Args:
-            payload: JSON response from server
+            payload: JSON response data from server
 
         Returns:
             StepResult with __ENV_CLASS_NAME__Observation
@@ -89,7 +86,7 @@ def _parse_state(self, payload: Dict) -> State:
         Parse server response into State object.
 
         Args:
-            payload: JSON response from /state endpoint
+            payload: JSON response from state request
 
         Returns:
             State object with episode_id and step_count

src/openenv/cli/templates/openenv_env/models.py

@@ -10,22 +10,20 @@
 The __ENV_NAME__ environment is a simple test environment that echoes back messages.
 """
 
-from dataclasses import dataclass
+from pydantic import Field
 
 from openenv.core.env_server.types import Action, Observation
 
 
-@dataclass(kw_only=True)
 class __ENV_CLASS_NAME__Action(Action):
     """Action for the __ENV_TITLE_NAME__ environment - just a message to echo."""
 
-    message: str
+    message: str = Field(..., description="Message to echo back")
 
 
-@dataclass(kw_only=True)
 class __ENV_CLASS_NAME__Observation(Observation):
     """Observation from the __ENV_TITLE_NAME__ environment - the echoed message."""
 
-    echoed_message: str
-    message_length: int = 0
+    echoed_message: str = Field(default="", description="The echoed message")
+    message_length: int = Field(default=0, description="Length of the echoed message")
 

src/openenv/cli/templates/openenv_env/pyproject.toml

@@ -15,6 +15,8 @@ description = "__ENV_TITLE_NAME__ environment for OpenEnv"
 requires-python = ">=3.10"
 dependencies = [
     # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    # install from github
+    # "openenv[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
     "openenv[core]>=0.2.0",
     # Environment-specific dependencies
     # Add all dependencies needed for your environment here

src/openenv/cli/templates/openenv_env/server/__ENV_NAME___environment.py

@@ -36,6 +36,12 @@ class __ENV_CLASS_NAME__Environment(Environment):
         >>> print(obs.message_length)  # 5
     """
 
+    # Enable concurrent WebSocket sessions.
+    # Set to True if your environment isolates state between instances.
+    # When True, multiple WebSocket clients can connect simultaneously, each
+    # getting their own environment instance (when using factory mode in app.py).
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
     def __init__(self):
         """Initialize the __ENV_NAME__ environment."""
         self._state = State(episode_id=str(uuid4()), step_count=0)

src/openenv/cli/templates/openenv_env/server/app.py

@@ -8,7 +8,14 @@
 FastAPI application for the __ENV_TITLE_NAME__ Environment.
 
 This module creates an HTTP server that exposes the __ENV_CLASS_NAME__Environment
-over HTTP endpoints, making it compatible with HTTPEnvClient.
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment
+    - POST /step: Execute an action
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
 
 Usage:
     # Development (with auto-reload):
@@ -28,18 +35,18 @@
         "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
     ) from e
 
-from __ENV_NAME__.models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
+# Import from local models.py (PYTHONPATH includes /app/env in Docker)
+from models import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation
 from .__ENV_NAME___environment import __ENV_CLASS_NAME__Environment
 
-# Create the environment instance
-env = __ENV_CLASS_NAME__Environment()
 
 # Create the app with web interface and README integration
 app = create_app(
-    env,
+    __ENV_CLASS_NAME__Environment,
     __ENV_CLASS_NAME__Action,
     __ENV_CLASS_NAME__Observation,
     env_name="__ENV_NAME__",
+    max_concurrent_envs=1,  # increase this number to allow more concurrent WebSocket sessions
 )
 
 

src/openenv/core/README.md

@@ -22,8 +22,8 @@ Core components for OpenEnv - a framework for building HTTP-based agentic enviro
 
 ## Features
 
-- **HTTPEnvClient**: Generic HTTP client for interacting with remote environments
-- **HTTPEnvServer**: FastAPI-based server wrapper for exposing environments over HTTP
+- **EnvClient**: Generic client for interacting with remote environments
+- **HTTPEnvServer**: FastAPI-based server wrapper for exposing environments over HTTP/WebSocket
 - **Container Providers**: Pluggable architecture for running containers (Docker, Kubernetes, etc.)
 - **Type System**: Strongly-typed Action/Observation/State interfaces
 - **Web Interface**: Optional web UI for interacting with environments
@@ -44,7 +44,7 @@ pip install "openenv[core]"
 ### Creating an Environment Client
 
 ```python
-from openenv.core import HTTPEnvClient, StepResult
+from openenv.core import EnvClient, StepResult
 from dataclasses import dataclass
 
 @dataclass
@@ -55,7 +55,7 @@ class MyAction:
 class MyObservation:
     response: str
 
-class MyEnvClient(HTTPEnvClient[MyAction, MyObservation]):
+class MyEnvClient(EnvClient[MyAction, MyObservation]):
     def _step_payload(self, action: MyAction) -> dict:
         return {"text": action.text}
 
@@ -141,7 +141,7 @@ provider.stop_container()
 
 ## API Reference
 
-### HTTPEnvClient
+### EnvClient
 
 Base class for environment clients with these abstract methods:
 

src/openenv/core/__init__.py

@@ -7,13 +7,10 @@
 """Core components for agentic environments."""
 
 # Re-export main components from submodules for convenience
-from .env_server import *
-from .client_types import StepResult
-from .http_env_client import HTTPEnvClient
+from .env_server import *  # noqa: F403
+from . import env_server
+from .env_client import EnvClient
 
 # Note: MCP module doesn't export anything yet
 
-__all__ = [
-    "HTTPEnvClient",
-    "StepResult",
-]
+__all__ = ["EnvClient"] + env_server.__all__  # type: ignore

src/openenv/core/client_types.py

@@ -1,9 +1,10 @@
 # Type definitions for EnvTorch
 from dataclasses import dataclass
-from typing import Any, Generic, Optional, TypeVar
+from typing import Generic, Optional, TypeVar
 
 # Generic type for observations
-ObsT = TypeVar("ObsT")  # TypeVar for typehinting in IDEs
+ObsT = TypeVar("ObsT")
+StateT = TypeVar("StateT")
 
 
 @dataclass

src/openenv/core/containers/runtime/providers.py

@@ -8,7 +8,7 @@
 Container provider abstractions for running environment servers.
 
 This module provides a pluggable architecture for different container providers
-(local Docker, Kubernetes, cloud providers, etc.) to be used with HTTPEnvClient.
+(local Docker, Kubernetes, cloud providers, etc.) to be used with EnvClient.
 """
 
 from __future__ import annotations