refactor: unify tool scheduling, recording, and recipe export

docs/source-en/rst_source/development/add_primitive.rst

@@ -28,58 +28,66 @@ Two types of primitives
      - ``move_to``, ``rotate_wrist``, ``release``,
        ``back_project``
 
-From the LLM's perspective, both types expose the same interface: a
-tool schema, a primitives method, and a state dump after the
-call. They differ only in how the method is implemented.
+Both types are native tools: a typed handler receives the session's resources
+through ``ToolContext`` and returns ``ToolResult``. The toolkit validates
+arguments and captures the post-action observation.
 
 Add a scripted primitive
 ------------------------
 
-Adding a scripted primitive usually involves two steps:
+Define a module-level handler in ``robots/<robot>/tools.py`` and add the
+resulting ``Tool`` to the robot's tool tuple. For example, this LIBERO handler
+holds the current pose for a bounded number of environment steps:
 
-1. **Add a method to the primitives.** Add the method to the
-   current robot's primitives class, such as
-   ``LiberoPrimitives`` or ``MyRobotPrimitives``. The method accepts
-   the tool-call arguments, performs the work, usually through one or
-   more ``self._env.step(...)`` calls, and returns a small log ``dict``.
-
-  Primitive methods capture and re-render state (``get_env_state``)
-  automatically after they run:
-
-   .. code-block:: python
-
-      def open_drawer(self, dx: float = 0.15) -> dict:
-          # Move end-effector back by dx while gripper is closed.
-          for _ in range(N):
-              self._env.step(build_open_drawer_chunk(dx))
-          return {"ok": True, "dx": dx}
-
-  You can mark read-only tools (``view_env_state``, ``back_project``, ``segment``,
-  ...) with :func:`~rpent.tools.toolkit.readonly` so the toolkit skips state
-  capture for them, improving performance.
-
-2. **Add the tool schema.** Add an entry to ``TOOLS_SPEC`` in
-   ``robots/<robot>/tools.py``:
-
-   .. code-block:: python
-
-      {
-          "name": "open_drawer",
-          "description": "Pull the currently-grasped drawer handle "
-                         "backwards by ``dx`` meters.",
-          "input_schema": {
-              "type": "object",
-              "properties": {"dx": {"type": "number"}},
-              "required": [],
-          },
-      }
-
-Once both exist, the toolkit registers the tool automatically: it iterates
-``TOOLS_SPEC`` and binds each spec to the matching primitive-driver method
-(e.g. ``getattr(self._primitives, name)``).
+.. code-block:: python
 
-After these steps, the ``api``, ``claude_code``, and ``codex`` planners
-can all call the primitive without any other code changes.
+   from typing import Annotated
+
+   from pydantic import Field
+
+   from rpent.tools import ToolContext, ToolResult, tool
+
+   @tool
+   def hold_pose(
+       steps: Annotated[int, Field(ge=1, le=100)] = 10,
+       *,
+       ctx: ToolContext,
+   ) -> ToolResult:
+       """Hold the current pose with the gripper closed.
+
+       Args:
+           steps: Number of environment steps.
+       """
+       runtime = ctx.robot
+       for _ in range(steps):
+           ctx.check_cancelled()
+           obs, _, terminated, truncated, _ = runtime.env.step(
+               [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]
+           )
+           runtime.executed_steps += 1
+           runtime.set_obs(obs)
+           ctx.record_frame(obs["main_images"])
+           if terminated or truncated:
+               break
+       return ToolResult(data={"steps_requested": steps})
+
+   # Add hold_pose to the existing LIBERO_TOOLS tuple.
+
+Use the concrete runtime type in ``ToolContext[LiberoRuntime]`` in robot code.
+``@tool`` generates the parameter model and JSON schema from the same function
+signature. Google-style docstrings provide descriptions; use ``Annotated`` /
+``Field`` for constraints on model-supplied arguments. ``ctx`` is injected by
+the executor and is not part of the published schema.
+
+After adding the declaration to ``LIBERO_TOOLS``, all three planners can call
+it. The toolkit handles capture through ``_capture_observation``; handlers
+submit frames but do not save their own episode video or duplicate the state dump.
+
+For a tool that reads existing observations, place ``@readonly`` below
+``@tool`` to allow concurrent execution and skip automatic capture.
+``write_text_file`` and ``finish`` have no readonly marker and run exclusively;
+the executor skips observation capture for common tools and ``finish``.
+See :doc:`interfaces` for cancellation and scheduling.
 
 .. _add-primitive-model-based:
 
@@ -113,44 +121,26 @@ primitive requires a few additional components:
    ``rpent.robots.components.pi05_vla_client.Pi05VLAClient`` for the LIBERO
    implementation.
 
