synth-ai 0.2.9.dev11__py3-none-any.whl → 0.4.1__py3-none-any.whl

examples/swe/task_app/hosted/policy_routes.py

@@ -1,1079 +0,0 @@
-from __future__ import annotations
-
-import contextlib
-import logging
-import os
-from datetime import datetime
-from typing import Any
-
-from fastapi import APIRouter, HTTPException, Request
-from pydantic import BaseModel
-
-from .envs.crafter.policy import CrafterPolicy
-from .envs.mini_swe.policy import MiniSwePolicy
-from .inference.openai_client import create_inference_client
-from .registry import registry
-from .storage.volume import storage
-
-# Token budgeting (shared logic with inference server)
-try:
-    from ..core.algorithms.gspo.inference.token_limits import (
-        clamp_effective_max_ctx,
-    )
-except Exception:  # pragma: no cover - defensive import path fallback
-    clamp_effective_max_ctx = None  # type: ignore
-
-try:
-    import tiktoken  # type: ignore
-except Exception:  # pragma: no cover
-    tiktoken = None  # type: ignore
-
-logger = logging.getLogger(__name__)
-
-router = APIRouter()
-
-
-class PolicyCreateRequest(BaseModel):
-    policy_name: str
-    config: dict[str, Any] = {}
-    parent_policy_id: str | None = None
-    rl_run_id: str
-    bound_env_id: str | None = None
-
-
-class PolicyCreateResponse(BaseModel):
-    policy_id: str
-
-
-class PolicyStepRequest(BaseModel):
-    policy_id: str
-    observation: dict[str, Any]
-    state: dict[str, Any] | None = None
-    metadata: dict[str, Any] | None = None
-    dry_run: bool = False
-
-
-class PolicyStepResponse(BaseModel):
-    tool_calls: list[dict[str, Any]]
-    meta: dict[str, Any]
-
-
-class PolicySnapshotRequest(BaseModel):
-    policy_id: str
-
-
-class PolicySnapshotResponse(BaseModel):
-    snapshot_id: str
-    path: str
-    rl_run_id: str
-    size: int
-
-
-class PolicyRestoreRequest(BaseModel):
-    snapshot_id: str
-
-
-class PolicyRestoreResponse(BaseModel):
-    policy_id: str
-
-
-class PolicyTerminateRequest(BaseModel):
-    policy_id: str
-
-
-class PolicyTerminateResponse(BaseModel):
-    ok: bool
-
-
-@router.post("/create", response_model=PolicyCreateResponse)
-async def create_policy(
-    request: PolicyCreateRequest,
-    req: Request,
-) -> PolicyCreateResponse:
-    """Create a new policy instance."""
-    try:
-        task_app = getattr(req.app.state, "task_app", None)
-
-        # Set defaults from TaskApp / environment if not provided
-        config = dict(request.config or {})
-        inference_url_env = os.getenv("SWE_MINI_INFERENCE_URL") or os.getenv(
-            "SWE_MINI_VLLM_BASE_URL"
-        )
-        if "inference_url" not in config:
-            if inference_url_env:
-                config["inference_url"] = inference_url_env
-            elif task_app is not None and getattr(task_app, "vllm_base_url", None):
-                config["inference_url"] = task_app.vllm_base_url
-        if "model" not in config:
-            default_model = (
-                config.get("model")
-                or os.getenv("SWE_MINI_DEFAULT_MODEL")
-                or (task_app.default_model if task_app is not None else None)
-            )
-            if default_model:
-                config["model"] = default_model
-        if request.policy_name.lower() in {"swe-mini-react", "swe-mini", "swe_mini"}:
-            # Ensure tool-using policies surface run_command/submit_patch to the model.
-            # Compatibility hacks for legacy configs sometimes toggle use_tools off; for Groq
-            # and other OpenAI-compatible endpoints we hard-enable structured tool calls.
-            config["use_tools"] = True
-
-        # Create policy instance based on name
-        pname = request.policy_name.lower()
-        if pname in ["crafter-react", "crafter"]:
-            policy = CrafterPolicy(
-                inference_url=config["inference_url"],
-                model=config["model"],
-            )
-            await policy.initialize(config)
-        elif pname in ["swe-mini-react", "swe-mini", "swe_mini"]:
-            policy = MiniSwePolicy(
-                inference_url=config.get("inference_url"),
-                model=config.get("model"),
-            )
-            await policy.initialize(config)
-        elif pname in ["wordle-react", "wordle"]:
-            try:
-                from .envs.wordle.policy import WordlePolicy
-            except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Wordle policy unavailable: {e}"
-                ) from e
-
-            policy = WordlePolicy(
-                inference_url=config["inference_url"],
-                model=config["model"],
-                word_length=int(config["word_length"]),
-                max_guesses=int(config["max_guesses"]),
-            )
-            await policy.initialize(config)
-        elif pname in ["sokoban-react", "sokoban"]:
-            try:
-                from .envs.sokoban.policy import SokobanPolicy
-            except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Sokoban policy unavailable: {e}"
-                ) from e
-
-            policy = SokobanPolicy(
-                inference_url=config["inference_url"],
-                model=config["model"],
-            )
-            await policy.initialize(config)
-        elif pname in ["math-react", "math"]:
-            try:
-                from .envs.math.policy import MathPolicy
-            except Exception as e:
-                raise HTTPException(status_code=500, detail=f"Math policy unavailable: {e}") from e
-
-            policy = MathPolicy(
-                inference_url=config["inference_url"],
-                model=config["model"],
-            )
-            await policy.initialize(config)
-        else:
-            raise HTTPException(
-                status_code=422,
-                detail=f"Unknown policy name: {request.policy_name}",
-            )
-
-        # Register in memory
-        policy_id = registry.register_policy(
-            policy=policy,
-            rl_run_id=request.rl_run_id,
-            bound_env_id=request.bound_env_id,
-        )
-
-        return PolicyCreateResponse(policy_id=policy_id)
-
-    except Exception as e:
-        logger.error(f"Failed to create policy: {e}")
-        raise HTTPException(status_code=500, detail=str(e)) from e
-
-
-@router.post("/step", response_model=PolicyStepResponse)
-async def step_policy(
-    request: PolicyStepRequest,
-    req: Request,
-) -> PolicyStepResponse:
-    """Execute a policy step to generate actions."""
-    handle = registry.get_policy(request.policy_id)
-    if not handle:
-        raise HTTPException(status_code=404, detail=f"Policy {request.policy_id} not found")
-
-    try:
-        task_app = req.app.state.task_app
-        policy = handle.policy
-        tracing_context = getattr(req.state, "rollout_tracing", None)
-
-        # Format observation text conditionally for each env
-        if isinstance(request.observation, dict):
-            if isinstance(policy, CrafterPolicy):
-                from .envs.crafter.shared import format_observation as format_crafter
-
-                obs_text = format_crafter(request.observation)
-            elif isinstance(policy, MiniSwePolicy):
-                from .envs.mini_swe.shared import format_observation as format_mini_swe
-
-                obs_text = format_mini_swe(request.observation)
-            elif True:
-                try:
-                    from .envs.wordle.policy import WordlePolicy
-                except Exception:
-                    wordle_policy_cls = None  # type: ignore[assignment]
-                else:
-                    wordle_policy_cls = WordlePolicy
-
-                if wordle_policy_cls is not None and isinstance(policy, wordle_policy_cls):
-                    from .envs.wordle.shared import format_observation_wordle
-
-                # ASSERTION: Validate observation structure
-                assert request.observation is not None, "request.observation cannot be None"
-                assert isinstance(request.observation, dict), (
-                    f"request.observation must be dict, got {type(request.observation)}"
-                )
-
-                # Required keys for Wordle observation
-                required_keys = {
-                    "text",
-                    "status",
-                    "remaining_guesses",
-                    "guesses",
-                    "feedback",
-                    "reward_last",
-                    "total_reward",
-                    "terminated",
-                }
-                missing_keys = required_keys - set(request.observation.keys())
-                assert not missing_keys, f"Wordle observation missing required keys: {missing_keys}"
-
-                print("DEBUG POLICY_ROUTES: About to format Wordle observation")
-                print(f"DEBUG POLICY_ROUTES: Observation type: {type(request.observation)}")
-                print(f"DEBUG POLICY_ROUTES: Observation keys: {list(request.observation.keys())}")
-                feedback_val = request.observation["feedback"]
-                print(f"DEBUG POLICY_ROUTES: Observation feedback: {feedback_val}")
-                print(f"DEBUG POLICY_ROUTES: Observation guesses: {request.observation['guesses']}")
-                print(
-                    f"DEBUG POLICY_ROUTES: Observation text length: {len(request.observation['text'])}"
-                )
-
-                # ASSERTION: Validate feedback data
-                guesses = request.observation["guesses"]
-                feedback = request.observation["feedback"]
-                assert isinstance(guesses, list), f"guesses must be list, got {type(guesses)}"
-                assert isinstance(feedback, list), f"feedback must be list, got {type(feedback)}"
-                # Note: We don't assert equal lengths here since the environment is broken
-
-                obs_text = format_observation_wordle(request.observation)
-
-                # ASSERTION: Validate formatted output
-                assert isinstance(obs_text, str), f"obs_text must be string, got {type(obs_text)}"
-                assert len(obs_text) > 0, "obs_text cannot be empty"
-                assert "WORDLE" in obs_text, "obs_text must contain 'WORDLE' header"
-                assert "Respond with a single tool call" in obs_text, (
-                    "obs_text must contain instruction text"
-                )
-
-                print(f"DEBUG POLICY_ROUTES: Formatted obs_text length: {len(obs_text)}")
-                print(f"DEBUG POLICY_ROUTES: Formatted obs_text contains 🟩: {'🟩' in obs_text}")
-                print(f"DEBUG POLICY_ROUTES: Formatted obs_text contains 🟨: {'🟨' in obs_text}")
-                print(f"DEBUG POLICY_ROUTES: Formatted obs_text contains ⬛: {'⬛' in obs_text}")
-                print(f"DEBUG POLICY_ROUTES: Formatted obs_text first 200 chars: {obs_text[:200]}")
-            elif True:
-                try:
-                    from .envs.sokoban.policy import SokobanPolicy
-                except Exception:
-                    sokoban_policy_cls = None  # type: ignore[assignment]
-                else:
-                    sokoban_policy_cls = SokobanPolicy
-
-                if sokoban_policy_cls is not None and isinstance(policy, sokoban_policy_cls):
-                    from .envs.sokoban.shared import format_observation_sokoban
-
-                    obs_text = format_observation_sokoban(request.observation)
-            elif True:
-                try:
-                    from .envs.math.policy import MathPolicy
-                except Exception:
-                    math_policy_cls = None  # type: ignore[assignment]
-                else:
-                    math_policy_cls = MathPolicy
-                if math_policy_cls is not None and isinstance(policy, math_policy_cls):
-                    # Simple extraction of problem text
-                    try:
-                        obs_text = str(
-                            request.observation.get("problem_text") or request.observation
-                        )
-                    except Exception:
-                        obs_text = str(request.observation)
-            else:
-                obs_text = str(request.observation)
-        else:
-            obs_text = request.observation
-
-        # Merge metadata with raw observation for multimodal policies
-        step_metadata: dict[str, Any] = dict(request.metadata or {})
-        step_metadata["raw_observation"] = request.observation
-
-        # Execute policy step to get inference request
-        tool_calls, meta = await policy.step(
-            observation_text=obs_text,
-            state=request.state,
-            metadata=step_metadata,
-        )
-        # Compact tool call summary
-        with contextlib.suppress(Exception):
-            _summary: list[dict[str, Any]] = []
-            _tc = tool_calls or []
-            for _item in (_tc if isinstance(_tc, list) else []):
-                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})
-            logger.info(
-                "POLICY_STEP: tool_calls=%d summary=%s",
-                len(_tc),
-                _summary,
-            )
-
-        # If not dry run, perform inference
-        if not request.dry_run and "inference_request" in meta:
-            # CRITICAL: Validate that the inference request contains the correct prompts for the policy
-            inf_req = meta["inference_request"]
-            msgs = inf_req["messages"]
-            model_name = inf_req.get("model") or getattr(policy, "model", None) or ""
-            system_messages: list[str] = []
-            user_messages: list[str] = []
-            if msgs and len(msgs) > 0 and msgs[0]["role"] == "system":
-                sys_text = msgs[0]["content"]
-                policy_name = getattr(policy, "name", "") or type(policy).__name__.lower()
-
-                # Assert environment-specific prompts match the policy
-                if policy_name in ("wordle-react", "wordle"):
-                    if "Wordle" not in sys_text:
-                        raise ValueError(
-                            f"PROMPT MISMATCH: Wordle policy {policy_name} received system prompt without 'Wordle' keyword: {sys_text[:200]}..."
-                        )
-                    if "Crafter" in sys_text:
-                        raise ValueError(
-                            f"PROMPT MISMATCH: Wordle policy {policy_name} received Crafter system prompt: {sys_text[:200]}..."
-                        )
-
-                elif policy_name in ("crafter-react", "crafter") or isinstance(
-                    policy, CrafterPolicy
-                ):
-                    if "Crafter" not in sys_text:
-                        raise ValueError(
-                            f"PROMPT MISMATCH: Crafter policy {policy_name} received system prompt without 'Crafter' keyword: {sys_text[:200]}..."
-                        )
-                    if "Wordle" in sys_text:
-                        raise ValueError(
-                            f"PROMPT MISMATCH: Crafter policy {policy_name} received Wordle system prompt: {sys_text[:200]}..."
-                        )
-                elif policy_name in ("swe-mini-react", "swe-mini", "swe_mini") or isinstance(
-                    policy, MiniSwePolicy
-                ):
-                    if "bash" not in sys_text.lower() and "shell" not in sys_text.lower():
-                        raise ValueError(
-                            f"PROMPT MISMATCH: mini-swe policy {policy_name} missing shell instructions: {sys_text[:200]}..."
-                        )
-
-                elif policy_name in ("sokoban-react", "sokoban"):
-                    if "Sokoban" not in sys_text:
-                        raise ValueError(
-                            f"PROMPT MISMATCH: Sokoban policy {policy_name} received system prompt without 'Sokoban' keyword: {sys_text[:200]}..."
-                        )
-                    if "Crafter" in sys_text or "Wordle" in sys_text:
-                        raise ValueError(
-                            f"PROMPT MISMATCH: Sokoban policy {policy_name} received wrong environment system prompt: {sys_text[:200]}..."
-                        )
-
-                logger.info(
-                    f"✅ PROMPT VALIDATION: {policy_name} policy has correct system prompt containing expected environment keywords"
-                )
-            else:
-                logger.warning(
-                    f"⚠️ PROMPT VALIDATION: No system message found in inference request for policy {getattr(policy, 'name', type(policy).__name__)}"
-                )
-
-            # Emit full system/user prompts for observability (no secrets included)
-            try:
-
-                def _as_text(content: object) -> str:
-                    if isinstance(content, str):
-                        return content
-                    if isinstance(content, list):
-                        # Concatenate any dict segments that resemble OpenAI content parts
-                        parts: list[str] = []
-                        for seg in content:
-                            try:
-                                if isinstance(seg, dict):
-                                    txt = seg.get("text") or seg.get("content") or ""
-                                    if isinstance(txt, str):
-                                        parts.append(txt)
-                            except Exception:
-                                continue
-                        return "".join(parts)
-                    return str(content)
-
-                system_prompt_records: list[dict[str, Any]] = []
-                user_prompt_records: list[dict[str, Any]] = []
-                for message in msgs:
-                    role = message.get("role")
-                    raw_content = message.get("content")
-                    content = _as_text(raw_content)
-                    record = {"role": role, "text": content, "content": raw_content}
-                    if role == "system":
-                        system_prompt_records.append(record)
-                    elif role == "user":
-                        user_prompt_records.append(record)
-
-                logger.info(
-                    "PROMPTS: system_msgs=%d user_msgs=%d last_user_chars=%d",
-                    len(system_prompt_records),
-                    len(user_prompt_records),
-                    len(user_prompt_records[-1].get("text", "")) if user_prompt_records else 0,
-                )
-
-                if system_prompt_records:
-                    logger.info("PROMPT_DUMP_SYSTEM_BEGIN")
-                    for idx, rec in enumerate(system_prompt_records):
-                        smsg = rec.get("text", "")
-                        logger.info(f"SYSTEM[{idx}]\n{smsg}")
-                    logger.info("PROMPT_DUMP_SYSTEM_END")
-
-                if user_prompt_records:
-                    logger.info("PROMPT_DUMP_USER_BEGIN")
-                    for idx, rec in enumerate(user_prompt_records):
-                        umsg = rec.get("text", "")
-                        logger.info(f"USER[{idx}]\n{umsg}")
-                    logger.info("PROMPT_DUMP_USER_END")
-                    # Print concise preview for visibility in standard logs
-                    with contextlib.suppress(Exception):
-                        last_user = (
-                            user_prompt_records[-1].get("text", "")
-                            if user_prompt_records
-                            else ""
-                        )
-                        print(f"[task:crafter] user prompt: {last_user}", flush=True)
-            except Exception as e:
-                logger.warning(f"PROMPT_DUMP_FAILED: {e}")
-
-            if tracing_context is not None:
-                try:
-                    await tracing_context.record_policy_prompts(
-                        system_prompt_records, user_prompt_records
-                    )
-                except Exception as exc:
-                    logger.debug(f"TRACING_PROMPTS_FAIL: {exc}")
-
-            # Create inference client (choose API key by target provider)
-            # Require inference_url to be set explicitly by the rollout policy config.
-            target_url = (
-                meta.get("inference_url")
-                or getattr(policy, "inference_url", None)
-                or getattr(task_app, "vllm_base_url", None)
-            )
-
-            # Ensure meta carries the final target URL for downstream logging/clients
-            with contextlib.suppress(Exception):
-                meta["inference_url"] = target_url
-
-            # Select API key based on resolved target URL
-            api_key_override = None
-            try:
-                import os as _os
-
-                if isinstance(target_url, str):
-                    low_url = target_url.lower()
-                    # Proxy endpoints should not receive a bearer; the server-side proxy holds the vendor key
-                    if "/proxy/groq" in low_url or "/proxy/openai" in low_url:
-                        api_key_override = None
-                    elif "openai.com" in low_url:
-                        api_key_override = _os.getenv("OPENAI_API_KEY") or getattr(
-                            task_app, "openai_api_key", None
-                        )
-                    elif "groq.com" in low_url or "/proxy/groq" in low_url:
-                        api_key_override = _os.getenv("GROQ_API_KEY")
-                    else:
-                        api_key_override = (
-                            _os.getenv("SYNTH_API_KEY")
-                            or _os.getenv("OPENAI_API_KEY")
-                            or getattr(task_app, "openai_api_key", None)
-                        )
-                else:
-                    api_key_override = (
-                        _os.getenv("SYNTH_API_KEY")
-                        or _os.getenv("OPENAI_API_KEY")
-                        or getattr(task_app, "openai_api_key", None)
-                    )
-            except Exception:
-                api_key_override = None
-
-            if api_key_override:
-                try:
-                    masked = f"{api_key_override[:6]}…{api_key_override[-4:]}"
-                except Exception:
-                    masked = "<masked>"
-                logger.debug(f"INFERENCE_AUTH: Using bearer key {masked}")
-            else:
-                logger.warning(
-                    "INFERENCE_AUTH: No API key resolved for inference request; downstream may 401"
-                )
-
-            client = create_inference_client(task_app, api_key=api_key_override)
-
-            # Add policy identification header for observability
-            policy_name = getattr(policy, "name", "") or type(policy).__name__.lower()
-            extra_headers = {"X-Policy-Name": policy_name}
-
-            # Apply input truncation to avoid 422 from inference server
-            try:
-                model_name = inf_req.get("model") or getattr(policy, "model", None) or ""
-                env_max_ctx = None
-                try:
-                    _env_max = int(os.getenv("CHAT_MAX_MODEL_LEN", "0") or 0)
-                    env_max_ctx = _env_max if _env_max > 0 else None
-                except Exception:
-                    env_max_ctx = None
-                # Compute effective max context and safety margin
-                eff_ctx = None
-                if clamp_effective_max_ctx is not None:
-                    eff_ctx = clamp_effective_max_ctx(
-                        model_name=model_name,
-                        configured_max_model_len=None,
-                        env_max_model_len=env_max_ctx,
-                    )
-                # Hard lower-only chat input cap if provided
-                try:
-                    hard_input_cap = int(os.getenv("CHAT_MAX_INPUT_TOKENS", "0") or 0)
-                    hard_input_cap = hard_input_cap if hard_input_cap > 0 else None
-                except Exception:
-                    hard_input_cap = None
-                try:
-                    safety_margin = int(os.getenv("CHAT_BUDGET_SAFETY", "64").strip() or 64)
-                except Exception:
-                    safety_margin = 64
-
-                # Determine budget
-                budget = None
-                if isinstance(eff_ctx, int) and eff_ctx > 0:
-                    budget = max(256, eff_ctx - safety_margin)
-                if isinstance(hard_input_cap, int) and hard_input_cap > 0:
-                    budget = min(budget, hard_input_cap) if budget is not None else hard_input_cap
-
-                if budget is not None and budget > 0 and isinstance(msgs, list):
-                    # Choose tokenizer
-                    enc = None
-                    if tiktoken is not None:
-                        try:
-                            if model_name:
-                                enc = tiktoken.encoding_for_model(model_name)
-                            else:
-                                enc = tiktoken.get_encoding("cl100k_base")
-                        except Exception:
-                            try:
-                                enc = tiktoken.get_encoding("cl100k_base")
-                            except Exception:
-                                enc = None
-
-                    def _content_to_text(content: object) -> str:
-                        if isinstance(content, str):
-                            return content
-                        if isinstance(content, list):
-                            parts: list[str] = []
-                            for seg in content:
-                                try:
-                                    if isinstance(seg, dict):
-                                        txt = seg.get("text") or seg.get("content") or ""
-                                        if isinstance(txt, str):
-                                            parts.append(txt)
-                                except Exception:
-                                    continue
-                            return "".join(parts)
-                        try:
-                            return str(content)
-                        except Exception:
-                            return ""
-
-                    def _count_tokens(text: str) -> int:
-                        if enc is None:
-                            # Fall back to character count heuristic (~4 chars per token)
-                            try:
-                                return max(1, int(len(text) / 4))
-                            except Exception:
-                                return len(text)
-                        try:
-                            return len(enc.encode(text))
-                        except Exception:
-                            return max(1, int(len(text) / 4))
-
-                    def _count_messages_tokens(messages: list[dict[str, Any]]) -> int:
-                        total = 0
-                        for m in messages:
-                            total += _count_tokens(_content_to_text(m.get("content")))
-                        return total
-
-                    def _truncate_messages_to_budget(
-                        messages: list[dict[str, Any]],
-                        max_tokens: int,
-                    ) -> tuple[list[dict[str, Any]], int, int, int]:
-                        before = _count_messages_tokens(messages)
-                        if before <= max_tokens:
-                            return messages, before, before, len(messages)
-                        # Always try to preserve the first system message if present
-                        system_msg = None
-                        start_idx = 0
-                        if messages and messages[0].get("role") == "system":
-                            system_msg = messages[0]
-                            start_idx = 1
-                        kept_rev: list[dict[str, Any]] = []
-                        total = _count_messages_tokens([system_msg] if system_msg else [])
-                        # Walk from the end keeping most recent messages
-                        for m in reversed(messages[start_idx:]):
-                            t = _count_tokens(_content_to_text(m.get("content")))
-                            if total + t <= max_tokens:
-                                kept_rev.append(m)
-                                total += t
-                            else:
-                                # Try to keep a truncated version of this message if we have some budget left
-                                remaining = max_tokens - total
-                                if remaining > 16:  # keep at least a little context
-                                    txt = _content_to_text(m.get("content"))
-                                    # Binary search-ish trim by tokens
-                                    low, high = 0, len(txt)
-                                    best = None
-                                    while low <= high:
-                                        mid = (low + high) // 2
-                                        candidate = txt[-mid:]
-                                        if _count_tokens(candidate) <= remaining:
-                                            best = candidate
-                                            low = mid + 1
-                                        else:
-                                            high = mid - 1
-                                    if best is not None and best:
-                                        m2 = dict(m)
-                                        m2["content"] = best
-                                        kept_rev.append(m2)
-                                        total += _count_tokens(best)
-                                break
-                        kept = list(reversed(kept_rev))
-                        if system_msg is not None:
-                            kept = [system_msg] + kept
-                        after = _count_messages_tokens(kept)
-                        return kept, before, after, len(kept)
-
-                    new_msgs, before_toks, after_toks, kept_count = _truncate_messages_to_budget(
-                        msgs, int(budget)
-                    )
-                    if new_msgs is not msgs:
-                        inf_req["messages"] = new_msgs
-                        with contextlib.suppress(Exception):
-                            logger.info(
-                                {
-                                    "chat_truncated": True,
-                                    "token_budget": int(budget),
-                                    "before_tokens": int(before_toks),
-                                    "after_tokens": int(after_toks),
-                                    "kept_msgs": int(kept_count),
-                                }
-                            )
-            except Exception as _trunc_e:
-                logger.warning(f"CHAT_TRUNCATION_FAILED: {type(_trunc_e).__name__}: {_trunc_e}")
-
-            # Formal assertion: If tools are expected, ensure tool_choice and tools are set
-            if policy_name in (
-                "wordle-react",
-                "sokoban-react",
-                "crafter-react",
-            ) and getattr(policy, "use_tools", True):
-                req_tools = meta["inference_request"]["tools"]
-                req_tool_choice = meta["inference_request"]["tool_choice"]
-                req_stop_after = meta["inference_request"]["stop_after_tool_calls"]
-                logger.info(
-                    f"TOOLCALL_CONFIG: policy={policy_name} tools_present={bool(req_tools)} tool_choice={req_tool_choice} stop_after={req_stop_after}"
-                )
-                if not req_tools or req_tool_choice != "required":
-                    raise HTTPException(
-                        status_code=500,
-                        detail=f"TOOLCALL_ASSERTION_FAIL: Missing tools or tool_choice!=required for policy {policy_name}",
-                    )
-
-            # Call inference service with retries for Flash cold-start (503)
-            import time as _t
-
-            # Prompt diagnostics before sending to inference: build chat template locally,
-            # count tokens, and log the first 10k tokens if oversized. Also stash a
-            # compact preview in meta so the trainer can surface it.
-            with contextlib.suppress(Exception):
-                req_for_diag = meta.get("inference_request", {})
-                model_for_diag = req_for_diag.get("model") or getattr(policy, "model", None) or ""
-                messages_for_diag = req_for_diag.get("messages") or []
-                if model_for_diag and messages_for_diag:
-                    from transformers import AutoTokenizer
-
-                    tok = AutoTokenizer.from_pretrained(model_for_diag)
-                    prompt_preview = tok.apply_chat_template(
-                        messages_for_diag,
-                        add_generation_prompt=True,
-                        tokenize=False,
-                    )
-                    ids = tok.encode(prompt_preview, add_special_tokens=False)
-                    max_len = getattr(tok, "model_max_length", None)
-                    over_limit = False
-                    with contextlib.suppress(Exception):
-                        over_limit = (
-                            isinstance(max_len, int) and max_len > 0 and len(ids) > int(max_len)
-                        )
-                    if over_limit or len(ids) > 10000:
-                        preview_ids = ids[:10000]
-                        preview_text = tok.decode(
-                            preview_ids,
-                            skip_special_tokens=False,
-                        )
-                        with contextlib.suppress(Exception):
-                            logger.warning(
-                                {
-                                    "prompt_token_overflow_local": True,
-                                    "model": str(model_for_diag),
-                                    "token_count": int(len(ids)),
-                                    "model_max_length": int(max_len)
-                                    if isinstance(max_len, int)
-                                    else None,
-                                    "preview_tokens_logged": int(len(preview_ids)),
-                                    "prompt_preview_first_10k_tokens": preview_text,
-                                }
-                            )
-                        with contextlib.suppress(Exception):
-                            meta["prompt_debug"] = {
-                                "token_count": int(len(ids)),
-                                "model_max_length": int(max_len)
-                                if isinstance(max_len, int)
-                                else None,
-                                "preview_first_10k_tokens": preview_text,
-                            }
-
-            # Emit the exact prompt/messages and tools before calling the LLM (bounded preview)
-            with contextlib.suppress(Exception):
-                req_dump = meta.get("inference_request", {})
-                msgs = req_dump.get("messages")
-                tools_dump = req_dump.get("tools")
-                if isinstance(msgs, list):
-                    # Print compact messages structure and tool schema with bounded length
-                    import json as _json
-
-                    msgs_compact = _json.dumps(msgs)[:20000]
-                    tools_compact = (
-                        _json.dumps(tools_dump)[:8000] if tools_dump is not None else None
-                    )
-                    print(
-                        {
-                            "llm.call": True,
-                            "policy": str(policy_name),
-                            "messages_preview": msgs_compact,
-                            "tools_preview": tools_compact,
-                        }
-                    )
-
-            # Normalize request for non-OpenAI endpoints (strict schemas)
-            with contextlib.suppress(Exception):
-                base = str(target_url or "")
-                is_openai_dotcom = "openai.com" in base.lower()
-                if not is_openai_dotcom:
-                    req_body = meta.get("inference_request", {})
-                    if isinstance(req_body, dict):
-                        # Force structured tool_choice if a bare "required" is present
-                        if req_body.get("tool_choice") == "required":
-                            func_name = "interact_many"
-                            with contextlib.suppress(Exception):
-                                tools_arr = req_body.get("tools") or []
-                                if isinstance(tools_arr, list) and tools_arr:
-                                    f = (
-                                        tools_arr[0].get("function")
-                                        if isinstance(tools_arr[0], dict)
-                                        else None
-                                    )
-                                    cand = (f or {}).get("name") if isinstance(f, dict) else None
-                                    if isinstance(cand, str) and cand:
-                                        func_name = cand
-                            req_body["tool_choice"] = {
-                                "type": "function",
-                                "function": {"name": func_name},
-                            }
-                            req_body["parallel_tool_calls"] = False
-                            req_body.setdefault("function_call", {"name": func_name})
-                        # Inject extra_body for thinking controls expected by Modal service
-                        with contextlib.suppress(Exception):
-                            tb = req_body.get("thinking_budget")
-                            tm = str(req_body.get("thinking_mode") or "").lower()
-                            enable_thinking = bool(tb) or tm == "think"
-                            extra = dict(req_body.get("extra_body") or {})
-                            chat_kwargs = dict(extra.get("chat_template_kwargs") or {})
-                            if enable_thinking:
-                                chat_kwargs["enable_thinking"] = True
-                            if isinstance(tb, int | float | str) and str(tb).strip():
-                                with contextlib.suppress(Exception):
-                                    chat_kwargs["thinking_budget"] = int(tb)
-                            if chat_kwargs:
-                                extra["chat_template_kwargs"] = chat_kwargs
-                            # Ensure stop_after_tool_calls honored via extra_body for stricter servers
-                            extra.setdefault("stop_after_tool_calls", 1)
-                            if extra:
-                                req_body["extra_body"] = extra
-                        # Provide a conservative default temperature if missing
-                        if "temperature" not in req_body:
-                            req_body["temperature"] = 0.1
-                        meta["inference_request"] = req_body
-
-                # Strip image parts: Crafter policy currently only uses text prompts.
-                # Some providers reject image_url payloads entirely, so always flatten to plain text.
-                req_body2 = meta.get("inference_request", {})
-                if isinstance(req_body2, dict):
-                    msgs = req_body2.get("messages")
-                    if isinstance(msgs, list):
-                        new_msgs = []
-                        changed = False
-                        for m in msgs:
-                            try:
-                                if isinstance(m, dict):
-                                    content = m.get("content")
-                                    if isinstance(content, list):
-                                        parts: list[str] = []
-                                        for seg in content:
-                                            if isinstance(seg, dict):
-                                                txt = seg.get("text") or seg.get("content")
-                                                if isinstance(txt, str) and txt:
-                                                    parts.append(txt)
-                                        m2 = dict(m)
-                                        m2["content"] = "\n".join(parts)
-                                        new_msgs.append(m2)
-                                        changed = True
-                                    else:
-                                        new_msgs.append(m)
-                                else:
-                                    new_msgs.append(m)
-                            except Exception:
-                                new_msgs.append(m)
-                        if changed:
-                            req_body2["messages"] = new_msgs
-                            meta["inference_request"] = req_body2
-
-            _t_start = _t.time()
-            call_started_at = datetime.utcnow()
-            inference_response = await client.generate_with_retries(
-                request=meta["inference_request"],
-                base_url=meta["inference_url"],
-                max_retries=12,
-                backoff_factor=2.0,
-                extra_headers=extra_headers,
-            )
-            meta["inference_ms"] = int((_t.time() - _t_start) * 1000)
-            call_completed_at = datetime.utcnow()
-
-            provider_url = str(meta.get("inference_url") or "")
-            low_url = provider_url.lower()
-            if "groq" in low_url:
-                provider_name = "groq"
-            elif "openai" in low_url:
-                provider_name = "openai"
-            else:
-                provider_name = "custom"
-
-            # Parse response to tool calls
-            tool_calls = policy.parse_response_to_tool_calls(
-                response=inference_response,
-                use_tools=getattr(policy, "use_tools", True),
-            )
-
-            # Debug logging (echo tool calls)
-            if not tool_calls:
-                # Structured error log with small preview; avoid dumping full response repeatedly
-                preview = str(inference_response)[:400]
-                logger.error(
-                    f"TOOLCALL_PARSE_FAIL: policy={policy_name} parsed=0 preview={preview}"
-                )
-            else:
-                try:
-                    import json as _json
-
-                    print(
-                        {
-                            "tool_calls_parsed": int(len(tool_calls)),
-                            "tool_calls_preview": _json.dumps(tool_calls)[:20000],
-                        }
-                    )
-                except Exception:
-                    logger.info(f"Parsed {len(tool_calls)} tool calls: {tool_calls}")
-
-            # Add response to metadata
-            # Parse tool calls from model response using policy-specific parser
-            try:
-                if hasattr(policy, "parse_response_to_tool_calls"):
-                    parsed = policy.parse_response_to_tool_calls(
-                        inference_response, getattr(policy, "use_tools", True)
-                    )
-                else:
-                    parsed = policy.parse_model_response(inference_response, request.observation)
-                # Replace tool_calls with parsed result
-                if isinstance(parsed, list):
-                    tool_calls = parsed
-                with contextlib.suppress(Exception):
-                    logger.info(
-                        "TOOLCALL_PARSE: parsed=%d has_tools=%s example=%r",
-                        len(tool_calls) if isinstance(tool_calls, list) else -1,
-                        bool(getattr(policy, "use_tools", True)),
-                        (tool_calls[0] if isinstance(tool_calls, list) and tool_calls else None),
-                    )
-            except Exception as _pe:
-                logger.warning(f"Failed to parse tool calls: {str(_pe)}")
-            # Attach raw response + usage for observability
-            meta["raw_response"] = inference_response
-            if "usage" in inference_response:
-                meta["usage"] = inference_response["usage"]
-
-            if tracing_context is not None:
-                try:
-                    await tracing_context.record_llm_call(
-                        inference_request=meta["inference_request"],
-                        inference_response=inference_response,
-                        tool_calls=tool_calls,
-                        provider=provider_name,
-                        model_name=model_name,
-                        started_at=call_started_at,
-                        completed_at=call_completed_at,
-                        latency_ms=meta.get("inference_ms"),
-                    )
-                except Exception as exc:
-                    logger.debug(f"TRACING_LLM_FAIL: {exc}")
-
-        return PolicyStepResponse(
-            tool_calls=tool_calls,
-            meta=meta,
-        )
-
-    except Exception as e:
-        logger.error(f"Failed to step policy {request.policy_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e)) from e
-
-
-@router.post("/snapshot", response_model=PolicySnapshotResponse)
-async def snapshot_policy(request: PolicySnapshotRequest) -> PolicySnapshotResponse:
-    """Create a snapshot of the policy state."""
-    handle = registry.get_policy(request.policy_id)
-    if not handle:
-        raise HTTPException(status_code=404, detail=f"Policy {request.policy_id} not found")
-
-    try:
-        # Serialize policy state
-        state_dict = await handle.policy.serialize()
-
-        # Save to volume
-        snapshot_id, path, size = storage.save_snapshot(
-            rl_run_id=handle.rl_run_id,
-            kind="policy",
-            state_dict=state_dict,
-        )
-
-        # Register snapshot
-        registry.register_snapshot(
-            kind="policy",
-            rl_run_id=handle.rl_run_id,
-            size=size,
-            path=path,
-        )
-
-        return PolicySnapshotResponse(
-            snapshot_id=snapshot_id,
-            path=path,
-            rl_run_id=handle.rl_run_id,
-            size=size,
-        )
-
-    except Exception as e:
-        logger.error(f"Failed to snapshot policy {request.policy_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e)) from e
-
-
-@router.post("/restore", response_model=PolicyRestoreResponse)
-async def restore_policy(request: PolicyRestoreRequest) -> PolicyRestoreResponse:
-    """Restore a policy from a snapshot."""
-    snapshot = registry.get_snapshot(request.snapshot_id)
-    if not snapshot:
-        raise HTTPException(status_code=404, detail=f"Snapshot {request.snapshot_id} not found")
-
-    if snapshot.kind != "policy":
-        raise HTTPException(
-            status_code=422,
-            detail=f"Snapshot {request.snapshot_id} is not a policy snapshot",
-        )
-
-    try:
-        # Load snapshot from volume
-        state_dict, meta = storage.load_snapshot(
-            rl_run_id=snapshot.rl_run_id,
-            kind="policy",
-            snapshot_id=request.snapshot_id,
-        )
-
-        # Recreate policy
-        policy_name = state_dict["name"]
-        low = policy_name.lower()
-        if low in ["crafter-react", "crafter"]:
-            policy = await CrafterPolicy.deserialize(state_dict)
-        elif low in ["swe-mini-react", "swe-mini", "swe_mini"]:
-            policy = await MiniSwePolicy.deserialize(state_dict)
-        elif low in ["wordle-react", "wordle"]:
-            try:
-                from .envs.wordle.policy import WordlePolicy
-            except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Wordle policy unavailable: {e}"
-                ) from e
-            policy = await WordlePolicy.deserialize(state_dict)
-        elif low in ["sokoban-react", "sokoban"]:
-            try:
-                from .envs.sokoban.policy import SokobanPolicy
-            except Exception as e:
-                raise HTTPException(
-                    status_code=500, detail=f"Sokoban policy unavailable: {e}"
-                ) from e
-            policy = await SokobanPolicy.deserialize(state_dict)
-        else:
-            raise HTTPException(
-                status_code=422,
-                detail=f"Unknown policy name in snapshot: {policy_name}",
-            )
-
-        # Register new instance
-        policy_id = registry.register_policy(
-            policy=policy,
-            rl_run_id=snapshot.rl_run_id,
-        )
-
-        return PolicyRestoreResponse(policy_id=policy_id)
-
-    except Exception as e:
-        logger.error(f"Failed to restore policy from snapshot {request.snapshot_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e)) from e
-
-
-@router.post("/terminate", response_model=PolicyTerminateResponse)
-async def terminate_policy(request: PolicyTerminateRequest) -> PolicyTerminateResponse:
-    """Terminate a policy and clean up resources."""
-    handle = registry.get_policy(request.policy_id)
-    if not handle:
-        raise HTTPException(status_code=404, detail=f"Policy {request.policy_id} not found")
-
-    try:
-        # Call terminate on the policy
-        await handle.policy.terminate()
-
-        # Remove from registry
-        registry.remove_policy(request.policy_id)
-
-        return PolicyTerminateResponse(ok=True)
-
-    except Exception as e:
-        logger.error(f"Failed to terminate policy {request.policy_id}: {e}")
-        raise HTTPException(status_code=500, detail=str(e)) from e