docs: document the Task API

docs/source/_toctree.yml

@@ -27,6 +27,8 @@
       title: Simulation vs Production
     - local: guides/mcp-environment-lifecycle
       title: MCP Environment Lifecycle
+    - local: guides/task-api
+      title: The Task API
     - local: guides/connecting
       title: Connecting to Servers
     - local: guides/runtime-providers

docs/source/guides/task-api.md

@@ -0,0 +1,273 @@
+# The Task API
+
+This guide explains the Task API: the optional discovery layer that lets a dataset-backed environment publish the *set of tasks* it can run, so a trainer or evaluator can enumerate them and choose which one each episode uses.
+
+It exists to answer a common question: OpenEnv's step loop gives you one episode at a time, so how does a training script find out that an environment holds 7,595 test problems, and how does it ask for problem number 12?
+
+## The Short Answer
+
+The Task API has two halves:
+
+- **Server side**: your environment optionally implements five methods — `list_splits()`, `list_tasks()`, `num_tasks()`, `get_task()`, `get_task_range()`. These are described by the `TaskProvider` protocol.
+- **HTTP side**: when those methods exist, the environment server automatically exposes them as routes under `/{env_name}/…`. You do not register anything.
+
+Task discovery is **metadata only**. It never starts an episode. Selecting a task happens the usual way, through `reset()`:
+
+```python
+env.list_splits()            # ["train", "test"]
+env.num_tasks("test")        # 7595
+env.get_task("test", 12)     # {"id": "test-12", "index": 12, "split": "test"}
+env.reset(split="test", index=12)   # <- this is what starts the episode
+```
+
+The API is shaped to be compatible with ORS/OpenReward task and split conventions, which is why environments imported with `openenv import` get it for free.
+
+## When You Need It
+
+Implement the Task API when your environment is **dataset-backed** — when an episode means "run one row of a dataset" and a trainer needs to iterate, shard, or shuffle over those rows.
+
+You do **not** need it for environments where every episode is generated or self-contained: Echo, Snake, a game, a REPL. Those environments simply omit the methods, and the routes report `501 Not Implemented`.
+
+## The Protocol
+
+`TaskProvider` is a `typing.Protocol`, not a base class. You do not inherit from it — you just define the methods on your `Environment` subclass and the server discovers them by name.
+
+```python
+from typing import Any, Optional
+
+class TaskProvider(Protocol):
+    def list_splits(self) -> list[Any]: ...
+    def list_tasks(self, split: str) -> list[Any]: ...
+    def num_tasks(self, split: str) -> int: ...
+    def get_task(self, split: str, index: int) -> Any: ...
+    def get_task_range(
+        self,
+        split: str,
+        start: Optional[int] = None,
+        stop: Optional[int] = None,
+    ) -> list[Any]: ...
+```
+
+Each method may be sync or `async`; the server awaits the result when it is awaitable, so you can back task discovery with async dataset clients.
+
+A "task spec" is deliberately untyped (`Any`). Return whatever your environment needs — a dict, a Pydantic model, a dataclass. The server converts it to JSON, handling Pydantic models, dataclasses, and plain objects. In practice a positional stub such as `{"id": "test-12", "index": 12, "split": "test"}` is enough, because `reset()` is what actually loads the row.
+
+### Two rules that are easy to miss
+
+1. **Task methods must be side-effect-free.** They are discovery, not control. They must not mutate episode state, consume a stream, or advance a cursor.
+2. **They must work on a freshly constructed environment.** Each task route builds a short-lived environment instance, calls the one method, and closes it again. Nothing you set up in `reset()` is available, so read configuration in `__init__` (or lazily inside the method) rather than relying on episode state.
+
+## HTTP Routes
+
+The routes are registered by `HTTPEnvServer` and appear in the server's OpenAPI schema under the `Task API` tag. All of them are namespaced by environment name except `/list_environments`.
+
+| Method | Route | Body | Response |
+|--------|-------|------|----------|
+| `GET`  | `/list_environments` | — | `["latex_ocr_env"]` |
+| `GET`  | `/{env_name}/splits` | — | `[{"name": "train", "type": "train"}, …]` |
+| `POST` | `/{env_name}/tasks` | `{"split": "test"}` | `{"tasks": [...], "env_name": "latex_ocr_env"}` |
+| `POST` | `/{env_name}/num_tasks` | `{"split": "test"}` | `{"num_tasks": 7595}` |
+| `POST` | `/{env_name}/task` | `{"split": "test", "index": 12}` | `{"task": {...}}` |
+| `POST` | `/{env_name}/task_range` | `{"split": "test", "start": 0, "stop": 32}` | `{"tasks": [...]}` |
+
+The request bodies are `ListTasksRequest`, `NumTasksRequest`, `GetTaskRequest`, and `GetTaskRangeRequest`. `start` and `stop` follow Python slice semantics: `start` is inclusive, `stop` is exclusive, and both may be omitted.
+
+These are HTTP-only. There are no WebSocket message types for task discovery — the `/ws` session protocol stays focused on `reset`, `step`, `state`, and `close`.
+
+### `env_name`
+
+`{env_name}` is the name passed to `create_app(..., env_name="latex_ocr_env")`. If you do not pass one, it defaults to the environment factory's class name (for example `LatexOCREnvironment`). Matching is case-insensitive, and any other name returns `404`. Pass `env_name` explicitly for any environment that implements the Task API — the generated URL is part of your public surface, and a class rename should not break a trainer.
+
+```python
+from openenv.core.env_server import create_app
+
+app = create_app(
+    LatexOCREnvironment,
+    LatexOCRAction,
+    LatexOCRObservation,
+    env_name="latex_ocr_env",
+)
+```
+
+### Split normalization
+
+`list_splits()` may return plain strings, dicts, or Pydantic models. The server normalizes each entry to `{"name": ..., "type": ...}`:
+
+- a dict or model is passed through as-is (after JSON conversion)
+- a string becomes `{"name": s, "type": s}` when `s` is `train`, `validation`, or `test`
+- any other string becomes `{"name": s, "type": "validation"}`
+
+So returning `["train", "holdout"]` yields `[{"name": "train", "type": "train"}, {"name": "holdout", "type": "validation"}]`. Return dicts yourself if you need to control the `type` for a custom split name.
+
+### Error semantics
+
+| Condition | Status |
+|-----------|--------|
+| `{env_name}` does not match the server's environment | `404 Not Found` |
+| Environment does not define the method, or it raises `NotImplementedError` | `501 Not Implemented` |
+| Method raises `IndexError` (index out of range) | `400 Bad Request` |
+
+Raise `IndexError` from `get_task()` for an out-of-range index so callers get a `400` rather than a `500`. Reserve `NotImplementedError` for capabilities the environment genuinely does not support — for example, random access on a streaming split.
+
+## Implementing It
+
+A dataset-backed environment reads its dataset in `__init__`, answers discovery from metadata, and loads the actual row in `reset()`:
+
+```python
+from typing import Any, Optional
+from uuid import uuid4
+
+from openenv.core.env_server import Environment
+from openenv.core.env_server.types import State
+
+
+class LatexOCREnvironment(Environment):
+    SPLITS = ["train", "test"]
+
+    def __init__(self, dataset_name: str = "unsloth/LaTeX_OCR") -> None:
+        super().__init__()
+        self.dataset_name = dataset_name
+
+    # --- Task API: discovery only, no side effects ---
+
+    def list_splits(self) -> list[str]:
+        return self.SPLITS
+
+    def num_tasks(self, split: str) -> int:
+        return len(self._load_split(split))
+
+    def list_tasks(self, split: str) -> list[dict[str, Any]]:
+        return [
+            {"id": f"{split}-{i}", "index": i}
+            for i in range(self.num_tasks(split))
+        ]
+
+    def get_task(self, split: str, index: int) -> dict[str, Any]:
+        n = self.num_tasks(split)
+        if index < 0 or index >= n:
+            raise IndexError(f"index {index} out of range for split {split} (n={n})")
+        return {"id": f"{split}-{index}", "index": index, "split": split}
+
+    def get_task_range(
+        self, split: str, start: Optional[int] = None, stop: Optional[int] = None
+    ) -> list[dict[str, Any]]:
+        n = self.num_tasks(split)
+        start = 0 if start is None else start
+        stop = n if stop is None else min(stop, n)
+        return [
+            {"id": f"{split}-{i}", "index": i, "split": split}
+            for i in range(start, stop)
+        ]
+
+    # --- Episode: this is where a task is actually loaded ---
+
+    def reset(
+        self,
+        split: str = "test",
+        index: Optional[int] = None,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> LatexOCRObservation:
+        if split not in self.list_splits():
+            raise ValueError(f"unknown split {split!r}; expected {self.list_splits()}")
+        self._state = State(episode_id=episode_id or str(uuid4()), step_count=0)
+        row = self._load_row(split, index, seed)
+        return LatexOCRObservation(image_base64=row["image"], done=False)
+```
+
+### Selecting a task in `reset()`
+
+`ResetRequest` allows extra fields, and both the HTTP `/reset` route and the WebSocket `reset` message forward unknown keys to your `reset()` — after filtering them against its signature. So naming the parameters `split` and `index` in `reset()` is all it takes for this to work:
+
+```python
+result = env.reset(split="test", index=12)
+```
+
+The corresponding wire calls are `POST /reset` with `{"split": "test", "index": 12}`, or a WebSocket `{"type": "reset", "data": {"split": "test", "index": 12}}`. The server filters incoming keys against your `reset()` signature, so a key it does not declare is dropped silently unless the signature also has `**kwargs`. A misspelled parameter therefore shows up as "the environment ignored my selection" rather than an error — worth checking first when a task selection appears not to take effect.
+
+When `index` is omitted, pick a row from `seed` so episodes stay reproducible.
+
+### Streaming and very large splits
+
+`list_tasks()` and `get_task_range()` can be asked for more rows than you want to materialize. Two defenses are worth building in:
+
+- return positional stubs rather than real rows, and cap how many you generate (`list_tasks()` returning a bounded preview is fine — `num_tasks()` still reports the honest total)
+- if a split is streamed sequentially and cannot be randomly accessed, raise `ValueError` or `NotImplementedError` from `reset(index=...)` and say which mode does support indexing
+
+Report the true count from `num_tasks()` even when `list_tasks()` is truncated. A trainer that shards work by `num_tasks()` needs the real denominator.
+
+## Consuming the Task API
+
+The core clients (`EnvClient`, `SyncEnvClient`) do not ship Task API methods, because task specs are environment-specific. Environment clients add thin HTTP helpers alongside the inherited Gym-style `reset()`/`step()`:
+
+```python
+import requests
+from urllib.parse import urljoin
+
+ENV_NAME = "latex_ocr_env"
+
+
+class LatexOCREnv(EnvClient[LatexOCRAction, LatexOCRObservation, State]):
+    def list_splits(self) -> list[str]:
+        resp = requests.get(urljoin(self._http_base(), f"{ENV_NAME}/splits"), timeout=30)
+        resp.raise_for_status()
+        return [s["name"] for s in resp.json()]
+
+    def num_tasks(self, split: str) -> int:
+        resp = requests.post(
+            urljoin(self._http_base(), f"{ENV_NAME}/num_tasks"),
+            json={"split": split},
+            timeout=60,
+        )
+        resp.raise_for_status()
+        return int(resp.json()["num_tasks"])
+```
+
+Note that the Task API is served over HTTP while episodes typically run over the WebSocket `/ws` session. A client that was constructed from a `ws://` URL has to derive the HTTP base URL before calling these routes.
+
+From a training or evaluation loop, discovery and episode control then compose naturally:
+
+```python
+with LatexOCREnv.from_docker_image("latex-ocr-env:latest") as env:
+    total = env.num_tasks("test")
+    for index in range(total):
+        result = env.reset(split="test", index=index)
+        prediction = model(result.observation.image_base64)
+        result = env.step(LatexOCRAction(latex=prediction))
+        print(index, result.reward)
+```
+
+Or with `curl`, against a locally running server:
+
+```bash
+curl http://localhost:8000/list_environments
+curl http://localhost:8000/latex_ocr_env/splits
+curl -X POST http://localhost:8000/latex_ocr_env/num_tasks \
+  -H 'Content-Type: application/json' -d '{"split": "test"}'
+curl -X POST http://localhost:8000/latex_ocr_env/task_range \
+  -H 'Content-Type: application/json' -d '{"split": "test", "start": 0, "stop": 4}'
+```
+
+## Imported Environments
+
+`openenv import` generates wrappers that already implement the Task API by delegating to the source environment's own task and split methods, and whose `reset()` accepts `split` / `index` / `task_spec`. If you are bringing in an ORS/OpenReward or Prime Intellect Verifiers environment, you get task discovery without writing any of the above — see the [CLI reference](../reference/cli.md) for `openenv import`.
+
+## Checklist
+
+Before shipping a dataset-backed environment:
+
+1. All five methods defined, or a deliberate `NotImplementedError` for the ones you cannot support.
+2. Every method works on a fresh instance, with no reliance on `reset()` having run.
+3. No side effects — no cursor advanced, no stream consumed, no episode state mutated.
+4. `env_name` passed explicitly to `create_app()`.
+5. `get_task()` raises `IndexError` for out-of-range indices.
+6. `num_tasks()` reports the honest total, even if `list_tasks()` truncates.
+7. `reset()` accepts `split` and `index` with the same names the trainer will use, and falls back to `seed`-driven selection when `index` is omitted.
+
+## Related Reading
+
+- [Concepts](concepts.md)
+- [RL Training](rl-integration.md)
+- [Core API](../reference/core.md)
+- [CLI reference](../reference/cli.md)