-3. **Add a method to the primitives.** In the current
-   robot's primitives class, call the model client, pass
-   the returned action chunk to the environment, and return a log
-   ``dict``. The model client API is
-   :meth:`rpent.robots.components.pi05_vla_client.Pi05VLAClient.predict`,
-   which reads the instruction from ``env_obs["task_descriptions"]`` and
-   returns a ``[chunk, action_dim]`` numpy action chunk (batch dim already
-   stripped):
-
-   .. code-block:: python
-
-      def mymodel_pick(self, target: str) -> dict:
-          env_obs = self._env.get_obs()
-          env_obs["task_descriptions"] = f"pick {target}"
-          chunk = self._model.predict(env_obs)
-          self._env.chunk_step(chunk)
-          return {"model": "mymodel", "target": target}
-
-4. **Add the tool schema and register it in the toolkit.** Follow the
-   same pattern as for a scripted primitive.
-
-5. **Wire the components together in ``robot_spec.py``.** The
-   robot's ``get_toolkit`` builds the toolkit with
-   ``runtime_kwargs``:
-
-   .. code-block:: python
-
-      def get_toolkit(*, runtime_kwargs, dashboard_events):
-          from robots.myrobot.toolkit import MyRobotToolkit
-          return MyRobotToolkit(
-              runtime_kwargs=runtime_kwargs,
-              dashboard_events=dashboard_events,
-          )
-
-   The robot package's ``_init_runtime`` builds
-   ``runtime_kwargs``, for example
-   ``{"env": MyRobotEnvClient(...), "model": MyModelClient(...)}``.
-   The toolkit constructor then forwards it to the primitives.
+3. **Write a native tool handler.** Use ``ctx.robot`` to access the model and
+   environment clients, request a prediction, execute the returned actions,
+   and return ``ToolResult(data={...})``. Follow the robot's observation and
+   action conventions: Pi0.5 reads the instruction from
+   ``env_obs["task_descriptions"]`` and returns a ``[chunk, action_dim]``
+   NumPy array. Check cancellation before inference and at safe action
+   boundaries, and submit environment frames through ``ctx.record_frame``.
+   See ``pi0_pick`` in ``robots/libero/tools.py`` and ``rldx_skill`` in
+   ``robots/robocasa/tools.py`` for concrete implementations.
+
+4. **Add the tool to the robot's tuple.** The toolkit receives that tuple in
+   its constructor and handles observation capture after execution, just as
+   for a scripted primitive.
+
+5. **Wire the clients in ``robot_spec.py``.** ``_init_runtime`` returns
+   ``(owned_daemons, runtime_kwargs)``, with entries such as ``env`` and
+   ``model``. ``get_toolkit(*, runtime_kwargs, dashboard_events, config)``
+   passes those inputs, ``config.output_dir``, and a ``MemoryManager`` to the
+   robot toolkit. The toolkit constructs the session runtime from
+   ``runtime_kwargs``. See :doc:`add_robot` for the complete factory.
 
 Reuse an existing vla_server across runs
 ----------------------------------------
@@ -190,7 +180,7 @@ parts:
   the server-side handler by the facade — the client does **not** pass it,
   and must not forge ``session_ids`` inside ``predict``'s ``options``.
 
-- **Primitives side**: call ``reset_session`` before a task starts to clear
+- **Runtime / tool side**: call ``reset_session`` before a task starts to clear
   policy state left over from the previous episode, so consecutive runs do
   not leak state into each other.
 
@@ -234,11 +224,12 @@ Design principles for a new primitive
 
 - **Tools describe intent, not motion.** A good tool name is
   ``pi0_pick``, not ``execute_action_chunk_of_length_20``.
-- **Every tool ends with a state dump.** The next turn depends on
-  the state dump reflecting the post-action world. Don't let the
-  primitive return before the render finishes.
-- **Return small dicts.** Tool return values are fed back to the LLM
-  as text. Save larger observations through ``EnvState.save``; ``EnvState``
+- **Action calls include a fresh observation.** The toolkit captures it after
+  the handler finishes and before returning to the planner. Readonly tools
+  reuse recorded observations.
+- **Return small ``ToolResult.data`` payloads.** The planner serializes them
+  as text and sends ``ToolResult.images`` as PNG content. Save larger observations
+  through ``EnvState.save``; ``EnvState``
   automatically records each logical base name in its owned
   ``StepRecord.artifacts`` set. Expose images through ``view_env_state`` and
   geometry through environment tools rather than returning raw paths.
@@ -262,5 +253,5 @@ The same pattern extends to non-VLA model primitives:
   head to call via a ``model`` kwarg on ``predict``.
 
 Regardless of the implementation, the framework contract remains
