molt-rl 0.1.10__py3-none-any.whl

molt/__init__.py

@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

molt/agents/__init__.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Molt agents — Gymnasium-aligned public surface.
+
+User writes ONE of:
+  - `class MyEnv(Env)`         with async step() / reset() returning `Result`
+  - `class MyAgent(ChatAgent)`  with async run(ctx) returning `Result`
+
+User binds it via a Runner subclass and exports as `AgentRunner`:
+  - `class AgentRunner(StepEnvRunner): super().__init__(MyEnv)`
+  - `class AgentRunner(ChatAgentRunner): super().__init__(MyAgent)`
+
+For the chat-agent path, build your own client from the session URL — either
+`openai.AsyncOpenAI(base_url=ctx.base_url)` or
+`anthropic.AsyncAnthropic(base_url=ctx.session_url)`. We don't wrap the SDKs —
+use whichever directly; the server speaks both wires against the same engine.
+"""
+
+from molt.agents.base import Env, Result, Runner, StepEnvRunner, Trajectory
+from molt.agents.chat_agent import ChatAgent, ChatAgentRunner, ChatContext
+
+__all__ = [
+    "Env",
+    "ChatAgent",
+    "ChatAgentRunner",
+    "ChatContext",
+    "Result",
+    "Runner",
+    "StepEnvRunner",
+    "Trajectory",
+]

molt/agents/_chat_server.py

@@ -0,0 +1,466 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Internal chat forwarder over the vllm-router (imported by ``chat_agent.py``).
+
+The agent talks plain **chat messages** to a loopback endpoint with a stock OpenAI or
+Anthropic client. This server accumulates token-exact ``Trajectory`` segments per session (TITO),
+exactly like the Env step runner: a segment renders its first prompt CLIENT-SIDE (the same
+``apply_chat_template`` -> ``process_prompt_with_images`` path the RL dataset uses, so the
+image-placeholder count is self-consistent with the vision encoder); each continuation turn appends
+only the NEW turn's chat-template delta and generates over the EXACT accumulated tokens via the
+shared router transport (``RouterGenerateClient`` — ``/inference/v1/generate``). Because we train
+on the very tokens vLLM generated over, there is zero re-template drift even for models whose
+chat template doesn't round-trip generated bytes (e.g. a ``</think>`` trim), so no fragmentation.
+A normal multi-turn episode is ONE segment with one action range per turn; a fresh segment starts on
+a compaction (the agent rewrites/shortens history to fit context) or when the template re-renders
+prior turns non-prefix (e.g. Qwen3 drops a prior turn's <think> block once a newer user query
+follows). All segments share the one ``rollout_id`` + reward the trainer assigns (``stitch_session``).
+
+Three concerns, each a section below:
+  1. WIRE CODECS — OpenAI is canonical (identity decode); Anthropic is translated
+     to the same shape (``_decode_anthropic``) and the reply re-encoded per wire.
+  2. FORWARD — ``_run_turn`` grows the session trajectory by one turn and generates.
+  3. HTTP SERVER — uvicorn on the actor's OWN event loop + the session routes.
+
+Runs on the rollout runner actor's own event loop. A stock client reaches a session
+through the URL path prefix ``/s/<session_id>/v1`` that ``ChatAgentRunner`` builds.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from copy import deepcopy
+from dataclasses import dataclass, field
+from typing import Any
+from uuid import uuid4
+
+import uvicorn
+from fastapi import APIRouter, HTTPException, Request
+from fastapi.responses import JSONResponse
+
+from molt.agents.base import (
+    _TOKENIZE_EXECUTOR,
+    Trajectory,
+    _extract_generation_logprobs,
+    _tokenize_feedback,
+    _tokenize_observation,
+)
+from molt.utils.vlm_utils import (
+    estimate_vllm_input_expansion_delta,
+    load_images,
+    should_expand_image_placeholder,
+    split_image_placeholder,
+)
+
+
+# ===========================================================================
+# 1. Wire codecs — Anthropic request -> canonical (OpenAI) messages; replies back.
+#    OpenAI is the canonical shape, so its decode is identity (``lambda b: b``).
+# ===========================================================================
+def _content_to_text_and_images(content) -> tuple[str, list]:
+    """One OpenAI message's ``content`` -> (flat text with a literal ``<image>`` per image,
+    loaded PIL images). Mirrors the dataset convention (user content is ``"<image>...problem"``);
+    the client-side chat template + processor then expand each ``<image>`` exactly as the step
+    runner does. A bare-string content passes through with no images."""
+    if isinstance(content, str):
+        return content, []
+    if not isinstance(content, list):
+        return ("" if content is None else str(content)), []
+    parts, pil = [], []
+    for item in content:
+        if not isinstance(item, dict):
+            parts.append(str(item))
+        elif item.get("type") == "text":
+            parts.append(item.get("text") or "")
+        elif item.get("type") == "image_url":
+            ref = item.get("image_url") or {}
+            url = ref.get("url") if isinstance(ref, dict) else ref
+            imgs = load_images(url) if url else []
+            pil.extend(imgs)
+            parts.append("<image>" * len(imgs))
+    return "".join(parts), pil
+
+
+def _decode_anthropic(body: dict) -> dict:
+    """Anthropic Messages request -> canonical (OpenAI-shaped) body the router
+    understands. Top-level ``system`` becomes the first canonical message; content
+    blocks are normalized below."""
+    messages = [{"role": "system", "content": body["system"]}] if body.get("system") else []
+    for m in body.get("messages", []):
+        content = m.get("content")
+        if isinstance(content, list):
+            content = [out for b in content for out in _anthropic_blocks(b)]
+        messages.append({"role": m.get("role"), "content": content})
+    return {**body, "messages": messages}
+
+
+def _anthropic_blocks(block) -> list:
+    """One Anthropic content block -> list of canonical (OpenAI) blocks. image ->
+    ``image_url`` (base64 source -> data URI); ``tool_result`` -> its inner
+    text/image blocks (an Anthropic harness returns tool output as a ``tool_result``
+    block, so without unwrapping the observation would be dropped)."""
+    if not isinstance(block, dict):
+        return [block]
+    btype = block.get("type")
+    if btype == "image":
+        src = block.get("source") or {}
+        url = src.get("url") if src.get("type") == "url" else f"data:{src.get('media_type')};base64,{src.get('data')}"
+        return [{"type": "image_url", "image_url": {"url": url}}]
+    if btype == "tool_result":
+        inner = block.get("content")  # bare string, or a list of text/image blocks
+        if isinstance(inner, str):
+            return [{"type": "text", "text": inner}]
+        if isinstance(inner, list):
+            return [out for b in inner for out in _anthropic_blocks(b)]
+        return []
+    return [block]
+
+
+# OpenAI's ChatCompletion schema accepts only these finish_reason values and the SDK
+# validates them client-side; vLLM also emits abort/error/repetition (and None
+# mid-stream), so any non-OpenAI reason maps to "stop" purely to keep the reply
+# SDK-parseable (the agent never reads it; truncation uses the RAW reason elsewhere).
+_OPENAI_FINISH_REASONS = frozenset({"stop", "length", "tool_calls", "content_filter", "function_call"})
+
+
+def _chat_completion_body(model_name: str, content: str, finish_reason: str) -> dict:
+    return {
+        "id": f"chatcmpl-{uuid4().hex[:24]}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": model_name,
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": content},
+                "finish_reason": finish_reason if finish_reason in _OPENAI_FINISH_REASONS else "stop",
+            }
+        ],
+    }
+
+
+# Anthropic's Message schema validates stop_reason client-side too; map the raw vLLM
+# reason into its vocab (anything else -> "end_turn").
+_ANTHROPIC_STOP_REASONS = {"stop": "end_turn", "length": "max_tokens", "tool_calls": "tool_use"}
+
+
+def _anthropic_message_body(model_name: str, content: str, finish_reason: str) -> dict:
+    return {
+        "id": f"msg_{uuid4().hex[:24]}",
+        "type": "message",
+        "role": "assistant",
+        "model": model_name,
+        "content": [{"type": "text", "text": content}],
+        "stop_reason": _ANTHROPIC_STOP_REASONS.get(finish_reason, "end_turn"),
+        "stop_sequence": None,
+        "usage": {"input_tokens": 0, "output_tokens": 0},
+    }
+
+
+def _build_sampling_params(base_sampling, body: dict):
+    """One vLLM SamplingParams per chat call: inherit THIS rollout's defaults (its train-vs-eval
+    temperature / max_tokens — see ``_Session.sampling_params``), honor the caller's temperature /
+    max_tokens, and force logprobs on (importance-sampling correction needs token logprobs)."""
+    sp = deepcopy(base_sampling)
+    requested = body.get("max_completion_tokens", body.get("max_tokens"))
+    if requested is not None:
+        sp.max_tokens = int(requested)
+    if body.get("temperature") is not None:
+        sp.temperature = float(body["temperature"])
+    if body.get("top_p") is not None:
+        sp.top_p = float(body["top_p"])
+    if sp.logprobs is None:
+        sp.logprobs = 1
+    return sp
+
+
+# ===========================================================================
+# 2. Per-session state — token-exact Trajectory segments per session (TITO).
+# ===========================================================================
+@dataclass
+class _Session:
+    prompt: str
+    label: Any
+    images: list | None
+    # This rollout's id, reused as the router ``x-session-id`` so all turns (and each turn's
+    # render + generate) pin to one engine (consistent_hash) — see ``RouterGenerateClient.generate``.
+    session_id: str | None = None
+    # THIS rollout's SamplingParams (carries the train-vs-eval temperature / max_tokens for this
+    # call). Used as the per-turn default before the agent's own body overrides it; the fallback is
+    # ChatServerState.default_sampling, which is frozen at the FIRST execute — so an agent that
+    # doesn't re-send temperature would otherwise sample every later train/eval call at the
+    # first-caller's temperature (silent train/eval sampling-temperature mismatch).
+    sampling_params: Any = None
+    # Token-exact Trajectory segments for this episode. Normal multi-turn = ONE segment that later
+    # turns append to in place; a compaction (the agent rewrites/shortens history) or a non-prefix
+    # template re-render of prior turns starts a fresh segment. Empty until the first turn generates.
+    trajectories: list = field(default_factory=list)
+    lock: asyncio.Lock = field(default_factory=asyncio.Lock)  # serialize turns within a session
+    # Idempotent replay: a stock client retries a turn whose response was lost (e.g. a weight
+    # broadcast stalled the loop). The retry re-sends the SAME messages; we return the cached
+    # reply WITHOUT recording a second action, so retries can't double-count a turn.
+    last_messages: list | None = None
+    last_response: tuple | None = None
+
+
+class ChatServerState:
+    """Per-actor state: the unified HTTP transport (over the router) + processor +
+    live session accumulators. ``transport`` is the shared ``RouterGenerateClient``
+    (``.generate(token_ids, sp, mm)`` + the shared ``.http`` session, base-URL'd to the router)."""
+
+    def __init__(self, transport, processor, model_name: str, max_length: int, default_sampling):
+        self.transport = transport  # RouterGenerateClient: .generate(token_ids, sp, mm) over /inference/v1/generate
+        self.processor = processor
+        self.model_name = model_name
+        self.max_length = max_length
+        self.default_sampling = default_sampling
+        # omni3 uses a literal ``<image>`` (pass content through); qwen3.6-VL uses ``<|image_pad|>``
+        # and needs the structured-content split. Mirrors the RL prompt dataset's rendering.
+        self.expand_image_placeholder = should_expand_image_placeholder(processor)
+        self.sessions: dict[str, _Session] = {}
+
+    def open(self, session_id: str, prompt, label, images, sampling_params=None):
+        self.sessions[session_id] = _Session(
+            prompt=prompt, label=label, images=images, session_id=session_id, sampling_params=sampling_params
+        )
+
+    def discard(self, session_id: str):
+        self.sessions.pop(session_id, None)
+
+
+# ===========================================================================
+# 3. Forward one turn -> grow the session's ONE token-exact Trajectory.
+# ===========================================================================
+def _messages_to_chat(state: ChatServerState, messages: list) -> tuple[list, list]:
+    """OpenAI messages -> (ChatML rows with a literal ``<image>`` per image, loaded PIL images).
+    Split the literal placeholder into structured content for models whose image token isn't
+    ``<image>`` (Qwen-VL), so the template renders it — mirrors the RL dataset's rendering."""
+    chat, pil_images = [], []
+    for m in messages:
+        text, imgs = _content_to_text_and_images(m.get("content"))
+        chat.append({"role": m.get("role"), "content": text})
+        pil_images.extend(imgs)
+    if state.expand_image_placeholder:
+        chat = [split_image_placeholder(m) for m in chat]
+    return chat, pil_images
+
+
+async def _run_turn(state: ChatServerState, session: _Session, body: dict) -> tuple[str, str]:
+    """Grow the session's current token-exact Trajectory segment by one turn, then generate over its
+    EXACT accumulated tokens (TITO). A new segment (full CLIENT-SIDE render, the same
+    ``apply_chat_template`` -> ``process_prompt_with_images`` path the RL dataset uses) starts the
+    episode AND resumes after a compaction; a normal continuation appends only the new chat-template
+    delta, so the trained tokens stay identical to what vLLM generated over — no re-template drift.
+    Returns ``(action_text, finish_reason)``."""
+    messages = body.get("messages") or []
+    if not messages:
+        raise HTTPException(status_code=400, detail="No messages provided")
+    if messages == session.last_messages:  # retry of an already-recorded turn -> replay, don't re-record
+        return session.last_response
+    sp = _build_sampling_params(session.sampling_params or state.default_sampling, body)
+    kwargs = {"tools": body["tools"]} if body.get("tools") else {}
+    loop = asyncio.get_running_loop()
+
+    # Does this turn continue the running segment, or did the agent rewrite/shorten the history
+    # (compaction)? A normal turn only appends (assistant + new user/tool) onto the SAME prefix, so
+    # its messages extend last_messages; a compaction changes that prefix. Also false on turn 1.
+    extends = bool(session.trajectories) and messages[: len(session.last_messages)] == session.last_messages
+    chat, pil_images = _messages_to_chat(state, messages)
+    full = state.processor.apply_chat_template(chat, tokenize=False, add_generation_prompt=True, **kwargs)
+    if extends:
+        prefix = state.processor.apply_chat_template(chat[:-1], tokenize=False, add_generation_prompt=False, **kwargs)
+        # The suffix slice below is only sound when the prior turns render the same in `full` and
+        # `prefix`. A template can instead rewrite history in place — e.g. Qwen3 re-renders a prior
+        # assistant turn WITHOUT its <think> block once a newer user query follows — and the slice
+        # would silently drop or garble this turn's feedback. Treat that like a compaction.
+        extends = full.startswith(prefix)
+    if not extends:
+        # New segment (episode start OR post-compaction/re-render): render the prompt in full +
+        # tokenize (image count self-consistent with the vision encoder). process_prompt_with_images
+        # is a multi-second CPU op per image; run it off the loop so other sessions' generations overlap.
+        prompt_ids, mm_train_inputs, pil_images = await loop.run_in_executor(
+            _TOKENIZE_EXECUTOR, _tokenize_observation, state.processor, full, pil_images
+        )
+        image_budget = (
+            estimate_vllm_input_expansion_delta(state.processor, prompt_ids, mm_train_inputs, pil_images)
+            if pil_images
+            else 0
+        )
+        remaining = state.max_length - len(prompt_ids) - image_budget
+        if remaining <= 0:  # prompt alone fills the context -> generate nothing; stitch drops empty rollouts
+            session.last_messages, session.last_response = messages, ("", "length")
+            return "", "length"
+        # Built but NOT yet attached to the session: a failed generate below must leave the session
+        # untouched (see the commit note) so a client retry re-runs cleanly.
+        traj = Trajectory(
+            prompt=session.prompt,
+            label=session.label,
+            images=session.images,
+            observation_text="",
+            observation_tokens=list(prompt_ids),
+            mm_train_inputs=mm_train_inputs,
+            pil_images=pil_images,
+            image_budget=image_budget,
+            rollout_log_probs=[0.0] * len(prompt_ids),
+        )
+        gen_tokens = traj.observation_tokens
+    else:
+        # Continuation: the new turn's tokens are the chat-template STRING suffix — the prior-turn
+        # renders land in the shared prefix and cancel (verified above), so the suffix is the new
+        # user/tool turn + generation prompt (committed as masked feedback -> token-exact segment).
+        traj = session.trajectories[-1]
+        delta_text = full[len(prefix) :]
+        new_pil = _content_to_text_and_images(messages[-1].get("content"))[1]  # images added THIS turn
+        delta_ids = await loop.run_in_executor(
+            _TOKENIZE_EXECUTOR, _tokenize_feedback, state.processor, delta_text, new_pil, traj, state.max_length
+        )
+        # Generate over the accumulated tokens PLUS this turn's delta, but do NOT append the delta yet.
+        gen_tokens = traj.observation_tokens + delta_ids
+        remaining = state.max_length - len(gen_tokens) - traj.image_budget
+        if remaining <= 0:  # grown context fills the window -> end truncated; prior turns are kept
+            traj.truncated = True
+            session.last_messages, session.last_response = messages, ("", "length")
+            return "", "length"
+
+    # Clamp generation to the remaining budget (the step runner clamps the same way).
+    if sp.max_tokens is None or sp.max_tokens > remaining:
+        sp.max_tokens = remaining
+    if getattr(sp, "min_tokens", 0) and sp.min_tokens > sp.max_tokens:
+        sp.min_tokens = sp.max_tokens
+
+    mm_data = {"image": traj.pil_images} if traj.pil_images else None
+    request_output, off_policy_len = await state.transport.generate(
+        gen_tokens, sp, multi_modal_data=mm_data, session_id=session.session_id
+    )
+    generation = request_output.outputs[0]
+    action_ids = list(generation.token_ids)
+    # /inference/v1/generate is token-only (generation.text == ""); decode the action text from the
+    # ids so the agent/grader sees it. skip_special_tokens=False keeps answer markers (<answer>,
+    # \boxed, tool tags) — matches observation_text decoding in base.py.
+    action_text = generation.text or (
+        state.processor.decode(action_ids, skip_special_tokens=False) if action_ids else ""
+    )
+    finish_reason = generation.finish_reason or "stop"
+    # a non-empty completion MUST carry aligned per-token logprobs — fails fast otherwise (IS correction).
+    action_logprobs = _extract_generation_logprobs(action_ids, generation.logprobs) if action_ids else []
+
+    # Commit the turn only now that generate succeeded. A failed generate above left the session
+    # untouched, so the stock client's retry (same messages) re-runs this turn cleanly — no
+    # double-appended delta, and no half-built segment to trip the extends check / crash on turn 1.
+    if not extends:
+        session.trajectories.append(traj)
+    else:
+        prev_action_text = session.last_response[0] if session.last_response else ""
+        traj.append_feedback(prev_action_text, delta_text, delta_ids)
+    traj.truncated = traj.truncated or finish_reason == "length"
+    # off_policy_len: leading tokens generated under stale weights when a broadcast landed
+    # mid-request (per-token IS still corrects them); the transport reports it, same as the step runner.
+    traj.append_action(action_ids, action_logprobs, off_policy_len=off_policy_len)
+    # R3: absorb this turn's prefill+gen routing by absolute position (first-writer-wins). The full
+    # prefill covers the prior action's trailing token too, so the turn-boundary backfills naturally.
+    traj.absorb_routing(request_output)
+    session.last_messages, session.last_response = messages, (action_text, finish_reason)
+    return action_text, finish_reason
+
+
+def stitch_session(state: ChatServerState, session_id: str, result):
+    """Return the rollout's token-exact trajectory segments (usually ONE; more when the agent
+    compacted history mid-episode or the template re-rendered prior turns non-prefix), each stamped
+    with the terminal reward/score/info. Multi-turn
+    actions are separate action ranges within a segment (like the step runner); all segments share the
+    one ``rollout_id`` + reward the trainer assigns, so group baselines dedup them to one reward per
+    rollout. Empty list when nothing was generated (the first prompt alone exceeded max_length)."""
+    session = state.sessions.get(session_id)
+    if session is None:
+        raise RuntimeError(f"Unknown session {session_id} at stitch")
+    for traj in session.trajectories:
+        traj.reward = result.reward
+        traj.scores = result.score if result.score is not None else result.reward
+        traj.extra_logs = result.info or {}
+        if result.images is not None:
+            traj.images = result.images
+    # The agent's own cut (turn cap, budget) ends the episode, so it lands on the last segment
+    # only; truncated_rate counts segments, and the earlier ones were generated in full.
+    if session.trajectories and result.truncated:
+        session.trajectories[-1].truncated = True
+    return session.trajectories
+
+
+# ===========================================================================
+# 4. HTTP server — uvicorn on the runner actor's own event loop + session routes.
+# ===========================================================================
+class _AutoPortServer(uvicorn.Server):
+    """uvicorn Server that reports the OS-assigned port when configured ``port=0``."""
+
+    def __init__(self, config: uvicorn.Config):
+        super().__init__(config)
+        self.actual_port: int | None = None
+        self._ready = asyncio.Event()
+
+    async def startup(self, sockets=None) -> None:
+        try:
+            await super().startup(sockets=sockets)
+            if self.servers and self.config.port == 0:
+                self.actual_port = self.servers[0].sockets[0].getsockname()[1]
+            else:
+                self.actual_port = self.config.port
+        finally:
+            self._ready.set()
+
+    async def get_port(self) -> int | None:
+        await self._ready.wait()
+        return self.actual_port
+
+
+async def serve_on_current_loop(app, host: str = "127.0.0.1") -> tuple[int, asyncio.Task]:
+    """Start uvicorn as a task on the CURRENT event loop. Returns ``(port, task)``."""
+    config = uvicorn.Config(app, host=host, port=0, log_level="warning")
+    server = _AutoPortServer(config)
+    task = asyncio.create_task(server.serve())
+    port = await server.get_port()
+    if port is None:
+        await task  # surfaces the startup exception
+        raise RuntimeError("Chat server failed to start (no port reported)")
+    return port, task
+
+
+def mount_session_capture(app, state: ChatServerState) -> None:
+    """Mount the session-scoped chat routes. The OpenAI and Anthropic wires forward
+    the SAME turn (``_run_turn``); they differ only in how the body is decoded to the
+    canonical shape and how the reply is encoded back."""
+    router = APIRouter()
+
+    async def _serve(session_id, request, decode, encode):
+        body = decode(await request.json())
+        session = state.sessions.get(session_id)
+        if session is None:
+            raise HTTPException(status_code=400, detail=f"Unknown or unopened session {session_id}")
+        # Serialize turns within a session: a stock client may retry, which would
+        # otherwise run a second turn concurrently and double-record the sample.
+        async with session.lock:
+            action_text, finish_reason = await _run_turn(state, session, body)
+        return JSONResponse(encode(state.model_name, action_text, finish_reason))
+
+    @router.post("/s/{session_id}/v1/chat/completions")
+    async def chat_sess(session_id: str, request: Request):
+        return await _serve(session_id, request, lambda b: b, _chat_completion_body)
+
+    @router.post("/s/{session_id}/v1/messages")
+    async def messages_sess(session_id: str, request: Request):
+        return await _serve(session_id, request, _decode_anthropic, _anthropic_message_body)
+
+    app.include_router(router)