docs/source/reference/core.md

@@ -6,6 +6,8 @@ If you are trying to understand when OpenEnv exposes the training loop versus di
 
 For a high-level explanation of how MCP-backed environments move through `step()`, `step_async()`, and convenience tool helpers, see the [MCP environment lifecycle](../guides/mcp-environment-lifecycle.md) guide.
 
+For dataset-backed environments that publish enumerable tasks and splits, see the [Task API](../guides/task-api.md) guide.
+
 ## Server
 
 ### Environment server primitives
@@ -16,6 +18,8 @@ For a high-level explanation of how MCP-backed environments move through `step()
 
 [[autodoc]] openenv.core.env_server.interfaces.Transform
 
+[[autodoc]] openenv.core.env_server.interfaces.TaskProvider
+
 [[autodoc]] openenv.core.env_server.interfaces.Environment
 
 ### Types
@@ -50,6 +54,14 @@ For a high-level explanation of how MCP-backed environments move through `step()
 
 [[autodoc]] openenv.core.env_server.types.HealthResponse
 
+[[autodoc]] openenv.core.env_server.types.ListTasksRequest
+
+[[autodoc]] openenv.core.env_server.types.NumTasksRequest
+
+[[autodoc]] openenv.core.env_server.types.GetTaskRequest
+
+[[autodoc]] openenv.core.env_server.types.GetTaskRangeRequest
+
 [[autodoc]] openenv.core.env_server.types.WSResetMessage
 
 [[autodoc]] openenv.core.env_server.types.WSStepMessage