-unchanged: model process → model client → primitives method →
-tool schema → ``Toolkit.add_tool``.
+unchanged: model process → model client → native ``@tool`` handler →
+robot tool tuple → ``Toolkit.execute_tool``.

docs/source-en/rst_source/development/add_robot.rst

@@ -57,7 +57,8 @@ For a new robot named ``myrobot``, use the following directory layout:
        robot_spec.py          # RobotSpec, factories, Dashboard spec, runtime hooks
        env_client.py          # MyEnvClient — agent-side RPC stub (§1)
        prompt_bundle.py       # system()/user() prompt factories         (§2)
-       toolkit.py             # MyRobotToolkit + primitives + tool definitions (§3)
+       toolkit.py             # Runtime, MyRobotToolkit, observations (§3)
+       tools.py               # Native tool declarations and handlers (§3)
        env_server.py          # server-side facade + RPC server (§1)
        vla_server.py          # (optional) VLA model server
 
@@ -100,6 +101,7 @@ these two functions:
        from robots.myrobot.toolkit import MyRobotToolkit
        return MyRobotToolkit(
            runtime_kwargs=runtime_kwargs,
+           output_dir=config.output_dir,
            dashboard_events=dashboard_events,
            memory=MemoryManager(
                root=config.prompt_vars.get("memory_dir") or get_memory_dir("myrobot"),
@@ -279,69 +281,69 @@ render time.
 3. ``toolkit.py``
 ------------------
 
-This module owns everything the LLM can call: the tool schemas, the primitives,
-the per-step state dump, and the MCP allowlist. (In the LIBERO robot these
-are split between ``tools.py`` and ``toolkit.py`` for historical reasons; for a
-new robot it is fine to keep them all in ``toolkit.py``.)
-
-A toolkit module typically contains four pieces:
-
-**Primitives class** (e.g. ``MyRobotPrimitives``) — a Python object owned
-by the toolkit. It holds the ``EnvClient``, the VLA ``model`` client, and any
-state needed for the current run. It exposes one method per primitive tool
-(``move_to``, ``pi0_pick``, ``release``, …), with each method returning a
-``dict`` log.
-
-**Tool definitions and handlers** — a module-level ``TOOLS_SPEC`` list of
-Anthropic-style tool definitions (``name``, ``description``, ``input_schema``),
-plus any module-level functions referenced by the toolkit (e.g.
-``view_env_state``, ``back_project``, ``finish``).
-
-**Per-step state dump** — ``dump_state(driver, env_state, log)`` opens
-``env_state.record_step(...)`` and receives the allocated step index; the
-``StepRecord`` is appended and committed immediately. Save large observations
-through ``env_state.save(...)`` — inside a ``record_step`` block the ``step``
-argument may be omitted (it defaults to the new step), pass an explicit
-``step=<int>`` to target a different step, and ``step=None`` for run-level
-artifacts. ``EnvState`` adds every successfully saved base name to the step's
-flat ``artifacts`` set automatically. Readers use the canonical artifact
-filenames rather than maintaining a parallel observation index.
-
-**Toolkit class** — subclass ``rpent.tools.toolkit.Toolkit``:
-
-- forward ``memory`` (a :class:`~rpent.memory.MemoryManager`) and ``state`` to
-  ``super().__init__(...)``. Configure ``memory_access`` and
-  ``inbox_cell_tag`` on the ``MemoryManager``; eval uses read-only access by
-  default.
-- build the primitives in ``__init__`` through a custom initialization
-  helper (named ``init_primitives`` in LIBERO; it calls
-  ``EnvState.reset()``, constructs the primitives, and dumps step 0),
-- register each tool with ``self.add_tool(name, spec, handler)`` — stateless
-  readers (``view_env_state``, ``finish``, …) bind directly to module-level
-  functions; primitive tools route through ``_step(name, **kwargs)`` which
-  calls ``getattr(self._primitives, name)(**kwargs)`` and re-renders state,
-- override ``close()`` to save remaining agent-side artifacts through
-  ``EnvState`` (for example ``state.save("episode.mp4", frames, step=None)``).
-
-``runtime_kwargs`` (forwarded from ``robot_spec.py:get_toolkit``) is the dict
-the toolkit passes verbatim to your primitives' ``__init__`` — typically
-``{"env": MyEnvClient(...), "model": VLAClient(...), ...}``.
+Split robot behavior between ``tools.py`` and ``toolkit.py``, following
+``robots/libero/``, ``robots/robocasa/``, or ``robots/robotwin/``:
+
+**Runtime** — a per-session object such as ``MyRobotRuntime`` holds the env
+and model clients, cached observations, and task progress. Construct it from
+``runtime_kwargs`` returned by ``init_runtime``. Tools access it through
+``ctx.robot``.
+
+**Native tools** — define typed functions in ``tools.py`` with ``@tool`` and a
+required keyword-only ``ctx: ToolContext[MyRobotRuntime]``. Return
+``ToolResult`` and collect the declarations in a tuple such as
+``MYROBOT_TOOLS``. Tool schemas come from signatures and Google-style
+docstrings; planner adapters handle SDK and MCP serialization. See
+:doc:`add_primitive` for an example.
+
+**Observations** — ``dump_state(runtime, env_state, log)`` opens
+``env_state.record_step(...)`` to allocate and commit a ``StepRecord``.
+Save arrays, images, and metadata through ``env_state.save(...)``: inside
+the block, omitting ``step`` targets the new step; an explicit integer targets
+another step and ``step=None`` creates a run-level artifact. ``EnvState``
+adds each saved logical name to the record's ``artifacts`` set.
+``build_observation(state, record)`` returns JSON data and ordered PNG bytes
+for both action responses and ``view_env_state``.
+
+**Toolkit** — subclass ``Toolkit[MyRobotRuntime]``:
+
+- Construct the runtime and a fresh ``EnvState(output_dir)``; pass ``state``,
+  ``memory``, ``robot=runtime``, ``output_dir``, ``tools=MYROBOT_TOOLS``, and
+  ``dashboard_events`` to ``super().__init__``. Configure ``memory_access`` and
+  ``inbox_cell_tag`` on ``MemoryManager``; evaluation defaults to read-only memory.
+- Initialize the environment and record step 0. Publish its ``StepRecordEvent``
+  so the Dashboard can display the initial observation.
+- Implement ``_capture_observation(*, command, result, elapsed_s)`` using
+  ``dump_state`` and ``build_observation``. Save the original ``result.to_dict()``
+  in the step log. Return an observation containing the step, artifact names,
+  state, and action log, together with its PNG images. The base executor calls
+  this after non-readonly robot handlers, including failed handlers, except
+  ``finish``. Common tools also skip capture.
+- Implement ``solved()`` using the environment's success criterion. Provide a
+  robot-specific ``finish`` tool returning ``status`` and ``summary``; the
+  executor stores accepted results in ``toolkit.finish_result``.
+
+Tools submit environment-step RGB frames via ``ctx.record_frame(rgb)``.
+The base toolkit owns the frame buffer, Dashboard action clips, episode video,
+recipe export, and cancellation lifecycle. Use the inherited ``close()`` for
+recording cleanup; if a robot needs additional cleanup, preserve the base call.
 
 Conventions worth keeping
 -------------------------
 
-- ``output_dir`` is the working directory that the runner creates for each
-  run. Environment observations are owned by ``EnvState``; callers use logical
-  base names and never construct storage paths. Transcripts and other
-  run-management outputs share the same run directory.
-- Tool definitions use the Anthropic format (``name`` / ``description`` /
-  ``input_schema``). Every tool registered with ``self.add_tool(...)`` is
-  exposed to all planners.
-- Server-side return values must be picklable and torch-free.
-- Each primitive tool dumps a fresh state snapshot after running so the next
-  ``view_env_state`` call reflects the post-action world.
-- Treat ``dump_state`` as the source of truth for what the agent sees — any new
-  modality (e.g. tactile, force) goes through it.
+- ``output_dir`` is the runner-created working directory. Environment artifacts
+  are managed by ``EnvState`` through logical names; transcripts share the same
+  run directory.
+- ``@readonly`` below ``@tool`` enables shared execution and skips automatic
+  observation capture. Keep actions and ``finish`` exclusive and call
+  ``ctx.check_cancelled()`` at safe boundaries in long loops.
+- Server-side values use transport-supported Python / NumPy types and remain
+  torch-free.
+- Add new observation modalities in ``dump_state`` and expose them through
+  ``build_observation``. ``view_env_state`` reads a saved step without stepping
+  or re-rendering the environment.
+
+See :doc:`interfaces` for the full tool execution and lifecycle contract.
 
 .. _add-robot-config:
 
@@ -404,7 +406,7 @@ validates those fields and returns a
   by this process. The active runner stops them during cleanup. A client for an
   external endpoint must not add that external service to this list.
 - ``runtime_kwargs: dict`` is passed to the toolkit constructor, which
-  forwards it to the primitives' ``__init__``. A complete set commonly
+  uses it to construct the robot runtime. A complete set commonly
   contains ``{"env": MyEnvClient(...), "model": VLAClient(...)}`` plus any
   supporting clients.