synth-ai 0.2.14__py3-none-any.whl → 0.4.4__py3-none-any.whl

examples/swe/task_app/hosted/rollout.py

@@ -1,1911 +0,0 @@
-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.lm.vendors.base import BaseLMResponse
-from synth_ai.task.tracing_utils import unique_sft_path
-from synth_ai.task.contracts import RolloutMode
-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
-    mode: RolloutMode  # Required: explicit RL vs EVAL mode
-
-
-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)
-        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:
-            return
-        try:
-            await self.tracer.initialize()
-        except Exception as exc:
-            logger.debug("TRACING_INIT_FAIL: %s", exc)
-        try:
-            await self.tracer.start_session(
-                session_id=self.run_id, metadata=dict(self.metadata_base)
-            )
-        except Exception as exc:
-            logger.warning("TRACING_START_FAIL: %s", exc)
-            self.enabled = False
-            self.tracer = None
-
-    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)
-
-    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(),
-                )
-            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:
-                self.session_trace = await self.tracer.end_session()
-                if self.session_trace is not None:
-                    self.session_trace.metadata.update(self.metadata_updates)
-            except Exception as exc:
-                logger.debug("TRACING_END_SESSION_FAIL: %s", exc)
-                self.session_trace = None
-            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 == "full":
-            payload = session_trace.to_dict()
-            payload.setdefault("metadata", {}).update(self.metadata_updates)
-            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()
-    # Print whether tracing is active for this rollout
-    try:
-        print(
-            f"[rollout] tracing enabled={bool(tracing_context.enabled)} run_id={request.run_id}",
-            flush=True,
-        )
-    except Exception:
-        pass
-
-    # 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
-                    # Extract inference_url from policy meta (best effort)
-                    inference_url = None
-                    if policy_handle is not None:
-                        try:
-                            policy_snapshot = policy_handle.snapshot()
-                            inference_url = policy_snapshot.get("config", {}).get("inference_url")
-                        except Exception:
-                            pass
-                    
-                    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),
-                        inference_url=inference_url,  # NEW: Required for trace correlation
-                        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)
-                    # Extract inference_url from policy meta (best effort)
-                    inference_url = None
-                    if policy_handle is not None:
-                        try:
-                            policy_snapshot = policy_handle.snapshot()
-                            inference_url = policy_snapshot.get("config", {}).get("inference_url")
-                        except Exception:
-                            pass
-                    
-                    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),
-                        inference_url=inference_url,  # NEW: Required for trace correlation
-                        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
-                    # Extract inference_url from policy meta (best effort)
-                    inference_url = None
-                    if policy_handle is not None:
-                        try:
-                            policy_snapshot = policy_handle.snapshot()
-                            inference_url = policy_snapshot.get("config", {}).get("inference_url")
-                        except Exception:
-                            pass
-                    
-                    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),
-                        inference_url=inference_url,  # NEW: Required for trace correlation
-                        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
-        # Extract inference_url from policy meta
-        inference_url = None
-        if policy_handle is not None:
-            try:
-                policy_snapshot = policy_handle.snapshot()
-                inference_url = policy_snapshot.get("config", {}).get("inference_url")
-            except Exception:
-                pass
-        
-        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),
-            inference_url=inference_url,  # NEW: Required for trace correlation
-            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,
-    )