synth-ai 0.2.8.dev4__py3-none-any.whl → 0.2.23.dev3__py3-none-any.whl

examples/warming_up_to_rl/task_app/synth_envs_hosted/rollout.py

@@ -0,0 +1,1891 @@
+from __future__ import annotations
+
+import contextlib
+import json
+import logging
+import os
+import time as _time
+from datetime import datetime
+from typing import Any
+
+from fastapi import APIRouter, HTTPException, Request, status
+from pydantic import BaseModel
+from synth_ai.tracing_v3 import BaseLMResponse
+from synth_ai.task.tracing_utils import unique_sft_path
+from synth_ai.tracing_v3.abstractions import EnvironmentEvent, LMCAISEvent, TimeRecord
+from synth_ai.tracing_v3.llm_call_record_helpers import create_llm_call_record_from_response
+from synth_ai.tracing_v3.session_tracer import SessionTracer
+
+from .registry import registry
+
+logger = logging.getLogger(__name__)
+
+
+# --- Seeding utilities (robust, optional deps) ---
+def _set_global_seed(seed_value: int) -> dict[str, Any]:
+    """Set global RNG seeds across common libraries; return details for logging/restoration.
+
+    Returns a dict containing which libraries were seeded and prior states if obtainable.
+    """
+    seeded: dict[str, Any] = {"seed": int(seed_value), "libs": []}
+    with contextlib.suppress(Exception):
+        import random as _random  # type: ignore
+
+        _random.seed(seed_value)
+        seeded["libs"].append("random")
+    with contextlib.suppress(Exception):
+        import numpy as _np  # type: ignore
+
+        _np.random.seed(seed_value)
+        seeded["libs"].append("numpy")
+    with contextlib.suppress(Exception):
+        import torch as _torch  # type: ignore
+
+        if hasattr(_torch, "manual_seed"):
+            _torch.manual_seed(seed_value)
+            seeded["libs"].append("torch")
+        # Make CUDA deterministic if present (best-effort)
+        with contextlib.suppress(Exception):
+            if getattr(_torch, "cuda", None) and _torch.cuda.is_available():
+                _torch.cuda.manual_seed_all(seed_value)
+                seeded.setdefault("cuda", True)
+        # CUDNN deterministic flags (optional)
+        with contextlib.suppress(Exception):
+            if getattr(_torch, "backends", None) and getattr(_torch.backends, "cudnn", None):
+                _torch.backends.cudnn.deterministic = True  # type: ignore[attr-defined]
+                _torch.backends.cudnn.benchmark = False  # type: ignore[attr-defined]
+    return seeded
+
+
+def _clear_seed_side_effects() -> None:
+    """Best-effort cleanup to avoid global deterministic side-effects between requests."""
+    # We cannot truly restore prior RNG states without capturing them; we just avoid
+    # leaving aggressive deterministic flags enabled where it matters.
+    with contextlib.suppress(Exception):
+        import torch as _torch  # type: ignore
+
+        with contextlib.suppress(Exception):
+            if getattr(_torch, "backends", None) and getattr(_torch.backends, "cudnn", None):
+                # Re-enable cudnn.benchmark default True only if it was True; safest is False -> leave as is.
+                # We'll keep deterministic False to avoid global impact; benchmark left False for stability.
+                _torch.backends.cudnn.deterministic = False  # type: ignore[attr-defined]
+
+
+router = APIRouter()
+
+
+class RolloutEnvSpec(BaseModel):
+    env_id: str | None = None
+    env_name: str | None = None
+    config: dict[str, Any] = {}
+    seed: int | None = None
+
+
+class RolloutPolicySpec(BaseModel):
+    policy_id: str | None = None
+    policy_name: str | None = None
+    config: dict[str, Any] = {}
+
+
+class RolloutBranchConfig(BaseModel):
+    branch_every_n_steps: int = 0
+    branch_on_condition: str | None = None
+    max_branches: int = 0
+    branch_policy: bool = False
+    branch_env: bool = False
+
+
+class RolloutRecordConfig(BaseModel):
+    trajectories: bool = True
+    logprobs: bool = False
+    value: bool = False
+    return_trace: bool = False
+    trace_format: str = "compact"
+
+
+class RolloutSafetyConfig(BaseModel):
+    max_ops: int = 100000
+    max_time_s: float = 3600.0
+
+
+class RolloutRequest(BaseModel):
+    run_id: str
+    env: RolloutEnvSpec
+    policy: RolloutPolicySpec
+    ops: list[str]  # ["agent", "env", ...]
+    record: RolloutRecordConfig = RolloutRecordConfig()
+    on_done: str = "reset"  # "reset" | "terminate"
+    branch: RolloutBranchConfig | None = None
+    safety: RolloutSafetyConfig = RolloutSafetyConfig()
+    # Optional run/session context
+    training_session_id: str | None = None
+    synth_base_url: str | None = None
+
+
+class RolloutStep(BaseModel):
+    obs: dict[str, Any]
+    tool_calls: list[dict[str, Any]]
+    reward: float | None = None
+    done: bool = False
+    truncated: bool | None = None
+    logprob: float | None = None
+    value: float | None = None
+    info: dict[str, Any] | None = None
+
+
+class RolloutTrajectory(BaseModel):
+    env_id: str
+    policy_id: str
+    steps: list[RolloutStep]
+    final: dict[str, Any] | None = None
+    length: int
+    decision_samples: list[dict[str, Any]] | None = None
+
+
+def compute_stepwise_reward(
+    prev_achievements: dict[str, bool],
+    new_achievements: dict[str, bool],
+    decision_index: int,
+    actions_summary: list[dict[str, Any]],
+    indicator_lambda: float,
+) -> tuple[dict[str, Any], dict[str, Any], dict[str, float]]:
+    """Compute stepwise reward metadata given achievement states before/after a decision."""
+
+    prev_map = prev_achievements or {}
+    next_map = new_achievements or {}
+
+    unlocked = [name for name, value in next_map.items() if value and not prev_map.get(name, False)]
+    indicator = 1 if unlocked else 0
+    reward_value = float(indicator_lambda) * indicator
+
+    stepwise_info = {
+        "decision_index": decision_index,
+        "indicator": indicator,
+        "new_achievements": unlocked,
+        "reward": reward_value,
+    }
+    decision_sample = {
+        "decision_index": decision_index,
+        "indicator": indicator,
+        "r_i": reward_value,
+        "actions": actions_summary,
+    }
+    stats = {
+        "indicator": float(indicator),
+        "reward": reward_value,
+        "new_achievements_count": float(len(unlocked)),
+    }
+    return stepwise_info, decision_sample, stats
+
+
+class RolloutMetrics(BaseModel):
+    episode_returns: list[float]
+    mean_return: float
+    num_steps: int
+    num_episodes: int = 0
+
+
+class RolloutResponse(BaseModel):
+    run_id: str
+    trajectories: list[RolloutTrajectory]
+    branches: dict[str, list[str]] = {}
+    metrics: RolloutMetrics
+    aborted: bool = False
+    ops_executed: int = 0
+    trace: dict[str, Any] | None = None
+
+
+class RolloutTracingContext:
+    """Helper managing tracing_v3 recording and optional SFT dumps for a rollout."""
+
+    def __init__(
+        self,
+        tracer: SessionTracer | None,
+        request: RolloutRequest,
+        fastapi_request: Request,
+    ) -> None:
+        self.tracer = tracer
+        self.enabled = tracer is not None
+        self.request = request
+        self.fastapi_request = fastapi_request
+        self.run_id = request.run_id
+        self.current_step_id: str | None = None
+        self.current_turn: int | None = None
+        self.lm_calls_summary: list[dict[str, Any]] = []
+        self.decision_rewards: list[dict[str, Any]] = []
+        self.sft_records: list[dict[str, Any]] = []
+        self.latest_system_messages: list[str] = []
+        self.latest_user_messages: list[str] = []
+        self.latest_system_prompt_content: list[Any] = []
+        self.latest_user_prompt_content: list[Any] = []
+        self.trace_format = (
+            getattr(request.record, "trace_format", "compact") or "compact"
+        ).lower()
+        self.return_trace = bool(getattr(request.record, "return_trace", False))
+        self.sft_output_dir = getattr(fastapi_request.app.state, "sft_output_dir", None)
+        print(f"[TRACE_DEBUG] RolloutTracingContext init: trace_format={self.trace_format} return_trace={self.return_trace}", flush=True)
+        self.session_trace = None
+        self.metadata_updates: dict[str, Any] = {}
+        self.policy_name = request.policy.policy_name or ""
+        self.env_name = request.env.env_name or ""
+        self.metadata_base: dict[str, Any] = {
+            "run_id": self.run_id,
+            "policy_name": self.policy_name,
+            "policy_id": request.policy.policy_id,
+            "env_name": self.env_name,
+            "env_id": request.env.env_id,
+            "seed": request.env.seed,
+            "training_session_id": request.training_session_id,
+            "synth_base_url": request.synth_base_url,
+        }
+
+        # Expose context for downstream calls inside this request lifecycle
+        fastapi_request.state.rollout_tracing = self
+        fastapi_request.state.rollout_run_id = self.run_id
+
+    async def start_session(self) -> None:
+        if not self.enabled or self.tracer is None:
+            print("[TRACE_DEBUG] start_session skipped: tracer disabled", flush=True)
+            return
+        try:
+            await self.tracer.initialize()
+            print("[TRACE_DEBUG] tracer initialized", flush=True)
+        except Exception as exc:
+            logger.debug("TRACING_INIT_FAIL: %s", exc)
+            # Hard fail: tracing requested but cannot initialize
+            raise
+        try:
+            await self.tracer.start_session(
+                session_id=self.run_id, metadata=dict(self.metadata_base)
+            )
+            print(f"[TRACE_DEBUG] start_session succeeded for run_id={self.run_id}", flush=True)
+        except Exception as exc:
+            logger.warning("TRACING_START_FAIL: %s", exc)
+            # Hard fail: tracing requested but cannot start session
+            raise
+
+    async def start_decision(self, turn_number: int) -> None:
+        self.current_turn = turn_number
+        self.current_step_id = f"decision_{turn_number}"
+        if not self.enabled or self.tracer is None:
+            return
+        try:
+            await self.tracer.start_timestep(step_id=self.current_step_id, turn_number=turn_number)
+        except Exception as exc:
+            logger.debug("TRACING_STEP_START_FAIL: %s", exc)
+
+    async def end_decision(self) -> None:
+        if not self.enabled or self.tracer is None:
+            return
+        try:
+            await self.tracer.end_timestep(step_id=self.current_step_id)
+        except Exception as exc:
+            logger.debug("TRACING_STEP_END_FAIL: %s", exc)
+        finally:
+            self.current_step_id = None
+
+    def _message_metadata(self) -> dict[str, Any]:
+        return {
+            "turn": self.current_turn,
+            "step_id": self.current_step_id,
+        }
+
+    async def record_policy_prompts(
+        self,
+        system_messages: list[Any],
+        user_messages: list[Any],
+    ) -> None:
+        self.latest_system_messages = [self._prompt_text(entry) for entry in system_messages]
+        self.latest_user_messages = [self._prompt_text(entry) for entry in user_messages]
+        self.latest_system_prompt_content = [
+            self._prompt_content(entry, role="system") for entry in system_messages
+        ]
+        self.latest_user_prompt_content = [
+            self._prompt_content(entry, role="user") for entry in user_messages
+        ]
+        if not self.enabled or self.tracer is None:
+            return
+        for entry in system_messages:
+            try:
+                await self.tracer.record_message(
+                    content=self._prompt_payload(entry, role="system"),
+                    message_type="policy_system_prompt",
+                    metadata=self._message_metadata(),
+                )
+            except Exception as exc:
+                logger.debug("TRACING_SYSTEM_MSG_FAIL: %s", exc)
+        for entry in user_messages:
+            try:
+                await self.tracer.record_message(
+                    content=self._prompt_payload(entry, role="user"),
+                    message_type="policy_user_prompt",
+                    metadata=self._message_metadata(),
+                )
+            except Exception as exc:
+                logger.debug("TRACING_USER_MSG_FAIL: %s", exc)
+        if self.tracer and self.tracer._current_trace:
+            msg_count = len(self.tracer._current_trace.markov_blanket_message_history)
+            print(f"[TRACE_DEBUG] After record_policy_prompts: {msg_count} messages", flush=True)
+
+    def _content_to_text(self, content: Any) -> str:
+        if isinstance(content, str):
+            return content
+        if isinstance(content, list):
+            parts: list[str] = []
+            for seg in content:
+                if isinstance(seg, dict):
+                    text_val = seg.get("text") or seg.get("content")
+                    if isinstance(text_val, str):
+                        parts.append(text_val)
+            return "".join(parts)
+        if content is None:
+            return ""
+        return str(content)
+
+    def _prompt_text(self, entry: Any) -> str:
+        if isinstance(entry, dict):
+            text = entry.get("text")
+            if isinstance(text, str):
+                return text
+            content = entry.get("content")
+            return self._content_to_text(content)
+        return self._content_to_text(entry)
+
+    def _prompt_payload(self, entry: Any, *, role: str) -> dict[str, Any]:
+        if isinstance(entry, dict):
+            payload = dict(entry)
+            payload.setdefault("role", role)
+            return payload
+        return {
+            "role": role,
+            "text": self._prompt_text(entry),
+            "content": entry,
+        }
+
+    def _prompt_content(self, entry: Any, *, role: str) -> Any:
+        payload = self._prompt_payload(entry, role=role)
+        return payload.get("content", payload.get("text"))
+
+    def _content_has_image(self, content: Any) -> bool:
+        if isinstance(content, list):
+            return any(
+                isinstance(seg, dict)
+                and seg.get("type") in {"image", "image_url"}
+                for seg in content
+            )
+        if isinstance(content, dict):
+            if content.get("type") in {"image", "image_url"}:
+                return True
+            inner = content.get("content")
+            if isinstance(inner, list):
+                return any(
+                    isinstance(seg, dict)
+                    and seg.get("type") in {"image", "image_url"}
+                    for seg in inner
+                )
+        return False
+
+    def _safe_json(self, payload: Any, limit: int = 4000) -> str:
+        try:
+            text = json.dumps(payload, ensure_ascii=False)
+        except Exception:
+            text = str(payload)
+        if len(text) > limit:
+            return text[:limit] + "…"
+        return text
+
+    async def record_tool_invocation(self, tool_calls: list[dict[str, Any]] | None) -> None:
+        if tool_calls is None:
+            return
+        if self.enabled and self.tracer is not None:
+            try:
+                await self.tracer.record_message(
+                    content=self._safe_json(tool_calls),
+                    message_type="policy_tool_call",
+                    metadata=self._message_metadata(),
+                )
+                if self.tracer._current_trace:
+                    print(
+                        f"[TRACE_DEBUG] After tool invocation: messages={len(self.tracer._current_trace.markov_blanket_message_history)}",
+                        flush=True,
+                    )
+            except Exception as exc:
+                logger.debug("TRACING_TOOL_MSG_FAIL: %s", exc)
+
+    async def _record_event(self, event: Any) -> int | None:
+        if not self.enabled or self.tracer is None:
+            return None
+        try:
+            return await self.tracer.record_event(event)
+        except Exception as exc:
+            logger.debug("TRACING_EVENT_FAIL: %s", exc)
+            return None
+
+    async def record_llm_call(
+        self,
+        *,
+        inference_request: dict[str, Any],
+        inference_response: dict[str, Any],
+        tool_calls: list[dict[str, Any]] | None,
+        provider: str,
+        model_name: str,
+        started_at: datetime,
+        completed_at: datetime,
+        latency_ms: int | None,
+    ) -> None:
+        usage = inference_response.get("usage") or {}
+        input_tokens = usage.get("input_tokens") or usage.get("prompt_tokens")
+        output_tokens = usage.get("output_tokens") or usage.get("completion_tokens")
+        total_tokens = usage.get("total_tokens")
+        cost_usd = usage.get("cost_usd") or usage.get("cost") or usage.get("total_cost")
+
+        assistant_message = None
+        choices = inference_response.get("choices") or []
+        if choices:
+            assistant_message = choices[0].get("message") or {}
+        assistant_content = (
+            assistant_message.get("content") if isinstance(assistant_message, dict) else None
+        )
+
+        raw_response = self._content_to_text(assistant_content)
+        if not raw_response:
+            raw_response = self._safe_json(inference_response, limit=2000)
+
+        base_response = BaseLMResponse(
+            raw_response=raw_response,
+            tool_calls=assistant_message.get("tool_calls")
+            if isinstance(assistant_message, dict)
+            else None,
+            usage=usage or None,
+            api_type="chat_completions",
+        )
+
+        request_messages = inference_request.get("messages") or []
+        try:
+            temperature = float(inference_request.get("temperature"))
+        except Exception:
+            temperature = 0.0
+
+        call_record = create_llm_call_record_from_response(
+            response=base_response,
+            model_name=model_name,
+            provider=provider,
+            messages=request_messages,
+            temperature=temperature,
+            request_params=inference_request,
+            tools=inference_request.get("tools"),
+            started_at=started_at,
+            completed_at=completed_at,
+            latency_ms=latency_ms,
+        )
+
+        event_metadata = {
+            "policy_id": self.request.policy.policy_id,
+            "turn": self.current_turn,
+            "run_id": self.run_id,
+        }
+
+        event = LMCAISEvent(
+            system_instance_id=f"policy:{self.policy_name or 'unknown'}",
+            time_record=TimeRecord(event_time=completed_at.timestamp()),
+            model_name=model_name,
+            provider=provider,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+            cost_usd=cost_usd,
+            latency_ms=latency_ms,
+            call_records=[call_record],
+            metadata=event_metadata,
+        )
+
+        await self._record_event(event)
+
+        self.lm_calls_summary.append(
+            {
+                "turn": self.current_turn,
+                "model": model_name,
+                "provider": provider,
+                "total_tokens": total_tokens,
+                "input_tokens": input_tokens,
+                "output_tokens": output_tokens,
+                "latency_ms": latency_ms,
+                "tool_calls": len(tool_calls or []),
+            }
+        )
+
+        if self.sft_output_dir is not None:
+            assistant_structured = assistant_content if assistant_content is not None else ""
+            assistant_text = self._content_to_text(assistant_content)
+            dialogue_structured: list[dict[str, Any]] = []
+            for content in self.latest_system_prompt_content:
+                if content is None:
+                    continue
+                dialogue_structured.append({"role": "system", "content": content})
+            for content in self.latest_user_prompt_content:
+                if content is None:
+                    continue
+                dialogue_structured.append({"role": "user", "content": content})
+            dialogue_text = (
+                [{"role": "system", "content": s} for s in self.latest_system_messages]
+                + [{"role": "user", "content": u} for u in self.latest_user_messages]
+            )
+            user_has_image = any(
+                self._content_has_image(content) for content in self.latest_user_prompt_content
+            )
+            assistant_has_image = self._content_has_image(assistant_structured)
+            record = {
+                "run_id": self.run_id,
+                "turn": self.current_turn,
+                "model": model_name,
+                "provider": provider,
+                "dialogue": dialogue_structured,
+                "dialogue_text": dialogue_text,
+                "assistant": {
+                    "content": assistant_structured,
+                    "content_text": assistant_text,
+                    "tool_calls": assistant_message.get("tool_calls")
+                    if isinstance(assistant_message, dict)
+                    else [],
+                    "has_image": assistant_has_image,
+                },
+                "metadata": {
+                    "user_has_image": user_has_image,
+                    "assistant_has_image": assistant_has_image,
+                    "has_image": user_has_image or assistant_has_image,
+                },
+                "timestamp": datetime.utcnow().isoformat(),
+            }
+            self.sft_records.append(record)
+
+    async def record_environment_event(
+        self,
+        *,
+        env_handle: Any,
+        prev_obs: dict[str, Any] | None,
+        env_response: Any,
+        next_obs: dict[str, Any] | None,
+        metadata: dict[str, Any] | None = None,
+    ) -> int | None:
+        if not self.enabled or self.tracer is None:
+            return None
+
+        try:
+            prev_summary = (
+                _summarize_observation_for_storage(env_handle, prev_obs or {})
+                if prev_obs is not None
+                else None
+            )
+        except Exception:
+            prev_summary = None
+        try:
+            next_summary = (
+                _summarize_observation_for_storage(env_handle, next_obs or {})
+                if next_obs is not None
+                else None
+            )
+        except Exception:
+            next_summary = None
+
+        reward_val = getattr(env_response, "reward", None)
+        try:
+            reward_float = float(reward_val) if reward_val is not None else 0.0
+        except Exception:
+            reward_float = 0.0
+
+        event = EnvironmentEvent(
+            system_instance_id=f"environment:{self.env_name or 'unknown'}",
+            time_record=TimeRecord(event_time=datetime.utcnow().timestamp()),
+            reward=reward_float,
+            terminated=bool(getattr(env_response, "done", False)),
+            truncated=bool(getattr(env_response, "truncated", False)),
+            system_state_before=prev_summary,
+            system_state_after=next_summary,
+            metadata={
+                "turn": self.current_turn,
+                "run_id": self.run_id,
+                **(metadata or {}),
+            },
+        )
+
+        return await self._record_event(event)
+
+    async def record_decision_reward(
+        self,
+        *,
+        event_id: int | None,
+        decision_meta: dict[str, Any] | None,
+    ) -> None:
+        decision_meta = decision_meta or {}
+        ach_delta = int(decision_meta.get("ach_delta", 0))
+        unique_delta = int(decision_meta.get("unique_delta", 0))
+        all_ach = list(decision_meta.get("all") or [])
+        unique_ach = list(decision_meta.get("unique") or [])
+
+        self.decision_rewards.append(
+            {
+                "turn": self.current_turn,
+                "ach_delta": ach_delta,
+                "unique_delta": unique_delta,
+                "achievements": all_ach,
+                "unique_achievements": unique_ach,
+            }
+        )
+
+        if not self.enabled or self.tracer is None or event_id is None:
+            return
+        try:
+            await self.tracer.record_event_reward(
+                event_id=event_id,
+                turn_number=self.current_turn,
+                reward_value=float(ach_delta),
+                reward_type="achievement_delta",
+                annotation={"achievements": all_ach},
+                source="environment",
+            )
+            if unique_delta:
+                await self.tracer.record_event_reward(
+                    event_id=event_id,
+                    turn_number=self.current_turn,
+                    reward_value=float(unique_delta),
+                    reward_type="unique_achievement_delta",
+                    annotation={"achievements": unique_ach},
+                    source="environment",
+                )
+        except Exception as exc:
+            logger.debug("TRACING_REWARD_FAIL: %s", exc)
+
+    def update_metadata(self, **kwargs: Any) -> None:
+        self.metadata_updates.update({k: v for k, v in kwargs.items() if v is not None})
+
+    async def finalize(
+        self,
+        *,
+        total_reward: float,
+        achievement_state: dict[str, bool] | None,
+        total_steps: int,
+    ) -> Any:
+        final_achievements = [key for key, val in (achievement_state or {}).items() if val]
+        self.metadata_updates.setdefault("final_achievements", final_achievements)
+        if self.enabled and self.tracer is not None:
+            try:
+                await self.tracer.record_outcome_reward(
+                    total_reward=int(total_reward),
+                    achievements_count=len(final_achievements),
+                    total_steps=int(total_steps),
+                    reward_metadata=dict(self.metadata_updates),
+                )
+            except Exception as exc:
+                logger.debug("TRACING_OUTCOME_FAIL: %s", exc)
+            try:
+                if self.tracer._current_trace:
+                    msg_count = len(self.tracer._current_trace.markov_blanket_message_history)
+                    print(f"[TRACE_DEBUG] Before end_session: {msg_count} messages in trace", flush=True)
+                self.session_trace = await self.tracer.end_session()
+                if self.session_trace is not None:
+                    self.session_trace.metadata.update(self.metadata_updates)
+                    print(
+                        f"[TRACE_DEBUG] Session ended successfully, session_id={self.session_trace.session_id}",
+                        flush=True,
+                    )
+                    print(
+                        f"[TRACE_DEBUG] session_trace.metadata keys: {list(self.session_trace.metadata.keys())}",
+                        flush=True,
+                    )
+            except Exception as exc:
+                logger.debug("TRACING_END_SESSION_FAIL: %s", exc)
+                self.session_trace = None
+                print(f"[TRACE_DEBUG] end_session failed for run_id={self.run_id}: {exc}", flush=True)
+            with contextlib.suppress(Exception):
+                await self.tracer.close()
+
+        if self.sft_records and self.sft_output_dir:
+            self.write_sft_records()
+
+        # Clear context from request state to avoid leaks
+        self.fastapi_request.state.rollout_tracing = None
+
+        return self.session_trace
+
+    def write_sft_records(self) -> None:
+        if not self.sft_output_dir or not self.sft_records:
+            return
+        try:
+            path = unique_sft_path(self.sft_output_dir, run_id=self.run_id)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            with path.open("w", encoding="utf-8") as fh:
+                for record in self.sft_records:
+                    json.dump(record, fh, ensure_ascii=False)
+                    fh.write("\n")
+            logger.info(f"SFT_WRITTEN: {path}")
+        except Exception as exc:
+            logger.warning(f"SFT_WRITE_FAIL: {exc}")
+        finally:
+            self.sft_records.clear()
+
+    def build_trace_payload(self, session_trace: Any) -> dict[str, Any] | None:
+        if not self.return_trace or session_trace is None:
+            return None
+        if self.trace_format in ("full", "structured"):
+            payload = session_trace.to_dict()
+            payload.setdefault("metadata", {}).update(self.metadata_updates)
+            print(
+                f"[TRACE_DEBUG] build_trace_payload returning structured trace with messages={len(payload.get('markov_blanket_message_history') or [])}",
+                flush=True,
+            )
+            return payload
+        metadata = dict(session_trace.metadata)
+        metadata.update(self.metadata_updates)
+        return {
+            "session_id": session_trace.session_id,
+            "created_at": session_trace.created_at.isoformat(),
+            "metadata": metadata,
+            "events_count": len(session_trace.event_history),
+            "messages_count": len(session_trace.markov_blanket_message_history),
+            "lm_calls": self.lm_calls_summary,
+            "decision_rewards": self.decision_rewards,
+        }
+
+
+def _summarize_observation_for_storage(
+    env_handle: Any, observation: dict[str, Any]
+) -> dict[str, Any]:
+    """Return a compact dict for trajectory storage instead of the raw observation.
+
+    - For Crafter, use the same summary used for the policy user prompt
+    - For others, keep a minimal subset or plain text preview
+    """
+    # Try Crafter-specific formatter
+    crafter_wrapper = None
+    with contextlib.suppress(Exception):
+        from .envs.crafter.environment import (
+            CrafterEnvironmentWrapper as _CrafterWrapper,  # type: ignore
+        )
+
+        crafter_wrapper = _CrafterWrapper  # type: ignore[assignment]
+
+    if crafter_wrapper is not None and isinstance(
+        getattr(env_handle, "env", None), crafter_wrapper
+    ):
+        with contextlib.suppress(Exception):
+            from .envs.crafter.shared import format_observation as _fmt  # type: ignore
+
+            text = _fmt(observation or {})
+            return {"text": text}
+
+    # Generic fallback: extract a few small fields if present; avoid huge arrays
+    with contextlib.suppress(Exception):
+        inv = observation.get("inventory") if isinstance(observation, dict) else None
+        ach = observation.get("achievements_status") if isinstance(observation, dict) else None
+        pos = observation.get("player_position") if isinstance(observation, dict) else None
+        health = None
+        if isinstance(inv, dict):
+            health = inv.get("health")
+        summary = {
+            "position": pos,
+            "health": health,
+            "inventory_keys": sorted(k for k, v in (inv or {}).items() if v)[:10]
+            if isinstance(inv, dict)
+            else None,
+            "achievements_unlocked": sorted(k for k, v in (ach or {}).items() if v)[:10]
+            if isinstance(ach, dict)
+            else None,
+        }
+        return {"text": json.dumps(summary, ensure_ascii=False)}
+
+    # Last resort: plain string preview
+    try:
+        return {"text": str(observation)[:10000]}
+    except Exception:
+        return {"text": ""}
+
+
+class RunAbortRequest(BaseModel):
+    run_id: str
+
+
+class RunAbortResponse(BaseModel):
+    ok: bool
+    run_id: str
+
+
+class RunStatusResponse(BaseModel):
+    run_id: str
+    status: str
+    started_at: datetime
+    finished_at: datetime | None = None
+
+
+@router.post("/rollout", response_model=RolloutResponse)
+async def execute_rollout(
+    request: RolloutRequest,
+    req: Request,
+) -> RolloutResponse:
+    """Execute a rollout with coordinated environment and policy steps."""
+    # Emit rollout identifier early for correlation
+    with contextlib.suppress(Exception):
+        _rid = getattr(request, "run_id", None)
+        _pol = getattr(request.policy, "policy_name", None) or getattr(request.policy, "policy_id", None)
+        _env = getattr(request.env, "env_name", None) or getattr(request.env, "env_id", None)
+        logger.info("ROLLOUT_BEGIN: run_id=%s policy=%s env=%s", _rid, _pol, _env)
+        print(f"[rollout] begin run_id={_rid} policy={_pol} env={_env}", flush=True)
+    # Enforce per-episode step cap via env-specific parameters; default to 20 if omitted
+    try:
+        _env_params = {}
+        if isinstance(request.env, RolloutEnvSpec) and isinstance(request.env.config, dict):
+            _env_params = dict(request.env.config.get("env_params") or {})
+        max_steps_per_episode = int(_env_params.get("max_steps_per_episode") or 20)
+        assert max_steps_per_episode > 0, "max_steps_per_episode must be a positive integer"
+    except Exception as _mse:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail={
+                "error": "invalid_env_params",
+                "message": f"Invalid or missing env_params.max_steps_per_episode: {_mse}",
+            },
+        ) from _mse
+    # Truncate incoming ops to the enforced cap (each step is [agent, env])
+    ops_seq: list[str] = list(request.ops or [])
+    allowed_ops = max(0, int(max_steps_per_episode) * 2)
+    if len(ops_seq) > allowed_ops:
+        with contextlib.suppress(Exception):
+            logger.info(
+                "ROLL_OUT: truncating ops to cap: requested_ops=%s allowed_ops=%s",
+                str(len(ops_seq)),
+                str(allowed_ops),
+            )
+        ops_seq = ops_seq[:allowed_ops]
+    # Simple API key auth for inbound rollout
+    header_key = req.headers.get("x-api-key")
+    env_key = os.getenv("ENVIRONMENT_API_KEY")
+    dev_key = os.getenv("DEV_ENVIRONMENT_API_KEY")
+    # Accept either ENVIRONMENT_API_KEY or DEV_ENVIRONMENT_API_KEY
+    expected_keys = [k for k in (env_key, dev_key) if k]
+    if not expected_keys:
+        missing = []
+        if not env_key:
+            missing.append("ENVIRONMENT_API_KEY")
+        if not dev_key:
+            missing.append("DEV_ENVIRONMENT_API_KEY")
+        msg = f"Auth not configured: missing {', '.join(missing)} in task service environment"
+        logger.error(msg)
+        raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail=msg)
+    if not header_key:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid or missing API key: X-API-Key header not provided",
+        )
+    if header_key not in expected_keys:
+        # Do not leak secrets; include short prefix for diagnostics
+        exp_src = env_key if env_key else (dev_key or "")
+        exp_prefix = (exp_src[:7] + "…") if len(exp_src) >= 7 else "set"
+        got_prefix = (header_key[:7] + "…") if len(header_key) >= 7 else "set"
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=f"Invalid API key: header does not match expected (got={got_prefix}, expected_prefix={exp_prefix})",
+        )
+
+    # Log contextual fields for traceability
+    if request.training_session_id:
+        logger.info(f"ROLL_OUT: training_session_id={request.training_session_id}")
+    if request.synth_base_url:
+        logger.info(f"ROLL_OUT: synth_base_url={request.synth_base_url}")
+
+    # Log masked OpenAI API key presence for diagnostics
+    with contextlib.suppress(Exception):
+        _oa = os.getenv("OPENAI_API_KEY")
+        if _oa:
+            _pref = (_oa[:6] + "…") if len(_oa) >= 6 else "set"
+            logger.info(f"ROLL_OUT: OPENAI_API_KEY present (prefix={_pref})")
+        else:
+            logger.warning("ROLL_OUT: OPENAI_API_KEY missing")
+
+    # Make synth_base_url available for outbound calls in this app
+    with contextlib.suppress(Exception):
+        task_app = req.app.state.task_app
+        if request.synth_base_url:
+            task_app.synth_base_url = request.synth_base_url
+
+    tracer_factory = getattr(req.app.state, "session_tracer_factory", None)
+    tracer_instance: SessionTracer | None = None
+    if callable(tracer_factory):
+        try:
+            inst = tracer_factory()
+            tracer_instance = inst if isinstance(inst, SessionTracer) else None
+        except Exception as exc:
+            logger.debug(f"TRACER_FACTORY_FAIL: {exc}")
+    tracing_context = RolloutTracingContext(tracer_instance, request, req)
+    await tracing_context.start_session()
+
+    # Register run
+    registry.register_run(request.run_id)
+
+    # Track resources created during this rollout so we can guarantee cleanup
+    created_env_id: str | None = None
+    created_policy_id: str | None = None
+    env_seed_used: int | None = None
+    trajectory_steps: list[RolloutStep] = []
+    decision_samples: list[dict[str, Any]] = []
+    pending_tool_calls: Any = None
+    current_obs: Any = {}
+    total_reward: float = 0.0
+    ops_executed = 0
+    last_agent_response_ts: float | None = None
+    last_policy_meta: dict[str, Any] | None = None
+    last_env_step_ms: float | None = None
+    last_env_step_completed_ts: float | None = None
+    decision_open = False
+    finalized = False
+    prev_achievements: dict[str, bool] = {}
+    session_trace = None
+    step_rewards_active = False
+
+    try:
+        # Initialize deterministic seed early for the entire rollout
+        seed_value: int | None = None
+        try:
+            if request.env and request.env.seed is not None:
+                seed_value = int(request.env.seed)
+            else:
+                # Derive a stable seed from run_id
+                import hashlib as _hashlib  # local import to avoid global deps
+
+                _digest = _hashlib.sha256(request.run_id.encode("utf-8")).hexdigest()
+                # Use lower 32 bits to fit common RNG ranges
+                seed_value = int(_digest[:8], 16)
+        except Exception:
+            # Fallback to time-based seed if anything goes wrong
+            try:
+                seed_value = int((_time.time_ns() // 1_000_000) % (2**31 - 1))
+            except Exception:
+                seed_value = 42
+
+        _seed_info = _set_global_seed(int(seed_value))
+        with contextlib.suppress(Exception):
+            logger.info(
+                "ROLL_OUT: RNG seeded seed=%s libs=%s",
+                str(_seed_info.get("seed")),
+                ",".join(_seed_info.get("libs", [])),
+            )
+        # Resolve or create environment
+        if request.env.env_id:
+            env_handle = registry.get_env(request.env.env_id)
+            if not env_handle:
+                raise HTTPException(
+                    status_code=404,
+                    detail=f"Environment {request.env.env_id} not found",
+                )
+            env_id = request.env.env_id
+        else:
+            # Create new environment
+            from .environment_routes import EnvCreateRequest, create_environment
+
+            if not request.env.env_name:
+                raise ValueError("FATAL: env_name is required - NO FALLBACKS!")
+
+            # Propagate training_session_id via env config for downstream usage
+            _env_config = dict(request.env.config or {})
+            if request.training_session_id is not None:
+                _env_config.setdefault("training_session_id", request.training_session_id)
+            env_response = await create_environment(
+                EnvCreateRequest(
+                    env_name=request.env.env_name,
+                    config=_env_config,
+                    seed=request.env.seed,
+                    rl_run_id=request.run_id,
+                )
+            )
+            env_id = env_response.env_id
+            env_handle = registry.get_env(env_id)
+            created_env_id = env_id
+
+        tracing_context.update_metadata(env_id=env_id)
+
+        # Resolve or create policy
+        if request.policy.policy_id:
+            policy_handle = registry.get_policy(request.policy.policy_id)
+            if not policy_handle:
+                raise HTTPException(
+                    status_code=404,
+                    detail=f"Policy {request.policy.policy_id} not found",
+                )
+            policy_id = request.policy.policy_id
+        else:
+            # Create new policy
+            from .policy_routes import PolicyCreateRequest, create_policy
+
+            if not request.policy.policy_name:
+                raise ValueError("FATAL: policy_name is required - NO FALLBACKS!")
+
+            # Propagate training_session_id and synth_base_url via policy config
+            _policy_config = dict(request.policy.config or {})
+            if request.training_session_id is not None:
+                _policy_config.setdefault("training_session_id", request.training_session_id)
+            if request.synth_base_url is not None:
+                _policy_config.setdefault("synth_base_url", request.synth_base_url)
+            policy_response = await create_policy(
+                PolicyCreateRequest(
+                    policy_name=request.policy.policy_name,
+                    config=_policy_config,
+                    rl_run_id=request.run_id,
+                    bound_env_id=env_id,
+                ),
+                req,
+            )
+            policy_id = policy_response.policy_id
+            policy_handle = registry.get_policy(policy_id)
+            created_policy_id = policy_id
+
+        tracing_context.update_metadata(policy_id=policy_id)
+
+        # Bind policy to environment if not already bound
+        if policy_handle and not policy_handle.bound_env_id:
+            policy_handle.bound_env_id = env_id
+
+        # Record seed bound to environment for end-of-rollout verification/logging
+        try:
+            env_seed_used = int(getattr(env_handle, "seed", 0) or 0)
+        except Exception:
+            env_seed_used = None
+        tracing_context.update_metadata(env_seed=env_seed_used)
+        # Initialize trajectory
+        trajectory_steps = []
+        pending_tool_calls = None
+        current_obs = env_handle.last_observation
+        total_reward = 0.0
+        ops_executed = 0
+        last_agent_response_ts = None
+        last_policy_meta = None
+        last_env_step_ms = None
+        last_env_step_completed_ts = None
+
+        # Stepwise reward configuration (Crafter shaping; gate on explicit enable)
+        step_rewards_cfg_raw: dict[str, Any] = {}
+        try:
+            if isinstance(request.policy.config, dict):
+                step_rewards_cfg_raw = dict(request.policy.config.get("step_rewards") or {})
+        except Exception:
+            step_rewards_cfg_raw = {}
+        if not step_rewards_cfg_raw:
+            try:
+                if isinstance(request.env.config, dict):
+                    step_rewards_cfg_raw = dict(request.env.config.get("step_rewards") or {})
+            except Exception:
+                step_rewards_cfg_raw = {}
+
+        step_rewards_enabled = bool(step_rewards_cfg_raw.get("enabled", False))
+        step_rewards_mode = str(step_rewards_cfg_raw.get("mode") or "off").lower()
+        try:
+            step_rewards_indicator_lambda = float(
+                step_rewards_cfg_raw.get("indicator_lambda") or 0.0
+            )
+        except Exception:
+            step_rewards_indicator_lambda = 0.0
+        try:
+            step_rewards_beta = float(step_rewards_cfg_raw.get("step_beta") or 0.0)
+        except Exception:
+            step_rewards_beta = 0.0
+        step_rewards_active = step_rewards_enabled and step_rewards_mode == "decision_stepwise"
+
+        def _extract_achievements(obs: Any) -> dict[str, bool]:
+            if not isinstance(obs, dict):
+                return {}
+            ach = obs.get("achievements_status")
+            if isinstance(ach, dict):
+                return {str(k): bool(v) for k, v in ach.items()}
+            return {}
+
+        def _summarize_tool_calls(tool_calls: Any) -> list[dict[str, Any]]:
+            if not tool_calls:
+                return []
+            try:
+                items = (
+                    tool_calls
+                    if isinstance(tool_calls, list)
+                    else list(tool_calls)  # tolerates tuples or pydantic lists
+                )
+            except Exception:
+                return []
+            summary: list[dict[str, Any]] = []
+            for tc in items:
+                tool_name = None
+                args: Any = {}
+                if isinstance(tc, dict):
+                    tool_name = tc.get("tool") or tc.get("tool_name") or tc.get("name")
+                    raw_args = tc.get("arguments") or tc.get("args") or {}
+                else:
+                    tool_name = getattr(tc, "tool", None) or getattr(tc, "tool_name", None)
+                    raw_args = getattr(tc, "arguments", None) or getattr(tc, "args", None) or {}
+                args = raw_args
+                if isinstance(raw_args, str):
+                    try:
+                        args = json.loads(raw_args)
+                    except Exception:
+                        args = raw_args
+                summary.append({"tool": tool_name, "args": args})
+            return summary
+
+        decision_samples: list[dict[str, Any]] = []
+        decision_index = 0
+        decision_open = False
+        session_trace = None
+        finalized = False
+        prev_achievements = _extract_achievements(current_obs)
+        # Track episode-level achievements that have been seen as true at any point so far
+        episode_seen_achievements: set[str] = {
+            k for k, v in (prev_achievements or {}).items() if bool(v)
+        }
+        stepwise_indicator_sum = 0.0
+        stepwise_reward_sum = 0.0
+        stepwise_new_achievements_total = 0
+        final_achievement_count = sum(1 for v in prev_achievements.values() if v)
+
+        # Execute ops sequence (capped by env_params.max_steps_per_episode)
+        for op_idx, op in enumerate(ops_seq):
+            # Check for abort
+            if registry.is_run_aborted(request.run_id):
+                logger.info(f"Run {request.run_id} aborted at op {op_idx}")
+                break
+
+            # Check safety limits
+            if ops_executed >= request.safety.max_ops:
+                logger.warning(f"Reached max_ops limit ({request.safety.max_ops})")
+                break
+
+            if op == "agent":
+                # Policy step
+                from .policy_routes import PolicyStepRequest, step_policy
+
+                if not decision_open:
+                    await tracing_context.start_decision(decision_index)
+                    decision_open = True
+
+                agent_request_start = _time.perf_counter()
+                if last_agent_response_ts is not None and last_policy_meta is not None:
+                    with contextlib.suppress(Exception):
+                        timing_prev = last_policy_meta.setdefault("timing", {})
+                        decision_ms = max(
+                            0.0,
+                            (agent_request_start - float(last_agent_response_ts)) * 1000.0,
+                        )
+                        # Update timing on prior policy meta (kept by previous env step)
+                        timing_prev["decision_ms"] = decision_ms
+                        if last_env_step_ms is not None:
+                            timing_prev["env_step_ms"] = float(last_env_step_ms)
+                            timing_prev["overhead_ms"] = max(
+                                0.0, decision_ms - float(last_env_step_ms)
+                            )
+                        else:
+                            timing_prev.setdefault("overhead_ms", 0.0)
+                        timing_prev["decision_ready_s"] = agent_request_start
+                        # Also backfill the last appended trajectory step so the trainer
+                        # can always see decision_ms without relying on shared dict refs.
+                        if trajectory_steps:
+                            with contextlib.suppress(Exception):
+                                _last = trajectory_steps[-1]
+                                _info = dict(_last.info or {})
+                                _meta = dict(_info.get("meta") or {})
+                                _timing = dict(_meta.get("timing") or {})
+                                _timing["decision_ms"] = decision_ms
+                                if last_env_step_ms is not None:
+                                    _timing.setdefault("env_step_ms", float(last_env_step_ms))
+                                    _timing.setdefault(
+                                        "overhead_ms",
+                                        max(0.0, decision_ms - float(last_env_step_ms)),
+                                    )
+                                else:
+                                    _timing.setdefault("overhead_ms", 0.0)
+                                _meta["timing"] = _timing
+                                _info["meta"] = _meta
+                                _last.info = _info
+                last_env_step_ms = None
+                last_env_step_completed_ts = None
+
+                # Build metadata for policy (carry previous tool_calls and env result)
+                metadata = {}
+                if pending_tool_calls:
+                    metadata["prev_tool_calls"] = pending_tool_calls
+                if len(trajectory_steps) > 0:
+                    last_step = trajectory_steps[-1]
+                    # Prefer the last executed tool calls to seed history
+                    if last_step.tool_calls:
+                        metadata["prev_tool_calls"] = last_step.tool_calls
+                    # Provide a compact env result snapshot
+                    metadata["prev_env_result"] = {
+                        "observation": last_step.obs,
+                        "reward": last_step.reward,
+                        "done": last_step.done,
+                        "truncated": last_step.truncated,
+                        "info": last_step.info,
+                    }
+
+                # Log compact metadata summary to confirm history threading
+                with contextlib.suppress(Exception):
+                    _prev_calls = metadata.get("prev_tool_calls")
+                    _count = len(_prev_calls) if isinstance(_prev_calls, list) else 0
+                    _first_guess = None
+                    if _count > 0 and isinstance(_prev_calls[0], dict):
+                        _args = _prev_calls[0].get("arguments", None)
+                        if isinstance(_args, str):
+                            import json as _json
+                            with contextlib.suppress(Exception):
+                                _args = _json.loads(_args)
+                        if not isinstance(_args, dict):
+                            _args = {}
+                        _first_guess = _args.get("guess") or _args.get("word")
+                    logger.info(
+                        "POLICY_METADATA: prev_tool_calls=%d first_guess=%r has_prev_env_result=%s",
+                        _count,
+                        _first_guess,
+                        str("prev_env_result" in metadata),
+                    )
+
+                try:
+                    policy_response = await step_policy(
+                        PolicyStepRequest(
+                            policy_id=policy_id,
+                            observation=current_obs,
+                            metadata=metadata,
+                        ),
+                        req,
+                    )
+                except Exception as _pe:
+                    # Do not 500 the rollout; finalize with partial trajectory
+                    with contextlib.suppress(Exception):
+                        logger.warning(
+                            "POLICY_STEP_FAIL: terminating episode early run_id=%s op_idx=%s err=%s",
+                            request.run_id,
+                            str(op_idx),
+                            str(_pe),
+                        )
+
+                    # Build partial trajectory and return HTTP 200
+                    trajectory = RolloutTrajectory(
+                        env_id=env_id,
+                        policy_id=policy_id,
+                        steps=trajectory_steps,
+                        final={
+                            "observation": current_obs,
+                            "rollout_status": "partial_policy_error",
+                            "error": str(_pe),
+                            "at_op": op,
+                        },
+                        length=len(trajectory_steps),
+                        decision_samples=decision_samples if step_rewards_active else None,
+                    )
+                    metrics = RolloutMetrics(
+                        episode_returns=[total_reward],
+                        mean_return=total_reward,
+                        num_steps=len(trajectory_steps),
+                        num_episodes=1,
+                    )
+                    aborted = registry.is_run_aborted(request.run_id)
+                    if not aborted:
+                        registry.complete_run(request.run_id)
+                    if decision_open:
+                        await tracing_context.end_decision()
+                        decision_open = False
+                    if not finalized:
+                        session_trace = await tracing_context.finalize(
+                            total_reward=total_reward,
+                            achievement_state=prev_achievements,
+                            total_steps=len(trajectory_steps),
+                        )
+                        finalized = True
+                    trace_payload = tracing_context.build_trace_payload(session_trace)
+                    return RolloutResponse(
+                        run_id=request.run_id,
+                        trajectories=[trajectory],
+                        branches={},
+                        metrics=metrics,
+                        aborted=aborted,
+                        ops_executed=ops_executed,
+                        trace=trace_payload,
+                    )
+
+                agent_response_ts = _time.perf_counter()
+                if isinstance(policy_response.meta, dict):
+                    with contextlib.suppress(Exception):
+                        timing_cur = policy_response.meta.setdefault("timing", {})
+                        timing_cur["agent_request_start_s"] = agent_request_start
+                        timing_cur["agent_response_s"] = agent_response_ts
+                        if "inference_ms" in policy_response.meta:
+                            with contextlib.suppress(Exception):
+                                timing_cur.setdefault(
+                                    "inference_ms",
+                                    float(policy_response.meta["inference_ms"]),
+                                )
+                                timing_cur.setdefault(
+                                    "inference_s",
+                                    float(policy_response.meta["inference_ms"]) / 1000.0,
+                                )
+                    last_policy_meta = policy_response.meta
+                else:
+                    last_policy_meta = None
+                last_agent_response_ts = agent_response_ts
+
+                # Diagnostic: summarize policy step target and tool calls
+                try:
+                    model_name = None
+                    target_url = None
+                    if isinstance(policy_response.meta, dict):
+                        req_body = policy_response.meta.get("inference_request") or {}
+                        model_name = req_body.get("model")
+                        target_url = policy_response.meta.get("inference_url")
+                    _tc = policy_response.tool_calls or []
+                    print(
+                        {
+                            "rollout.policy_step": True,
+                            "run_id": request.run_id,
+                            "model": model_name,
+                            "inference_url": target_url,
+                            "tool_calls_count": len(_tc) if isinstance(_tc, list) else 0,
+                        },
+                        flush=True,
+                    )
+                except Exception:
+                    pass
+
+                pending_tool_calls = policy_response.tool_calls
+                # Log summarized agent tool calls
+                with contextlib.suppress(Exception):
+                    _tc = pending_tool_calls or []
+                    _summary = []
+                    for _item in (_tc if isinstance(_tc, list) else []):
+                        try:
+                            if isinstance(_item, dict):
+                                _tool = _item.get("tool")
+                                _args = _item.get("args")
+                                _keys = list(_args.keys()) if isinstance(_args, dict) else []
+                                _summary.append({"tool": _tool, "args_keys": _keys})
+                        except Exception:
+                            continue
+                    _rid = getattr(request, "run_id", None)
+                    logger.info("AGENT_TOOL_CALLS: run_id=%s count=%d summary=%s", _rid, len(_tc), _summary)
+                    print(f"[rollout] agent tool_calls run_id={_rid} count={len(_tc)} summary={_summary}", flush=True)
+                await tracing_context.record_tool_invocation(pending_tool_calls)
+                ops_executed += 1
+
+            elif op == "env":
+                if not pending_tool_calls:
+                    # Treat absence of tool calls as a soft terminal condition; yield partial trajectory
+                    with contextlib.suppress(Exception):
+                        logger.warning(
+                            "NO_TOOL_CALLS: terminating episode early run_id=%s op_idx=%s",
+                            request.run_id,
+                            str(op_idx),
+                        )
+                        print(
+                            f"[rollout] no tool_calls; terminating early run_id={request.run_id} op_idx={op_idx}",
+                            flush=True,
+                        )
+                    term_step = RolloutStep(
+                        obs=current_obs,
+                        tool_calls=[],
+                        reward=None,
+                        done=True,
+                        truncated=False,
+                        info={
+                            "terminated": True,
+                            "reason": "no_tool_calls",
+                        },
+                    )
+                    trajectory_steps.append(term_step)
+                    trajectory = RolloutTrajectory(
+                        env_id=env_id,
+                        policy_id=policy_id,
+                        steps=trajectory_steps,
+                        final={
+                            "observation": current_obs,
+                            "rollout_status": "partial_no_tool_calls",
+                            "at_op": op,
+                        },
+                        length=len(trajectory_steps),
+                        decision_samples=decision_samples if step_rewards_active else None,
+                    )
+                    metrics = RolloutMetrics(
+                        episode_returns=[total_reward],
+                        mean_return=total_reward,
+                        num_steps=len(trajectory_steps),
+                        num_episodes=1,
+                    )
+                    aborted = registry.is_run_aborted(request.run_id)
+                    if not aborted:
+                        registry.complete_run(request.run_id)
+                    if decision_open:
+                        await tracing_context.end_decision()
+                        decision_open = False
+                    if not finalized:
+                        session_trace = await tracing_context.finalize(
+                            total_reward=total_reward,
+                            achievement_state=prev_achievements,
+                            total_steps=len(trajectory_steps),
+                        )
+                        finalized = True
+                    trace_payload = tracing_context.build_trace_payload(session_trace)
+                    return RolloutResponse(
+                        run_id=request.run_id,
+                        trajectories=[trajectory],
+                        branches={},
+                        metrics=metrics,
+                        aborted=aborted,
+                        ops_executed=ops_executed,
+                        trace=trace_payload,
+                    )
+
+                # Environment step
+                from .environment_routes import EnvStepRequest, step_environment
+
+                env_step_error: Exception | None = None
+                env_response = None
+                env_step_start = _time.perf_counter()
+                try:
+                    env_response = await step_environment(
+                        EnvStepRequest(
+                            env_id=env_id,
+                            tool_calls=pending_tool_calls,
+                        )
+                    )
+                except Exception as _ee:
+                    env_step_error = _ee
+                env_step_end = _time.perf_counter()
+                env_step_duration_ms = (env_step_end - env_step_start) * 1000.0
+                last_env_step_ms = env_step_duration_ms
+                last_env_step_completed_ts = env_step_end
+                if last_policy_meta is not None:
+                    with contextlib.suppress(Exception):
+                        timing_env = last_policy_meta.setdefault("timing", {})
+                        timing_env["env_step_ms"] = env_step_duration_ms
+                        timing_env["env_step_end_s"] = env_step_end
+
+                if env_step_error is not None:
+                    # Invalid action or environment rejection — terminate episode early with partial trajectory
+                    with contextlib.suppress(Exception):
+                        logger.warning(
+                            "ENV_STEP_FAIL: terminating episode early run_id=%s op_idx=%s err=%s",
+                            request.run_id,
+                            str(op_idx),
+                            str(env_step_error),
+                        )
+
+                    term_step = RolloutStep(
+                        obs=current_obs,
+                        tool_calls=pending_tool_calls,
+                        reward=None,
+                        done=True,
+                        truncated=False,
+                        info={
+                            "terminated": True,
+                            "reason": "invalid_action",
+                            "error": str(env_step_error),
+                        },
+                    )
+                    trajectory_steps.append(term_step)
+                    # Build partial response
+                    trajectory = RolloutTrajectory(
+                        env_id=env_id,
+                        policy_id=policy_id,
+                        steps=trajectory_steps,
+                        final={
+                            "observation": current_obs,
+                            "rollout_status": "partial_invalid_action",
+                            "error": str(env_step_error),
+                            "at_op": op,
+                        },
+                        length=len(trajectory_steps),
+                        decision_samples=decision_samples if step_rewards_active else None,
+                    )
+                    metrics = RolloutMetrics(
+                        episode_returns=[total_reward],
+                        mean_return=total_reward,
+                        num_steps=len(trajectory_steps),
+                        num_episodes=1,
+                    )
+                    aborted = registry.is_run_aborted(request.run_id)
+                    if not aborted:
+                        registry.complete_run(request.run_id)
+                    if (
+                        last_policy_meta is not None
+                        and last_agent_response_ts is not None
+                        and "decision_ms" not in last_policy_meta.get("timing", {})
+                    ):
+                        with contextlib.suppress(Exception):
+                            timing_last = last_policy_meta.setdefault("timing", {})
+                            decision_ms = max(
+                                0.0,
+                                (env_step_end - float(last_agent_response_ts)) * 1000.0,
+                            )
+                            timing_last["decision_ms"] = decision_ms
+                            timing_last.setdefault(
+                                "overhead_ms", max(0.0, decision_ms - env_step_duration_ms)
+                            )
+                    if decision_open:
+                        await tracing_context.end_decision()
+                        decision_open = False
+                    if not finalized:
+                        session_trace = await tracing_context.finalize(
+                            total_reward=total_reward,
+                            achievement_state=prev_achievements,
+                            total_steps=len(trajectory_steps),
+                        )
+                        finalized = True
+                    trace_payload = tracing_context.build_trace_payload(session_trace)
+                    return RolloutResponse(
+                        run_id=request.run_id,
+                        trajectories=[trajectory],
+                        branches={},
+                        metrics=metrics,
+                        aborted=aborted,
+                        ops_executed=ops_executed,
+                        trace=trace_payload,
+                    )
+
+                # Reaching here means env step succeeded
+                assert env_response is not None
+
+                # Record step, including policy meta if present for timing/tokens observability
+                _info = env_response.info if isinstance(env_response.info, dict) else {}
+                # Attach policy meta from the immediately preceding agent step
+                with contextlib.suppress(Exception):
+                    prev_meta = {}
+                    if "policy_response" in locals() and isinstance(policy_response.meta, dict):  # type: ignore[name-defined]
+                        prev_meta = policy_response.meta
+                    if prev_meta:
+                        _info = dict(_info)
+                        _info["meta"] = prev_meta
+
+                event_metadata = {
+                    "op_index": op_idx,
+                }
+                event_id = await tracing_context.record_environment_event(
+                    env_handle=env_handle,
+                    prev_obs=current_obs,
+                    env_response=env_response,
+                    next_obs=getattr(env_response, "observation", None),
+                    metadata=event_metadata,
+                )
+
+                decision_index += 1
+                next_obs = env_response.observation
+                new_achievement_state = _extract_achievements(next_obs)
+                final_achievement_count = sum(
+                    1 for _, unlocked in new_achievement_state.items() if unlocked
+                )
+                indicator_val = 0
+                reward_stepwise = 0.0
+                decision_rewards_meta: dict[str, Any] | None = None
+                if step_rewards_active:
+                    decision_actions = _summarize_tool_calls(pending_tool_calls)
+                    stepwise_info, decision_record, stats = compute_stepwise_reward(
+                        prev_achievements or {},
+                        new_achievement_state,
+                        decision_index,
+                        decision_actions,
+                        step_rewards_indicator_lambda,
+                    )
+                    indicator_val = int(stats.get("indicator", 0.0))
+                    reward_stepwise = float(stats.get("reward", 0.0))
+                    stepwise_indicator_sum += float(stats.get("indicator", 0.0))
+                    stepwise_reward_sum += reward_stepwise
+                    stepwise_new_achievements_total += int(stats.get("new_achievements_count", 0.0))
+                    _info = {} if not isinstance(_info, dict) else dict(_info)
+                    _info["stepwise"] = stepwise_info
+                    # Compute decision-level rewards (absolute vs unique) and attach to metadata
+                    with contextlib.suppress(Exception):
+                        turned_true = set(stepwise_info.get("new_achievements") or [])
+                        seen_before = set(episode_seen_achievements)
+                        new_unique = sorted(turned_true - seen_before)
+                        ach_delta = int(len(turned_true))
+                        unique_delta = int(len(new_unique))
+                        # Prepare stable lists for logging/metadata
+                        all_list = sorted(turned_true)
+                        # Ensure nested meta exists
+                        meta_block = (
+                            _info.get("meta") if isinstance(_info.get("meta"), dict) else {}
+                        )
+                        decision_rewards = {
+                            "turn": int(decision_index),
+                            "ach_delta": ach_delta,
+                            "unique_delta": unique_delta,
+                            "all": all_list,
+                            "unique": new_unique,
+                        }
+                        decision_rewards_meta = decision_rewards
+                        meta_block["decision_rewards"] = decision_rewards
+                        _info["meta"] = meta_block
+                        # Update episode-level seen set after attributing uniqueness to this decision
+                        episode_seen_achievements.update(turned_true)
+                    decision_samples.append(decision_record)
+                prev_achievements = new_achievement_state
+
+                await tracing_context.record_decision_reward(
+                    event_id=event_id,
+                    decision_meta=decision_rewards_meta,
+                )
+
+                step = RolloutStep(
+                    obs=_summarize_observation_for_storage(env_handle, current_obs),
+                    tool_calls=pending_tool_calls,
+                    reward=env_response.reward,
+                    done=env_response.done,
+                    truncated=env_response.truncated,
+                    info=_info,
+                )
+                # Log summarized env application of tool calls and immediate reward/done
+                with contextlib.suppress(Exception):
+                    _tc = pending_tool_calls or []
+                    _summary = []
+                    for _item in (_tc if isinstance(_tc, list) else []):
+                        try:
+                            if isinstance(_item, dict):
+                                _tool = _item.get("tool")
+                                _args = _item.get("args")
+                                _keys = list(_args.keys()) if isinstance(_args, dict) else []
+                                _summary.append({"tool": _tool, "args_keys": _keys})
+                        except Exception:
+                            continue
+                    _rid = getattr(request, "run_id", None)
+                    logger.info(
+                        "ENV_APPLY: run_id=%s tool_calls=%d reward=%s done=%s summary=%s",
+                        _rid,
+                        len(_tc),
+                        str(env_response.reward),
+                        str(env_response.done),
+                        _summary,
+                    )
+                    print(
+                        f"[rollout] env apply run_id={_rid} tool_calls={len(_tc)} reward={env_response.reward} done={env_response.done} summary={_summary}",
+                        flush=True,
+                    )
+                trajectory_steps.append(step)
+
+                if env_response.reward is not None:
+                    total_reward += env_response.reward
+
+                # Update state
+                current_obs = next_obs
+                pending_tool_calls = None
+                ops_executed += 1
+
+                # Handle episode end
+                if env_response.done:
+                    if request.on_done == "reset":
+                        # Reset environment
+                        from .environment_routes import (
+                            EnvResetRequest,
+                            reset_environment,
+                        )
+
+                        reset_response = await reset_environment(EnvResetRequest(env_id=env_id))
+                        current_obs = reset_response.observation
+                    elif request.on_done == "terminate":
+                        break
+
+                if decision_open:
+                    await tracing_context.end_decision()
+                    decision_open = False
+
+            else:
+                logger.warning(f"Unknown op: {op}")
+
+        if (
+            last_policy_meta is not None
+            and last_agent_response_ts is not None
+            and "timing" in last_policy_meta
+            and isinstance(last_policy_meta["timing"], dict)
+            and "decision_ms" not in last_policy_meta["timing"]
+        ):
+            with contextlib.suppress(Exception):
+                final_now = last_env_step_completed_ts or _time.perf_counter()
+                final_decision_ms = max(0.0, (final_now - float(last_agent_response_ts)) * 1000.0)
+                timing_final = last_policy_meta.setdefault("timing", {})
+                timing_final["decision_ms"] = final_decision_ms
+                if last_env_step_ms is not None:
+                    timing_final.setdefault("env_step_ms", float(last_env_step_ms))
+                    timing_final.setdefault(
+                        "overhead_ms",
+                        max(0.0, final_decision_ms - float(last_env_step_ms)),
+                    )
+                else:
+                    timing_final.setdefault("overhead_ms", 0.0)
+
+        # Build trajectory
+        trajectory = RolloutTrajectory(
+            env_id=env_id,
+            policy_id=policy_id,
+            steps=trajectory_steps,
+            final={"observation": _summarize_observation_for_storage(env_handle, current_obs)},
+            length=len(trajectory_steps),
+            decision_samples=decision_samples if step_rewards_active else None,
+        )
+
+        # Build metrics
+        metrics = RolloutMetrics(
+            episode_returns=[total_reward],
+            mean_return=total_reward,
+            num_steps=len(trajectory_steps),
+            num_episodes=1,
+        )
+
+        # Environment-specific: Log summary if available
+        try:
+            # Check if this is a Wordle environment and use Wordle helpers (lazy import)
+            wordle_wrapper_cls = None
+            try:
+                from .envs.wordle.environment import WordleEnvironmentWrapper
+                from .envs.wordle.helpers import (
+                    get_wordle_rollout_summary,
+                    log_wordle_rollout_summary,
+                )
+
+                wordle_wrapper_cls = WordleEnvironmentWrapper
+            except Exception:
+                wordle_wrapper_cls = None  # type: ignore[assignment]
+                get_wordle_rollout_summary = None  # type: ignore
+                log_wordle_rollout_summary = None  # type: ignore
+
+            is_wordle = wordle_wrapper_cls is not None and isinstance(
+                env_handle.env,
+                wordle_wrapper_cls,  # type: ignore[arg-type]
+            )
+            if is_wordle:
+                # Convert trajectory steps to expected format
+                formatted_steps = []
+                for step in trajectory_steps:
+                    formatted_steps.append({"tool_calls": step.tool_calls or []})
+
+                if (
+                    get_wordle_rollout_summary is not None
+                    and log_wordle_rollout_summary is not None
+                ):
+                    summary = get_wordle_rollout_summary(formatted_steps, current_obs, env_handle)
+                    log_wordle_rollout_summary(request.run_id, summary)
+        except ImportError:
+            # Wordle helpers not available, skip Wordle-specific logging
+            pass
+        except Exception as e:
+            logger.warning(f"Failed to generate environment-specific summary: {e}")
+
+        # Mark run as completed
+        aborted = registry.is_run_aborted(request.run_id)
+        if not aborted:
+            registry.complete_run(request.run_id)
+        if decision_open:
+            await tracing_context.end_decision()
+            decision_open = False
+        if not finalized:
+            session_trace = await tracing_context.finalize(
+                total_reward=total_reward,
+                achievement_state=prev_achievements,
+                total_steps=len(trajectory_steps),
+            )
+            finalized = True
+        trace_payload = tracing_context.build_trace_payload(session_trace)
+
+        return RolloutResponse(
+            run_id=request.run_id,
+            trajectories=[trajectory],
+            branches={},
+            metrics=metrics,
+            aborted=aborted,
+            ops_executed=ops_executed,
+            trace=trace_payload,
+        )
+
+    except Exception as e:
+        logger.error(f"Rollout failed for run {request.run_id}: {e}")
+        registry.abort_run(request.run_id)
+        if decision_open:
+            with contextlib.suppress(Exception):
+                await tracing_context.end_decision()
+            decision_open = False
+        if not finalized:
+            session_trace = None
+            with contextlib.suppress(Exception):
+                session_trace = await tracing_context.finalize(
+                    total_reward=total_reward,
+                    achievement_state=prev_achievements,
+                    total_steps=len(trajectory_steps),
+                )
+            finalized = True
+        raise HTTPException(status_code=500, detail=str(e)) from e
+    finally:
+        # Ensure any environment created for this rollout is terminated (no reuse across rollouts)
+        try:
+            if created_env_id:
+                from .environment_routes import EnvTerminateRequest, terminate_environment
+
+                await terminate_environment(EnvTerminateRequest(env_id=created_env_id))
+                logger.info(
+                    "ROLL_OUT: terminated environment env_id=%s seed=%s",
+                    str(created_env_id),
+                    str(env_seed_used) if env_seed_used is not None else "unknown",
+                )
+                # Verify removal from registry
+                with contextlib.suppress(Exception):
+                    _post = registry.get_env(created_env_id)
+                    logger.info(
+                        "ROLL_OUT: env_killed=%s (post_lookup=%s)",
+                        str(_post is None),
+                        str(_post),
+                    )
+        except Exception as _te:
+            logger.warning(f"ROLL_OUT: failed to terminate environment {created_env_id}: {_te}")
+
+        # Best-effort policy cleanup if we created one (avoid reuse across rollouts)
+        with contextlib.suppress(Exception):
+            if created_policy_id:
+                from .policy_routes import PolicyTerminateRequest, terminate_policy
+
+                await terminate_policy(PolicyTerminateRequest(policy_id=created_policy_id))
+                logger.info("ROLL_OUT: terminated policy policy_id=%s", str(created_policy_id))
+
+        if not finalized:
+            session_trace = None
+            with contextlib.suppress(Exception):
+                session_trace = await tracing_context.finalize(
+                    total_reward=total_reward,
+                    achievement_state=prev_achievements,
+                    total_steps=len(trajectory_steps),
+                )
+            finalized = True
+
+        with contextlib.suppress(Exception):
+            _clear_seed_side_effects()
+            logger.info("ROLL_OUT: RNG seed terminated/cleared before conclusion")
+
+
+@router.post("/run/abort", response_model=RunAbortResponse)
+async def abort_run(request: RunAbortRequest) -> RunAbortResponse:
+    """Abort a running rollout."""
+    success = registry.abort_run(request.run_id)
+
+    if not success:
+        raise HTTPException(
+            status_code=404,
+            detail=f"Run {request.run_id} not found",
+        )
+
+    return RunAbortResponse(
+        ok=True,
+        run_id=request.run_id,
+    )
+
+
+@router.get("/run/status/{run_id}", response_model=RunStatusResponse)
+async def get_run_status(run_id: str) -> RunStatusResponse:
+    """Get the status of a run."""
+    run_handle = registry.get_run(run_id)
+
+    if not run_handle:
+        raise HTTPException(
+            status_code=404,
+            detail=f"Run {run_id} not found",
+        )
+
+    return RunStatusResponse(
+        run_id=run_id,
+        status=run_handle.status,
+        started_at=run_handle.started_at,
+        finished_at=run_handle.finished_at,
+    )