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

examples/task_apps/sokoban/task_app.py

@@ -1,1058 +0,0 @@
-from __future__ import annotations
-
-import json
-import os
-import re
-import contextlib
-import time
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Mapping, Optional, Sequence
-
-import httpx
-
-from fastapi import APIRouter, HTTPException, Request
-from fastapi.exceptions import RequestValidationError
-from fastapi.responses import JSONResponse
-
-from synth_ai.environments.environment.tools import EnvToolCall
-from synth_ai.environments.examples.sokoban.environment import SokobanEnvironment
-from synth_ai.environments.examples.sokoban.taskset import (
-    SokobanTaskInstance,
-    SokobanTaskSet,
-    create_task_instance_from_seed,
-)
-from synth_ai.task.apps import TaskAppEntry, register_task_app
-from synth_ai.task.contracts import (
-    RolloutMetrics,
-    RolloutRequest,
-    RolloutResponse,
-    RolloutStep,
-    RolloutTrajectory,
-    TaskInfo,
-)
-from synth_ai.task.auth import is_api_key_header_authorized, normalize_environment_api_key
-from synth_ai.task.server import TaskAppConfig, create_task_app
-
-
-ACTION_ID_TO_NAME = {0: "left", 1: "up", 2: "right", 3: "down"}
-ACTION_TOKEN_TO_ID = {
-    "0": 0,
-    "1": 1,
-    "2": 2,
-    "3": 3,
-    "left": 0,
-    "move_left": 0,
-    "west": 0,
-    "l": 0,
-    "up": 1,
-    "move_up": 1,
-    "north": 1,
-    "u": 1,
-    "right": 2,
-    "move_right": 2,
-    "east": 2,
-    "r": 2,
-    "down": 3,
-    "move_down": 3,
-    "south": 3,
-    "d": 3,
-}
-
-SOKOBAN_SYSTEM_PROMPT = """You are an agent playing Sokoban.
-The grid uses characters: '#' wall, '_' floor, 'O' box, '√' box on target, 'X' target, 'P' player.
-Always respond with a single tool call named interact_many containing 1-5 actions.
-Valid action tokens are digits 0/1/2/3 or their direction words (left/up/right/down).
-Mapping: 0=left, 1=up, 2=right, 3=down. Avoid undoing progress and focus on pushing boxes onto targets."""
-
-
-def _short_text(value: Any, *, limit: int = 280) -> str:
-    if value is None:
-        return ""
-    if isinstance(value, str):
-        text = value.strip()
-        return text if len(text) <= limit else text[: limit - 1] + "…"
-    if isinstance(value, (int, float, bool)):
-        return str(value)
-    try:
-        text = json.dumps(value, ensure_ascii=False)
-    except Exception:
-        text = str(value)
-    text = text.strip()
-    return text if len(text) <= limit else text[: limit - 1] + "…"
-
-
-def _summarize_observation(observation: Any) -> str:
-    if isinstance(observation, dict):
-        for key in ("room_text", "observation", "grid"):
-            value = observation.get(key)
-            if isinstance(value, str) and value.strip():
-                return _short_text(value, limit=512)
-        preview = {
-            key: observation.get(key)
-            for key in ("player_position", "boxes_on_target", "num_boxes", "steps_taken")
-            if key in observation
-        }
-        if preview:
-            return _short_text(preview, limit=512)
-    return _short_text(observation, limit=512)
-
-
-def _format_tool_calls(tool_calls: Sequence[Dict[str, Any]] | None) -> str:
-    if not tool_calls:
-        return "<noop>"
-    formatted: list[str] = []
-    for call in tool_calls:
-        args = call.get("args") if isinstance(call, dict) else None
-        if not isinstance(args, dict):
-            continue
-        if "actions" in args and isinstance(args["actions"], list):
-            parts: list[str] = []
-            for item in args["actions"]:
-                try:
-                    val = int(item)
-                except Exception:
-                    token = str(item).strip().lower()
-                    val = ACTION_TOKEN_TO_ID.get(token)
-                name = ACTION_ID_TO_NAME.get(val, str(item)) if val is not None else str(item)
-                parts.append(str(name))
-            if parts:
-                formatted.append("[" + ", ".join(parts) + "]")
-            continue
-        action = args.get("action")
-        if action is None:
-            continue
-        try:
-            action = int(action)
-        except Exception:
-            token = str(action).strip().lower()
-            action = ACTION_TOKEN_TO_ID.get(token, action)
-        name = ACTION_ID_TO_NAME.get(action, str(action))
-        formatted.append(str(name))
-    return ", ".join(formatted) if formatted else "<noop>"
-
-
-def _build_trace_payload(
-    request: RolloutRequest,
-    steps: Sequence[RolloutStep],
-    metrics: RolloutMetrics,
-    *,
-    difficulty: str,
-    initial_observation: Any,
-    provider: str = "local",
-) -> Dict[str, Any]:
-    created_at = datetime.now(timezone.utc)
-    base_time = time.time()
-    event_history: list[dict[str, Any]] = []
-    markov_messages: list[dict[str, Any]] = []
-    session_steps: list[dict[str, Any]] = []
-
-    if not steps:
-        observation_text = _summarize_observation(initial_observation)
-        event_time = base_time
-        observation_msg = {
-            "content": {"text": observation_text},
-            "message_type": "observation",
-            "time_record": {"event_time": event_time},
-            "metadata": {"step_index": 0},
-        }
-        markov_messages.append(observation_msg)
-        event_history.append(
-            {
-                "system_instance_id": "sokoban.step.0",
-                "time_record": {"event_time": event_time},
-                "reward": 0.0,
-                "terminated": True,
-                "truncated": False,
-                "metadata": {
-                    "tool_calls": [],
-                },
-            }
-        )
-        session_steps.append(
-            {
-                "step_id": "step_0",
-                "step_index": 0,
-                "events": [event_history[-1]],
-                "markov_blanket_messages": markov_messages[-1:],
-                "step_metadata": {"reward": 0.0, "done": True, "truncated": False},
-            }
-        )
-    else:
-        for idx, step in enumerate(steps):
-            event_time = base_time + idx * 0.01
-            observation_text = _summarize_observation(step.obs)
-            action_text = _format_tool_calls(step.tool_calls)
-            observation_msg = {
-                "content": {"text": observation_text},
-                "message_type": "observation",
-                "time_record": {"event_time": event_time},
-                "metadata": {"step_index": idx},
-            }
-            action_msg = {
-                "content": {"text": action_text},
-                "message_type": "action",
-                "time_record": {"event_time": event_time + 0.0005},
-                "metadata": {"step_index": idx},
-            }
-            markov_messages.extend([observation_msg, action_msg])
-            reward_val = float(step.reward or 0.0)
-            event_history.append(
-                {
-                    "system_instance_id": f"sokoban.step.{idx}",
-                    "time_record": {"event_time": event_time},
-                    "reward": reward_val,
-                    "terminated": bool(step.done),
-                    "truncated": bool(step.truncated),
-                    "metadata": {
-                        "tool_calls": step.tool_calls,
-                        "info": step.info or {},
-                    },
-                }
-            )
-            session_steps.append(
-                {
-                    "step_id": f"step_{idx}",
-                    "step_index": idx,
-                    "events": [event_history[-1]],
-                    "markov_blanket_messages": [observation_msg, action_msg],
-                    "step_metadata": {
-                        "reward": reward_val,
-                        "done": bool(step.done),
-                        "truncated": bool(step.truncated),
-                    },
-                }
-            )
-
-    session_trace = {
-        "session_id": str(request.run_id),
-        "created_at": created_at.isoformat(),
-        "metadata": {
-            "task": "sokoban",
-            "difficulty": difficulty,
-            "seed": request.env.seed,
-            "provider": provider,
-            "env": request.env.model_dump(),
-            "policy": request.policy.model_dump(),
-        },
-        "session_time_steps": session_steps,
-        "event_history": event_history,
-        "markov_blanket_message_history": markov_messages,
-    }
-
-    return {
-        "version": 3,
-        "session_trace": session_trace,
-        "run_id": request.run_id,
-        "policy_id": request.policy.policy_id or request.policy.policy_name,
-        "reward": metrics.mean_return,
-        "episode_returns": metrics.episode_returns,
-        "mean_return": metrics.mean_return,
-        "num_steps": metrics.num_steps,
-    }
-
-
-
-def _task_info() -> TaskInfo:
-    return TaskInfo(
-        task={"id": "sokoban", "name": "Sokoban", "version": "1.0.0"},
-        environment="sokoban",
-        action_space={
-            "type": "tool_call",
-            "tools": [{"name": "interact", "schema": {"action": "int"}}],
-            "max_calls": 1,
-        },
-        observation={"summary": "Sokoban grid observation", "keys": ["grid", "player"]},
-        dataset={"id": "sokoban", "name": "Sokoban", "version": "1.0.0"},
-        rubric={"version": "1", "criteria_count": 1, "source": "inline"},
-        inference={"supports_proxy": False},
-        limits={"max_turns": 200},
-    )
-
-
-router = APIRouter()
-
-
-async def rollout_executor(request: RolloutRequest, fastapi_request: Request) -> RolloutResponse:
-    policy_cfg = dict(request.policy.config or {})
-    provider = str(policy_cfg.get("provider") or "").strip().lower()
-    if provider == "groq":
-        return await _rollout_with_groq(request, fastapi_request, policy_cfg)
-    if provider == "openai":
-        return await _rollout_with_openai(request, fastapi_request, policy_cfg)
-
-    taskset: SokobanTaskSet = fastapi_request.app.state.sokoban_taskset
-    seed = request.env.seed or 0
-    difficulty = (request.env.config or {}).get("difficulty") or "easy"
-    # Create deterministic instance from seed
-    instance: SokobanTaskInstance = await create_task_instance_from_seed(str(difficulty), int(seed))
-    env = SokobanEnvironment(instance)
-    obs = await env.initialize()
-    initial_observation = obs
-
-    tool_calls: List[Dict[str, Any]] = []
-    # If a predefined action sequence is provided, execute it (evaluation-style)
-    actions: Optional[Sequence[int]] = None
-    try:
-        cfg = request.policy.config or {}
-        if isinstance(cfg.get("actions"), list):
-            actions = [int(a) for a in cfg["actions"]]
-    except Exception:
-        actions = None
-
-    last_obs: Any = obs
-    steps: List[RolloutStep] = []
-    max_steps = int((request.env.config or {}).get("max_steps") or 50)
-    executed = 0
-    if actions:
-        for a in actions[:max_steps]:
-            last_obs = await env.step(EnvToolCall(tool="interact", args={"action": int(a)}))
-            executed += 1
-            steps.append(
-                RolloutStep(obs=last_obs, tool_calls=[{"tool": "interact", "args": {"action": int(a)}}], reward=0.0, done=False, info={})
-            )
-    # Mark episode end (single-episode trajectory)
-    final = {"observation": last_obs, "reward": 0.0}
-    if not steps:
-        steps = [RolloutStep(obs=last_obs, tool_calls=[], reward=0.0, done=True, info={})]
-    
-    # Extract inference_url from policy config (None for manual rollouts)
-    inference_url = policy_cfg.get("inference_url")
-    
-    traj = RolloutTrajectory(
-        env_id="sokoban",
-        policy_id=request.policy.policy_id or "policy",
-        steps=steps,
-        final=final,
-        length=len(steps),
-        inference_url=inference_url,  # NEW: Required for trace correlation
-    )
-    metrics = RolloutMetrics(
-        episode_returns=[final.get("reward", 0.0) or 0.0],
-        mean_return=final.get("reward", 0.0) or 0.0,
-        num_steps=len(steps),
-        num_episodes=1,
-        outcome_score=None,
-        events_score=None,
-        details={},
-    )
-    trace_payload = _build_trace_payload(
-        request,
-        steps,
-        metrics,
-        difficulty=str(difficulty),
-        initial_observation=initial_observation,
-    )
-    return RolloutResponse(
-        run_id=request.run_id,
-        trajectories=[traj],
-        branches={},
-        metrics=metrics,
-        aborted=False,
-        ops_executed=1 + executed,
-        trace=trace_payload,
-    )
-
-
-def _format_sokoban_prompt(observation: dict[str, Any], last_actions: list[int]) -> str:
-    grid = observation.get("room_text", "")
-    boxes = observation.get("boxes_on_target", 0)
-    total_boxes = observation.get("num_boxes", boxes)
-    position = observation.get("player_position", ())
-    reward_last = observation.get("reward_last", 0.0)
-    steps_taken = observation.get("steps_taken", 0)
-    max_steps = observation.get("max_steps", 0)
-    last_str = (
-        ", ".join(ACTION_ID_TO_NAME.get(a, str(a)) for a in last_actions) if last_actions else "none"
-    )
-    return (
-        f"Step {steps_taken} / {max_steps}\n"
-        f"Player position: {position}\n"
-        f"Boxes on target: {boxes} / {total_boxes}\n"
-        f"Last reward: {reward_last}\n"
-        f"Previous actions: {last_str}\n"
-        "Grid:\n"
-        f"{grid}\n"
-        "Select up to five next actions via the interact_many tool."
-    )
-
-
-def _extract_actions_from_response(
-    response: dict[str, Any], max_actions: int
-) -> list[int]:
-    import json as json_lib
-    print(f"[extract] FULL RESPONSE:", flush=True)
-    print(json_lib.dumps(response, indent=2)[:2000], flush=True)
-    
-    actions: list[int] = []
-    choices = response.get("choices") or []
-    print(f"[extract] {len(choices)} choices", flush=True)
-    if choices:
-        msg = choices[0].get("message", {})
-        print(f"[extract] tool_calls: {msg.get('tool_calls')}", flush=True)
-        print(f"[extract] content: {msg.get('content')}", flush=True)
-        print(f"[extract] finish_reason: {choices[0].get('finish_reason')}", flush=True)
-    for choice in choices:
-        message = choice.get("message") or {}
-        tool_calls = message.get("tool_calls") or []
-        for tool_call in tool_calls:
-            function = tool_call.get("function") or {}
-            arguments = function.get("arguments")
-            payload: dict[str, Any] | None = None
-            if isinstance(arguments, str):
-                try:
-                    payload = json.loads(arguments)
-                except json.JSONDecodeError:
-                    payload = None
-            elif isinstance(arguments, dict):
-                payload = arguments
-            if not payload:
-                continue
-            raw_actions = payload.get("actions")
-            if isinstance(raw_actions, list):
-                for item in raw_actions:
-                    if isinstance(item, int) and item in ACTION_ID_TO_NAME:
-                        actions.append(int(item))
-                        continue
-                    if isinstance(item, str):
-                        token = item.strip().lower()
-                        if token in ACTION_TOKEN_TO_ID:
-                            actions.append(ACTION_TOKEN_TO_ID[token])
-            if actions:
-                break
-        if actions:
-            break
-
-    if not actions and choices:
-        # Fallback: parse tokens from assistant text
-        text = choices[0].get("message", {}).get("content") or ""
-        tokens = re.findall(r"[0-3a-zA-Z_]+", text)
-        for tok in tokens:
-            token = tok.strip().lower()
-            if token in ACTION_TOKEN_TO_ID:
-                actions.append(ACTION_TOKEN_TO_ID[token])
-
-    if len(actions) > max_actions:
-        return actions[:max_actions]
-    return actions
-
-
-async def _call_groq_chat(
-    client: httpx.AsyncClient,
-    api_key: str,
-    payload: dict[str, Any],
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    try:
-        response = await client.post(
-            "https://api.groq.com/openai/v1/chat/completions",
-            json=payload,
-            headers={"Authorization": f"Bearer {api_key}"},
-        )
-        response.raise_for_status()
-        data = response.json()
-        return data, {
-            "status": response.status_code,
-            "headers": dict(response.headers),
-            "body": data,
-        }
-    except httpx.HTTPStatusError as exc:
-        try:
-            body = exc.response.json()
-        except Exception:
-            body = {"raw": exc.response.text}
-        error_detail = {
-            "status": exc.response.status_code,
-            "body": body,
-            "headers": dict(exc.response.headers),
-        }
-        raise HTTPException(status_code=exc.response.status_code, detail=error_detail) from exc
-    except httpx.RequestError as exc:
-        raise HTTPException(status_code=502, detail=f"Groq request error: {exc}") from exc
-
-
-async def _call_openai_chat(
-    client: httpx.AsyncClient,
-    api_key: str,
-    payload: dict[str, Any],
-) -> tuple[dict[str, Any], dict[str, Any]]:
-    try:
-        response = await client.post(
-            "https://api.openai.com/v1/chat/completions",
-            json=payload,
-            headers={"Authorization": f"Bearer {api_key}"},
-        )
-        response.raise_for_status()
-        data = response.json()
-        return data, {
-            "status": response.status_code,
-            "headers": dict(response.headers),
-            "body": data,
-        }
-    except httpx.HTTPStatusError as exc:
-        try:
-            body = exc.response.json()
-        except Exception:
-            body = {"raw": exc.response.text}
-        error_detail = {
-            "status": exc.response.status_code,
-            "body": body,
-            "headers": dict(exc.response.headers),
-        }
-        try:
-            print("[openai:error]", error_detail, flush=True)
-        except Exception:
-            pass
-        raise HTTPException(status_code=exc.response.status_code, detail=error_detail) from exc
-    except httpx.RequestError as exc:
-        raise HTTPException(status_code=502, detail=f"OpenAI request error: {exc}") from exc
-
-
-async def _rollout_with_groq(
-    request: RolloutRequest,
-    fastapi_request: Request,
-    config: dict[str, Any],
-) -> RolloutResponse:
-    api_key = os.getenv("GROQ_API_KEY")
-    if not api_key:
-        raise HTTPException(
-            status_code=503,
-            detail="GROQ_API_KEY environment variable is required for Groq rollouts.",
-        )
-
-    seed = request.env.seed or 0
-    difficulty = (request.env.config or {}).get("difficulty") or "easy"
-    instance: SokobanTaskInstance = await create_task_instance_from_seed(str(difficulty), int(seed))
-    env = SokobanEnvironment(instance)
-    observation = await env.initialize()
-    initial_observation = observation
-
-    model = config.get("model") or "qwen/qwen3-32b"
-    temperature = float(config.get("temperature", 0.0) or 0.0)
-    top_p = float(config.get("top_p", 0.95) or 0.95)
-    max_tokens = int(config.get("max_tokens", 128) or 128)
-    actions_per_call = int(config.get("max_actions_per_call", 4) or 4)
-    actions_per_call = max(1, min(8, actions_per_call))
-
-    max_steps = int((request.env.config or {}).get("max_steps") or 50)
-
-    steps: List[RolloutStep] = []
-    last_actions: list[int] = []
-    total_reward = float(observation.get("total_reward") or 0.0)
-    executed = 0
-
-    tool_items_enum = sorted(set(ACTION_TOKEN_TO_ID.keys()))
-    tool_schema = {
-        "type": "function",
-        "function": {
-            "name": "interact_many",
-            "description": "Execute a short sequence of Sokoban moves in order.",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "actions": {
-                        "type": "array",
-                        "items": {"type": "string", "enum": tool_items_enum},
-                        "minItems": 1,
-                        "maxItems": actions_per_call,
-                    }
-                },
-                "required": ["actions"],
-                "additionalProperties": False,
-            },
-        },
-    }
-
-    async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
-        for _ in range(max_steps):
-            user_prompt = _format_sokoban_prompt(observation, last_actions)
-            messages = [
-                {"role": "system", "content": SOKOBAN_SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
-            ]
-            payload = {
-                "model": model,
-                "messages": messages,
-                "temperature": temperature,
-                "top_p": top_p,
-                "max_tokens": max_tokens,
-                "tools": [tool_schema],
-                "tool_choice": {"type": "function", "function": {"name": "interact_many"}},
-            }
-            vendor_attempts: list[dict[str, Any]] = []
-            try:
-                response, response_meta = await _call_groq_chat(client, api_key, payload)
-                vendor_attempts.append({"request": payload, "response": response_meta})
-            except HTTPException as exc:
-                detail = exc.detail
-                if isinstance(detail, dict):
-                    vendor_attempts.append({"request": payload, "error": detail})
-                else:
-                    vendor_attempts.append({"request": payload, "error": {"message": str(detail)}})
-                raise
-
-            actions = _extract_actions_from_response(response, actions_per_call)
-            if not actions:
-                break
-
-            aggregated_actions: list[int] = []
-            aggregated_reward = 0.0
-            done = False
-            truncated = False
-            intermediate_rewards: list[float] = []
-            if executed >= max_steps:
-                break
-
-            for action in actions:
-                if executed >= max_steps:
-                    break
-                aggregated_actions.append(int(action))
-                observation = await env.step(
-                    EnvToolCall(tool="interact", args={"action": int(action)})
-                )
-                current_total = float(observation.get("total_reward") or total_reward)
-                reward_delta = current_total - total_reward
-                total_reward = current_total
-                aggregated_reward += reward_delta
-                intermediate_rewards.append(reward_delta)
-                done = bool(observation.get("terminated"))
-                truncated = bool(observation.get("truncated"))
-                executed += 1
-                if done or truncated:
-                    break
-
-            if not aggregated_actions:
-                continue
-
-            last_actions = aggregated_actions
-            step = RolloutStep(
-                obs=observation,
-                tool_calls=[
-                    {
-                        "tool": "interact_many",
-                        "args": {"actions": [int(a) for a in aggregated_actions]},
-                        "source": "groq",
-                    }
-                ],
-                reward=aggregated_reward,
-                done=done,
-                truncated=truncated if truncated else None,
-                info={
-                    "provider": "groq",
-                    "model": model,
-                    "actions_executed": aggregated_actions,
-                    "prompt": user_prompt,
-                    "reward_deltas": intermediate_rewards,
-                    "vendor_attempts": vendor_attempts,
-                    "groq_attempts": vendor_attempts,
-                },
-            )
-            steps.append(step)
-
-            if step.done or (step.truncated or False):
-                break
-
-    final = {"observation": observation, "reward": total_reward}
-    inference_url_groq = "https://api.groq.com/openai/v1/chat/completions"
-    
-    trajectory = RolloutTrajectory(
-        env_id=request.env.env_id or request.env.env_name or "sokoban",
-        policy_id=request.policy.policy_id or request.policy.policy_name or "sokoban-groq",
-        steps=steps,
-        final=final,
-        length=len(steps),
-        inference_url=inference_url_groq,  # NEW: Required for trace correlation
-    )
-    metrics = RolloutMetrics(
-        episode_returns=[total_reward],
-        mean_return=total_reward if steps else 0.0,
-        num_steps=len(steps),
-        num_episodes=1,
-        outcome_score=None,
-        events_score=None,
-        details={"provider": "groq", "model": model},
-    )
-    trace_payload = _build_trace_payload(
-        request,
-        steps,
-        metrics,
-        difficulty=str(difficulty),
-        initial_observation=initial_observation,
-        provider="groq",
-    )
-    return RolloutResponse(
-        run_id=request.run_id,
-        trajectories=[trajectory],
-        branches={},
-        metrics=metrics,
-        aborted=False,
-        ops_executed=executed,
-        trace=trace_payload,
-    )
-
-
-async def _rollout_with_openai(
-    request: RolloutRequest,
-    fastapi_request: Request,
-    config: dict[str, Any],
-) -> RolloutResponse:
-    api_key = os.getenv("OPENAI_API_KEY")
-    if not api_key:
-        raise HTTPException(
-            status_code=503,
-            detail="OPENAI_API_KEY environment variable is required for OpenAI rollouts.",
-        )
-
-    seed = request.env.seed or 0
-    difficulty = (request.env.config or {}).get("difficulty") or "easy"
-    instance: SokobanTaskInstance = await create_task_instance_from_seed(str(difficulty), int(seed))
-    env = SokobanEnvironment(instance)
-    observation = await env.initialize()
-    initial_observation = observation
-
-    model = config.get("model") or "gpt-5"
-    temperature_cfg = config.get("temperature")
-    top_p_cfg = config.get("top_p")
-    completion_tokens = int(
-        config.get("max_completion_tokens")
-        or config.get("max_tokens")
-        or 4000
-    )
-    actions_per_call = int(config.get("max_actions_per_call", 4) or 4)
-    actions_per_call = max(1, min(8, actions_per_call))
-
-    max_steps = int((request.env.config or {}).get("max_steps") or 50)
-
-    steps: List[RolloutStep] = []
-    last_actions: list[int] = []
-    total_reward = float(observation.get("total_reward") or 0.0)
-    executed = 0
-
-    tool_items_enum = sorted(set(ACTION_TOKEN_TO_ID.keys()))
-    tool_schema = {
-        "type": "function",
-        "function": {
-            "name": "interact_many",
-            "description": "Execute a short sequence of Sokoban moves in order.",
-            "parameters": {
-                "type": "object",
-                "properties": {
-                    "actions": {
-                        "type": "array",
-                        "items": {"type": "string", "enum": tool_items_enum},
-                        "minItems": 1,
-                        "maxItems": actions_per_call,
-                    }
-                },
-                "required": ["actions"],
-                "additionalProperties": False,
-            },
-        },
-    }
-
-    async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client:
-        # Process ops array - each "policy" op triggers one LLM call
-        ops_to_process = request.ops or []
-        if not ops_to_process:
-            # If no ops provided, default to max_steps policy calls
-            ops_to_process = ["policy"] * max_steps
-        
-        for op_idx, op in enumerate(ops_to_process):
-            # Only process "policy" ops, skip explicit actions for now
-            if op != "policy" and not (isinstance(op, str) and op.lower() == "policy"):
-                continue
-                
-            user_prompt = _format_sokoban_prompt(observation, last_actions)
-            messages = [
-                {"role": "system", "content": SOKOBAN_SYSTEM_PROMPT},
-                {"role": "user", "content": user_prompt},
-            ]
-            payload_base: dict[str, Any] = {
-                "model": model,
-                "messages": messages,
-                "max_completion_tokens": completion_tokens,
-                "tools": [tool_schema],
-                "tool_choice": {"type": "function", "function": {"name": "interact_many"}},
-            }
-            # GPT-5 models don't support temperature/top_p (only default value of 1)
-            is_gpt5 = "gpt-5" in model.lower()
-            if temperature_cfg is not None and not is_gpt5:
-                with contextlib.suppress(Exception):
-                    payload_base["temperature"] = float(temperature_cfg)
-            if top_p_cfg is not None and not is_gpt5:
-                with contextlib.suppress(Exception):
-                    payload_base["top_p"] = float(top_p_cfg)
-
-            vendor_attempts: list[dict[str, Any]] = []
-            attempt_payload = dict(payload_base)
-            while True:
-                attempt_record: dict[str, Any] = {"request": dict(attempt_payload)}
-                try:
-                    response, response_meta = await _call_openai_chat(client, api_key, attempt_payload)
-                    attempt_record["response"] = response_meta
-                    vendor_attempts.append(attempt_record)
-                    break
-                except HTTPException as exc:
-                    detail = exc.detail
-                    attempt_record["error"] = detail if isinstance(detail, dict) else {"message": str(detail)}
-                    vendor_attempts.append(attempt_record)
-                    handled = False
-                    body = detail.get("body") if isinstance(detail, dict) else None
-                    error_info = body.get("error") if isinstance(body, dict) else None
-                    code = error_info.get("code") if isinstance(error_info, dict) else None
-                    param = error_info.get("param") if isinstance(error_info, dict) else None
-                    if code in {"unsupported_parameter", "unsupported_value"}:
-                        if param == "temperature" and "temperature" in attempt_payload:
-                            attempt_payload = dict(attempt_payload)
-                            attempt_payload.pop("temperature", None)
-                            handled = True
-                        elif param == "top_p" and "top_p" in attempt_payload:
-                            attempt_payload = dict(attempt_payload)
-                            attempt_payload.pop("top_p", None)
-                            handled = True
-                    if handled:
-                        continue
-                    raise
-
-            actions = _extract_actions_from_response(response, actions_per_call)
-            if not actions:
-                break
-
-            aggregated_actions: list[int] = []
-            aggregated_reward = 0.0
-            done = False
-            truncated = False
-            intermediate_rewards: list[float] = []
-            if executed >= max_steps:
-                break
-
-            print(f"[debug] Processing {len(actions)} actions from LLM", flush=True)
-            for action in actions:
-                if executed >= max_steps:
-                    break
-                aggregated_actions.append(int(action))
-                observation = await env.step(
-                    EnvToolCall(tool="interact", args={"action": int(action)})
-                )
-                current_total = float(observation.get("total_reward") or total_reward)
-                reward_delta = current_total - total_reward
-                total_reward = current_total
-                aggregated_reward += reward_delta
-                intermediate_rewards.append(reward_delta)
-                done = bool(observation.get("terminated"))
-                truncated = bool(observation.get("truncated"))
-                executed += 1
-                if done or truncated:
-                    break
-
-                print(f"[debug] After action {action}: done={done}, trunc={truncated}, exec={executed}", flush=True)
-            if not aggregated_actions:
-                continue
-
-            last_actions = aggregated_actions
-            step = RolloutStep(
-                obs=observation,
-                tool_calls=[
-                    {
-                        "tool": "interact_many",
-                        "args": {"actions": [int(a) for a in aggregated_actions]},
-                        "source": "openai",
-                    }
-                ],
-                reward=aggregated_reward,
-                done=done,
-                truncated=truncated if truncated else None,
-                info={
-                    "provider": "openai",
-                    "model": model,
-                    "actions_executed": aggregated_actions,
-                    "prompt": user_prompt,
-                    "reward_deltas": intermediate_rewards,
-                    "vendor_attempts": vendor_attempts,
-                    "openai_attempts": vendor_attempts,
-                    "max_completion_tokens": completion_tokens,
-                    "temperature_requested": temperature_cfg,
-                    "top_p_requested": top_p_cfg,
-                },
-            )
-            steps.append(step)
-
-            if step.done or (step.truncated or False):
-                break
-
-    final = {"observation": observation, "reward": total_reward}
-    inference_url_openai = "https://api.openai.com/v1/chat/completions"
-    
-    trajectory = RolloutTrajectory(
-        env_id=request.env.env_id or request.env.env_name or "sokoban",
-        policy_id=request.policy.policy_id or request.policy.policy_name or "sokoban-openai",
-        steps=steps,
-        final=final,
-        length=len(steps),
-        inference_url=inference_url_openai,  # NEW: Required for trace correlation
-    )
-    metrics = RolloutMetrics(
-        episode_returns=[total_reward],
-        mean_return=total_reward if steps else 0.0,
-        num_steps=len(steps),
-        num_episodes=1,
-        outcome_score=None,
-        events_score=None,
-        details={"provider": "openai", "model": model},
-    )
-    trace_payload = _build_trace_payload(
-        request,
-        steps,
-        metrics,
-        difficulty=str(difficulty),
-        initial_observation=initial_observation,
-        provider="openai",
-    )
-    return RolloutResponse(
-        run_id=request.run_id,
-        trajectories=[trajectory],
-        branches={},
-        metrics=metrics,
-        aborted=False,
-        ops_executed=executed,
-        trace=trace_payload,
-    )
-
-
-def build_config() -> TaskAppConfig:
-    taskset = SokobanTaskSet()
-    base = _task_info()
-    app_state: dict[str, Any] = {"sokoban_taskset": taskset, "sokoban_envs": {}}
-    config = TaskAppConfig(
-        app_id="sokoban",
-        name="Sokoban Task App",
-        description="Sokoban environment exposed as a Synth task app.",
-        base_task_info=base,
-        describe_taskset=lambda: {"id": "sokoban", "name": "Sokoban"},
-        provide_task_instances=lambda seeds: taskset.provide_task_instances(seeds),
-        rollout=rollout_executor,
-        dataset_registry=None,
-        rubrics=None,
-        proxy=None,
-        routers=(router,),
-        app_state=app_state,
-        cors_origins=["*"],
-    )
-    return config
-
-
-# --- Health routes (auth-tolerant) ---
-def fastapi_app():
-    app = create_task_app(build_config())
-
-    # Replace default health handlers to log expected ENVIRONMENT_API_KEY when unauthorized
-    filtered_routes = []
-    for route in app.router.routes:
-        path = getattr(route, "path", None)
-        methods = getattr(route, "methods", set()) or set()
-        if path in {"/health", "/health/rollout"} and "GET" in methods:
-            continue
-        filtered_routes.append(route)
-    app.router.routes = filtered_routes
-
-    def _key_prefix() -> Optional[str]:
-        key = normalize_environment_api_key()
-        return key[: max(1, len(key) // 2)] if key else None
-
-    @app.get("/health")
-    async def health(request: Request):
-        env_key = normalize_environment_api_key()
-        if not env_key:
-            return JSONResponse(status_code=503, content={"status": "unhealthy", "detail": "Missing ENVIRONMENT_API_KEY"})
-        if not is_api_key_header_authorized(request):
-            content: Dict[str, Any] = {"status": "healthy", "authorized": False}
-            prefix = _key_prefix()
-            if prefix:
-                content["expected_api_key_prefix"] = prefix
-            return JSONResponse(status_code=200, content=content)
-        return {"status": "healthy", "authorized": True}
-
-    @app.get("/health/rollout")
-    async def health_rollout(request: Request):
-        env_key = normalize_environment_api_key()
-        if not env_key:
-            return JSONResponse(status_code=503, content={"status": "unhealthy", "detail": "Missing ENVIRONMENT_API_KEY"})
-        if not is_api_key_header_authorized(request):
-            content: Dict[str, Any] = {"status": "healthy", "authorized": False}
-            prefix = _key_prefix()
-            if prefix:
-                content["expected_api_key_prefix"] = prefix
-            return JSONResponse(status_code=200, content=content)
-        return {"ok": True, "authorized": True}
-
-    # Basic env lifecycle routes (for local eval only)
-    @app.post("/env/sokoban/initialize")
-    async def initialize_env(request: Request, payload: Dict[str, Any]):
-        difficulty = str((payload.get("config") or {}).get("difficulty") or "easy")
-        seed = payload.get("seed")
-        try:
-            instance: SokobanTaskInstance = await create_task_instance_from_seed(difficulty, int(seed) if seed is not None else 0)
-        except Exception as exc:
-            raise HTTPException(status_code=400, detail=str(exc))
-        env = SokobanEnvironment(instance)
-        obs = await env.initialize()
-        envs: Dict[str, SokobanEnvironment] = request.app.state.sokoban_envs
-        env_id = f"{difficulty}:{seed or 0}"
-        envs[env_id] = env
-        return {"env_id": env_id, "observation": obs}
-
-    @app.post("/env/sokoban/step")
-    async def step_env(request: Request, payload: Dict[str, Any]):
-        env_id = str(payload.get("env_id") or "")
-        if not env_id:
-            raise HTTPException(status_code=400, detail="env_id required")
-        envs: Dict[str, SokobanEnvironment] = request.app.state.sokoban_envs
-        env = envs.get(env_id)
-        if not env:
-            raise HTTPException(status_code=404, detail="Unknown env_id")
-
-        action = None
-        tool_calls = payload.get("tool_calls") or []
-        if tool_calls:
-            try:
-                first = tool_calls[0] or {}
-                args = first.get("args") or {}
-                action = int(args.get("action")) if "action" in args else None
-            except Exception:
-                action = None
-        if action is None and "action" in payload:
-            try:
-                action = int(payload.get("action"))
-            except Exception:
-                action = None
-        if action is None:
-            raise HTTPException(status_code=400, detail="action required")
-        obs = await env.step(EnvToolCall(tool="interact", args={"action": int(action)}))
-        return {"observation": obs}
-
-    @app.post("/env/sokoban/terminate")
-    async def terminate_env(request: Request, payload: Dict[str, Any]):
-        env_id = str(payload.get("env_id") or "")
-        envs: Dict[str, SokobanEnvironment] = request.app.state.sokoban_envs
-        env = envs.pop(env_id, None)
-        if env:
-            obs = await env.terminate()
-        else:
-            obs = {"terminated": True}
-        return {"ok": True, "observation": obs}
-
-    @app.exception_handler(RequestValidationError)
-    async def _on_validation_error(_request: Request, exc: RequestValidationError):
-        return JSONResponse(status_code=422, content={"status": "invalid", "detail": exc.errors()[:5]})
-
-    return app
-
-
-register_task_app(
-    entry=TaskAppEntry(
-        app_id="sokoban",
-        description="Sokoban task app",
-        config_factory=build_config,
-        aliases=("sokoban-rl",),
-        env_files=(),
-        modal=None,
-    )
-)