playground/server.py

#!/usr/bin/env python3
"""Trace2Env playground: a local web page where you play the task agent and the harness plays the environment.

    export OPENROUTER_API_KEY=...
    PYTHONPATH=src python docs/playground/server.py --package PACKAGE --session work/playground/terminal \
        [--allow-unvalidated] [--initial-state @state.json] [--environment-prompt @prompt.md] [--port 8765]

Then open http://127.0.0.1:8765, type an action, and read the simulated observation. Every step goes through
``RuntimeHarness.step`` exactly as in ``trace2env simulate``: rule selection, the world-model agent's tool loop over
the worldbook (the package), verification, rendering, and an atomic commit to a persistent session.

Each visitor (a browser cookie) has their own episode, action budget and, optionally, their own model credentials:
a visitor may send ``--max-actions-per-visitor`` actions on the server's key, or any number on a key they provide,
which is held in this process's memory only (never written to disk, logs or audits) and forgotten on request or at
restart. The default backbone is deepseek/deepseek-v4.1-flash through OpenRouter (provider pinned to CoreWeave,
reasoning effort low, native tool calling); pass --model / --base-url / --api-key-env / --provider for another one,
or --no-model for rules only. Only the JSON API under /api/ and the static files are served; the server binds to
localhost by default. Put it behind HTTPS if you expose it, since visitors may type keys into it.
"""
from __future__ import annotations

import argparse
import json
import re
import secrets
import sys
import threading
import time
import traceback
from dataclasses import dataclass, field
from datetime import datetime, timezone
from http import HTTPStatus
from http.cookies import SimpleCookie
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from typing import Any

HERE = Path(__file__).resolve().parent
REPO = HERE.parents[1]
if str(REPO / "src") not in sys.path:  # run from a checkout without installing
    sys.path.insert(0, str(REPO / "src"))

from trace2env.agent import AgentProtocolError  # noqa: E402
from trace2env.agentworld import _from_call_text, _from_json_action, terminal_keystrokes_action  # noqa: E402
from trace2env.cli import (_add_agent_options, _add_model_options, _agent_llm, _chat_reasoning, _chat_reasoning_json,  # noqa: E402
                           _json_argument, _llm, _resolved_features, _text_argument, _trace_corpus)
from trace2env.engine import AmbiguousRuleError  # noqa: E402
from trace2env.llm import CachedAgentLLM, OpenAIChatStructuredLLM, OpenAIToolLLM, StructuredAgentLLM  # noqa: E402
from trace2env.models import EnvironmentState, NormalizedAction, StepRequest  # noqa: E402
from trace2env.package import EnvironmentPackage  # noqa: E402
from trace2env.runtime import InvalidAction, NoApplicableRule, RuntimeHarness, initial_state_from_package  # noqa: E402

STATIC = {"/": ("index.html", "text/html; charset=utf-8"), "/app.js": ("app.js", "text/javascript; charset=utf-8"),
          "/style.css": ("style.css", "text/css; charset=utf-8")}
INTERNAL_TOOLS = {"_usage", "_knowledge_gate"}
VISITOR_COOKIE = "t2e_visitor"
VISITOR_HEADER = "X-T2E-Visitor"
DEFAULT_TITLE = "Playground - Trace2Env (Agentic Language World Model)"
GITHUB_URL = "https://github.com/ruyue0001/trace2env"


# ----- text helpers ----------------------------------------------------------------------------------------------
def clip(value: Any, limit: int = 1500) -> Any:
    """Recursively shorten long strings so an API payload stays readable in the browser."""
    if isinstance(value, str):
        return value if len(value) <= limit else value[:limit] + f"... [{len(value) - limit} more characters]"
    if isinstance(value, dict):
        return {key: clip(item, limit) for key, item in value.items()}
    if isinstance(value, list):
        return [clip(item, limit) for item in value]
    return value


def flatten(value: Any, prefix: str = "") -> dict[str, Any]:
    if isinstance(value, dict) and value:
        out: dict[str, Any] = {}
        for key, item in value.items():
            out.update(flatten(item, f"{prefix}.{key}" if prefix else str(key)))
        return out
    return {prefix: value}


def state_diff(before: EnvironmentState, after: EnvironmentState) -> list[dict[str, Any]]:
    """Changed state paths between two states (revision and step excluded), values clipped for display."""
    old = flatten(before.model_dump(mode="json"))
    new = flatten(after.model_dump(mode="json"))
    changes = []
    for path in sorted(set(old) | set(new)):
        if path in ("revision", "step") or old.get(path, "\0missing") == new.get(path, "\0missing"):
            continue
        changes.append({"path": path, "kind": "added" if path not in old else "removed" if path not in new else "changed",
                        "before": clip(old.get(path), 300), "after": clip(new.get(path), 300)})
    return changes


# ----- typed input -> action -------------------------------------------------------------------------------------
KEY_NAMES = {"enter", "escape", "esc", "tab", "up", "down", "left", "right", "space", "bspace", "backspace", "delete", "home",
             "end", "pageup", "pagedown", "pgup", "pgdn"}
_KEY_CHORD = re.compile(r"^(?:[CMS]-)+[A-Za-z0-9\[\]\\^_]$")  # tmux-style chords: C-c, M-x, C-M-d
_WAIT = re.compile(r"^wait(?:\s+(\d+(?:\.\d+)?))?$", re.IGNORECASE)
_TOOL_CALL = re.compile(r"^([A-Za-z_][\w.\-]*)\s*(.*)$", re.S)
# How long the terminal is observed after typed keystrokes (the benchmark's per-action "duration"). A person typing a
# command expects to see its output, so the window is chosen for them: generous for ordinary commands, longer for
# installs, builds and downloads. A command that outlives its window is reported as partial and keeps the terminal
# until a later ``wait``. Plays and API callers may still pass an explicit duration.
DEFAULT_CAPTURE_S = 30.0
LONG_CAPTURE_S = 120.0
_LONG_RUNNING = re.compile(r"(^|[\s;&|(])(apt(-get)?|apt-cache|dpkg|yum|dnf|apk|brew|pip3?|python3? -m pip|uv|poetry|conda|npm|npx|yarn|pnpm|"
                           r"cargo|go (get|build|test)|mvn|gradle|make|cmake|gcc|g\+\+|clang|rustc|javac|wget|curl|git (clone|fetch|pull|push)|"
                           r"docker|tar|unzip|rsync|scp|pytest)\b")


def auto_capture(text: str) -> float:
    return LONG_CAPTURE_S if _LONG_RUNNING.search(text) else DEFAULT_CAPTURE_S


def input_mode_of(package: EnvironmentPackage) -> str:
    """``terminal`` when the worldbook models a shell (it has the generic shell/keys/wait actions), else ``tool``."""
    names = {action.name for action in package.action_schema.actions}
    return "terminal" if names & {"shell", "keys", "wait"} else "tool"


def _parse_scalar(text: str) -> Any:
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        return text


def parse_text(text: str, mode: str, duration: float | None = None) -> NormalizedAction:
    """Turn what a person types into the action the worldbook expects.

    Any mode: a JSON object with ``type`` is an action; one with ``name``/``arguments`` is a tool call.
    ``terminal``: the text is what you would type at the prompt. A line gets its newline and becomes one action per
    simple command (``rm report.txt`` -> type ``rm``), ``shell`` for pipelines, ``&&``, heredocs and multi-line input;
    a tmux-style chord (``C-c``) or key name (``Enter``, ``Escape``) is sent as ``keys``; ``wait`` or ``wait 5`` waits.
    ``duration`` is the capture window in seconds after the keystrokes; when absent it is chosen from the command
    (``auto_capture``): the observation is the screen at the end of it, so a command that runs longer needs a ``wait``.
    ``tool``: ``name(key=value, ...)``, ``name {"key": value}``, ``name key=value key2=value2`` or a bare ``name``.
    """
    body = text.strip("\n")
    if not body.strip():
        raise ValueError("type a command, a key chord such as C-c, or 'wait'")
    stripped = body.strip()
    if stripped.startswith("{"):
        parsed = json.loads(stripped)
        if not isinstance(parsed, dict):
            raise ValueError("a JSON action must be an object")
        if "type" in parsed:
            return NormalizedAction.model_validate({"type": parsed["type"], "arguments": parsed.get("arguments") or {}, "raw": parsed})
        action = _from_json_action("playground", parsed)
        if action is None:
            raise ValueError('JSON input needs "type" or "name" and "arguments"')
        return action
    if mode == "terminal":
        wait = _WAIT.match(stripped)
        if wait:
            return terminal_keystrokes_action([{"keystrokes": "", "duration": float(wait.group(1) or 5.0)}])
        explicit = float(duration) if duration is not None and float(duration) > 0 else None
        if _KEY_CHORD.match(stripped) or stripped.lower() in KEY_NAMES:
            return terminal_keystrokes_action([{"keystrokes": stripped, "duration": explicit if explicit is not None else 0.5}])
        keystrokes = body if body.endswith("\n") else body + "\n"
        return terminal_keystrokes_action([{"keystrokes": keystrokes, "duration": explicit if explicit is not None else auto_capture(stripped)}])
    match = _TOOL_CALL.match(stripped)
    if not match:
        raise ValueError("a tool call looks like name(key=value), name {\"key\": value}, name key=value, or just name")
    name, rest = match.group(1), match.group(2).strip()
    if rest.startswith("(") and rest.endswith(")"):
        action = _from_call_text(stripped)
        if action is not None:
            return action
        raise ValueError(f"could not parse the arguments of {stripped!r}")
    if rest.startswith("{"):
        arguments = json.loads(rest)
        if not isinstance(arguments, dict):
            raise ValueError("the arguments after the tool name must be a JSON object")
        return NormalizedAction(type=name, arguments=arguments, raw=stripped)
    arguments: dict[str, Any] = {}
    for token in re.findall(r'(\w[\w.\-]*)=("(?:[^"\\]|\\.)*"|\'(?:[^\'\\]|\\.)*\'|\S+)', rest):
        key, value = token
        if value[:1] in "\"'" and value[-1:] == value[:1]:
            value = value[1:-1]
        arguments[key] = _parse_scalar(value)
    if rest and not arguments:
        raise ValueError(f"could not parse {rest!r}: use key=value pairs, (key=value), or a JSON object")
    return NormalizedAction(type=name, arguments=arguments, raw=stripped)


def load_examples(path: Path) -> list[dict[str, Any]]:
    """Suggested plays: each has an id, mode, title, description, optional initial_state and package_match, and steps."""
    data = json.loads(path.read_text(encoding="utf-8"))
    examples = data.get("examples") if isinstance(data, dict) else data
    if not isinstance(examples, list):
        raise ValueError(f"{path}: expected a list of examples or an object with an 'examples' list")
    for example in examples:
        for key in ("id", "mode", "title", "steps"):
            if key not in example:
                raise ValueError(f"{path}: example {example.get('id', '?')!r} lacks {key!r}")
        if example.get("initial_state") is not None:
            EnvironmentState.model_validate(example["initial_state"])  # fail at startup, not when a play starts
    return examples


class Busy(RuntimeError):
    pass


class BudgetExhausted(RuntimeError):
    pass


_MODEL_NAME = re.compile(r"^[\w][\w.:/\-]{0,127}$")


@dataclass
class Visitor:
    """One browser's playground: its episode, its action count, and optionally its own model credentials."""
    id: str
    harness: RuntimeHarness
    episode_dir: Path
    current_example: dict[str, Any] | None = None
    free_actions: int = 0  # steps taken on the server's key (count against the free caps)
    own_actions: int = 0  # steps taken on the visitor's own key (never touch the free balance)
    credentials: dict[str, str] | None = None  # model, api_key, base_url: process memory only
    busy: bool = False
    lock: threading.Lock = field(default_factory=threading.Lock)
    created: float = field(default_factory=time.time)
    last_seen: float = field(default_factory=time.time)


class Playground:
    def __init__(self, args: argparse.Namespace):
        self.args = args
        self.package = EnvironmentPackage(args.package)
        self.environment_prompt = _text_argument(args.environment_prompt) if args.environment_prompt else ""
        self.initial_state, self.initial_state_source = self._initial_state()
        # --session is the root under which every episode gets its own directory; a visit never resumes an old one.
        self.session_root = Path(args.session)
        self.input_mode = input_mode_of(self.package)
        self.examples = load_examples(Path(args.examples)) if args.examples else load_examples(HERE / "examples.json")
        self.title = args.title or DEFAULT_TITLE
        self.tagline = args.tagline or self._default_tagline()
        self.max_actions_per_visitor = max(0, args.max_actions_per_visitor)
        self.max_actions_own_key = max(0, args.max_actions_own_key)
        self.max_actions_total = max(0, args.max_actions_total)
        self.total_actions = 0
        self.visitors: dict[str, Visitor] = {}
        self.registry_lock = threading.Lock()
        # Action counters are mirrored to <session root>/budget.json after every charged action and read back at
        # start-up, so restarting the server over the same session root does not refill anyone's allowance
        # (visitor ids and counts only; never credentials).
        self.budget_file = self.session_root / "budget.json"
        self.saved_counters: dict[str, dict[str, int]] = {}
        self._load_budget()
        self.steps = threading.Semaphore(max(1, args.max_concurrent_steps))
        self.server_llm_cache: dict[str, Any] = {}

    # ----- construction -------------------------------------------------------------------------------------
    def _load_budget(self) -> None:
        try:
            data = json.loads(self.budget_file.read_text(encoding="utf-8"))
        except (OSError, ValueError):
            return
        if not isinstance(data, dict):
            return
        total = data.get("total_actions")
        self.total_actions = total if isinstance(total, int) and total >= 0 else 0
        visitors = data.get("visitors")
        if isinstance(visitors, dict):
            for visitor_id, counts in visitors.items():
                if isinstance(visitor_id, str) and isinstance(counts, dict):
                    self.saved_counters[visitor_id] = {key: max(0, int(counts.get(key) or 0)) for key in ("free_actions", "own_actions")}

    def _save_budget(self) -> None:
        """Called with the registry lock held. Best effort: a disk problem is logged, never raised into a step."""
        visitors = dict(self.saved_counters)
        visitors.update({v.id: {"free_actions": v.free_actions, "own_actions": v.own_actions} for v in self.visitors.values()})
        try:
            self.session_root.mkdir(parents=True, exist_ok=True)
            temporary = self.budget_file.with_name("budget.json.tmp")
            temporary.write_text(json.dumps({"total_actions": self.total_actions, "visitors": visitors}, sort_keys=True), encoding="utf-8")
            temporary.replace(self.budget_file)
        except OSError as exc:
            sys.stderr.write(f"[playground] could not save the action counters: {exc}\n")

    def _default_tagline(self) -> str:
        manifest = self.package.manifest.model_dump(mode="json")
        name = str(manifest.get("name") or self.package.root.name)
        count = manifest.get("source_count")
        traces = f"{count} " if isinstance(count, int) and count > 0 else ""
        if "terminal-bench" in name.lower():
            return (f"A simulated terminal environment. The worldbook is built from {traces}Terminal-Bench 2.0 traces. "
                    f"GitHub: {GITHUB_URL}")
        if "ledger" in name.lower():
            return f"A simulated ledger. The worldbook is hand-authored. GitHub: {GITHUB_URL}"
        return f"A simulated environment. The worldbook is built from {traces}traces of {name}. GitHub: {GITHUB_URL}"

    def _initial_state(self) -> tuple[EnvironmentState, str]:
        if self.args.initial_state:
            return EnvironmentState.model_validate(_json_argument(self.args.initial_state)), "--initial-state"
        if self.package.manifest.metadata.get("construction_kind") == "authored":
            return initial_state_from_package(self.package), "schema defaults (authored package)"
        # A reconstructed package never starts from schema defaults: the episode's state is unknown until observed.
        return EnvironmentState(), "empty (nothing observed yet; pass --initial-state to seed the episode)"

    def _episode_dir(self, label: str) -> Path:
        stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
        candidate = self.session_root / f"{label}-{stamp}"
        n = 1
        while candidate.exists():  # two episodes within one second
            n += 1
            candidate = self.session_root / f"{label}-{stamp}-{n}"
        return candidate

    def _agent_llm(self, credentials: dict[str, str] | None):
        """The world-model agent's model: the server's (key from the environment) or a visitor's own credentials."""
        a = self.args
        if credentials is None:
            if a.no_model or a.prediction_mode == "single_shot":
                return None
            if "agent" not in self.server_llm_cache:
                self.server_llm_cache["agent"] = _agent_llm(a)
            return self.server_llm_cache["agent"]
        model, key = credentials["model"], credentials["api_key"]
        base_url = credentials.get("base_url") or a.base_url
        # The provider pin belongs to the server's default model; another model may not be served by that provider.
        routing = json.loads(a.chat_provider_json) if (a.chat_provider_json and model == a.model) else None
        if a.agent_transport == "tools":
            inner = OpenAIToolLLM(model=model, base_url=base_url, api_key=key, temperature=a.structured_temperature,
                                  reasoning_effort=_chat_reasoning(a), reasoning=_chat_reasoning_json(a), provider_routing=routing,
                                  **({"max_tokens": a.agent_max_output_tokens} if a.agent_max_output_tokens else {}))
        else:
            inner = StructuredAgentLLM(OpenAIChatStructuredLLM(model=model, base_url=base_url, api_key=key,
                                                               temperature=a.structured_temperature, reasoning_effort=_chat_reasoning(a),
                                                               reasoning=_chat_reasoning_json(a), provider_routing=routing,
                                                               max_tokens=a.agent_max_output_tokens or 65536, schema_mode=a.chat_schema_mode))
        # The cache stores model answers keyed by their inputs and model, never the key.
        return CachedAgentLLM(inner, a.cache_dir) if a.cache_dir else inner

    def _build(self, session_dir: Path, initial_state: EnvironmentState | None = None,
               credentials: dict[str, str] | None = None) -> RuntimeHarness:
        a = self.args
        initial_state = self.initial_state if initial_state is None else initial_state
        single_shot = a.prediction_mode == "single_shot"
        single_shot_llm = _llm(a, max_tokens=a.agent_max_output_tokens) if (not a.no_model and single_shot and credentials is None) else None
        return RuntimeHarness(
            str(self.package.root), str(session_dir), package=self.package,
            agent_llm=self._agent_llm(credentials), single_shot_llm=single_shot_llm, prediction_mode=a.prediction_mode,
            trace_corpus=_trace_corpus(a), state_tracking=a.state_tracking, package_knowledge=a.package_knowledge,
            knowledge_tools=a.knowledge_tools, allow_unvalidated=a.allow_unvalidated, initial_state=initial_state,
            trust_policy=a.trust_policy or "schema_checked", fast_path=a.fast_path, max_tool_calls=a.max_tool_calls,
            features=_resolved_features(a), history_turn_chars=a.history_turn_chars,
            history_total_chars=a.history_total_chars, knowledge_gate=a.knowledge_gate, knowledge_names=a.knowledge_names,
            effect_rejection=a.effect_rejection, memory_recent=a.history_window,
            environment_prompt_chars=a.environment_prompt_chars if a.environment_prompt_chars > 0 else None,
            require_observation=not a.allow_missing_observation,
        )

    def _metadata(self, context: dict[str, Any] | None) -> dict[str, Any]:
        metadata: dict[str, Any] = {"history_window": self.args.history_window}
        if self.environment_prompt:
            metadata["environment_prompt"] = self.environment_prompt
        for key, value in (context or {}).items():  # other keys reach the agent as context
            if key not in ("environment_prompt", "official_input", "history_window") and value not in (None, ""):
                metadata[str(key)] = value
        return metadata

    # ----- visitors -----------------------------------------------------------------------------------------
    def visitor(self, visitor_id: str) -> Visitor:
        with self.registry_lock:
            visitor = self.visitors.get(visitor_id)
            if visitor is None:
                episode = self._episode_dir("episode")
                saved = self.saved_counters.get(visitor_id) or {}
                visitor = Visitor(id=visitor_id, harness=self._build(episode), episode_dir=episode,
                                  free_actions=saved.get("free_actions", 0), own_actions=saved.get("own_actions", 0))
                self.visitors[visitor_id] = visitor
            visitor.last_seen = time.time()
            return visitor

    def _claim(self, visitor: Visitor) -> None:
        with visitor.lock:
            if visitor.busy:
                raise Busy("a step is already running for you")
            visitor.busy = True

    def _release(self, visitor: Visitor) -> None:
        with visitor.lock:
            visitor.busy = False

    def budget(self, visitor: Visitor) -> dict[str, Any]:
        own = visitor.credentials is not None
        free_cap, own_cap = self.max_actions_per_visitor, self.max_actions_own_key
        free_left = max(0, free_cap - visitor.free_actions) if free_cap else None
        own_left = max(0, own_cap - visitor.own_actions) if own_cap else None
        return {
            "own_key": own,
            # the counters of the mode the visitor is in now (what the header pill shows) ...
            "per_visitor": (own_cap if own else free_cap) or None,
            "used": visitor.own_actions if own else visitor.free_actions,
            "left": own_left if own else free_left,
            # ... and the free balance, which own-key play never touches
            "free_cap": free_cap or None, "free_used": visitor.free_actions, "free_left": free_left,
            "own_cap": own_cap or None, "own_used": visitor.own_actions, "own_left": own_left,
            "total_cap": self.max_actions_total or None, "total_used": self.total_actions,
            "total_left": (max(0, self.max_actions_total - self.total_actions) if self.max_actions_total else None),
        }

    def charge(self, visitor: Visitor) -> None:
        """Count one action for this visitor in the mode they are in; raise before the step runs when a cap is reached."""
        with self.registry_lock:
            if visitor.credentials is not None:
                if self.max_actions_own_key and visitor.own_actions >= self.max_actions_own_key:
                    raise BudgetExhausted(f"you have used the {self.max_actions_own_key} actions allowed per visitor on your own key")
                visitor.own_actions += 1
                self._save_budget()
                return
            if self.max_actions_total and self.total_actions >= self.max_actions_total:
                raise BudgetExhausted("this playground has used up its free action budget for now; use your own key to continue")
            if self.max_actions_per_visitor and visitor.free_actions >= self.max_actions_per_visitor:
                raise BudgetExhausted(f"you have used your {self.max_actions_per_visitor} free actions; add your own OpenRouter key to keep playing")
            visitor.free_actions += 1
            self.total_actions += 1
            self._save_budget()

    # ----- read-only views ----------------------------------------------------------------------------------
    def info(self, visitor: Visitor) -> dict[str, Any]:
        a = self.args
        manifest = self.package.manifest.model_dump(mode="json")
        metadata = manifest.get("metadata") or {}
        state = visitor.harness.session.load()
        summary = visitor.harness.inspector.summary()
        own = visitor.credentials
        model = (own["model"] if own else (None if a.no_model else a.model))
        return {
            "title": self.title, "tagline": self.tagline,
            "package": {
                "root": str(self.package.root), "name": manifest.get("name"), "environment_id": manifest.get("environment_id"),
                "version": manifest.get("version"), "validation_status": metadata.get("validation_status"),
                "construction_kind": metadata.get("construction_kind"),
                "counts": {k: summary[k] for k in ("state_field_count", "rule_count", "invariant_count", "renderer_count", "note_count")},
            },
            "input_mode": self.input_mode,
            "actions": visitor.harness.inspector.list_actions(),
            "state_fields": [f.model_dump(mode="json") for f in self.package.state_schema.fields],
            "settings": {
                "model": model, "default_model": None if a.no_model else a.model, "own_key": own is not None,
                "provider": None if (a.no_model and not own) else a.provider,
                "agent_transport": None if (a.no_model and not own) else a.agent_transport, "prediction_mode": a.prediction_mode,
                "trust_policy": a.trust_policy or "schema_checked", "fast_path": a.fast_path, "max_tool_calls": a.max_tool_calls,
                "features": sorted(_resolved_features(a)), "knowledge_gate": a.knowledge_gate, "effect_rejection": a.effect_rejection,
                "environment_prompt_chars": len(self.environment_prompt), "history_window": a.history_window,
                "initial_state_source": self.initial_state_source, "base_url": a.base_url,
            },
            "session": {"directory": str(visitor.harness.session.root), "root": str(self.session_root), "step": state.step,
                        "revision": state.revision, "turns": visitor.harness.memory.count()},
            "current_example": visitor.current_example["id"] if visitor.current_example else None,
            "budget": self.budget(visitor),
        }

    def state(self, visitor: Visitor) -> dict[str, Any]:
        return visitor.harness.session.load().model_dump(mode="json")

    def history(self, visitor: Visitor, limit: int = 200) -> list[dict[str, Any]]:
        records = visitor.harness.session.audit_records()[-limit:]
        return [self._audit_view(index, record) for index, record in enumerate(records, start=1)]

    @staticmethod
    def _audit_view(index: int, record: Any) -> dict[str, Any]:
        calls = [clip(call, 600) for call in record.tool_calls if call.get("tool") not in INTERNAL_TOOLS]
        usage = next((call.get("arguments") for call in record.tool_calls if call.get("tool") == "_usage"), None)
        gate = next((call.get("arguments") for call in record.tool_calls if call.get("tool") == "_knowledge_gate"), None)
        plan = record.plan
        return {
            "index": index, "timestamp": record.timestamp.isoformat(), "committed": record.committed, "route": record.route,
            "action": record.request.action.model_dump(mode="json"), "observation": record.observation, "error": record.error,
            "outcome": plan.outcome.value if plan else None, "rule_ids": plan.rule_ids if plan else [],
            "effects": [m.model_dump(mode="json") for m in plan.effects] if plan else [],
            "uncertainty": plan.uncertainty if plan else [], "rationale": plan.rationale if plan else "",
            "citations": record.citations, "retrieved_artifacts": record.retrieved_artifacts,
            "verification": record.verification.model_dump(mode="json") if record.verification else None,
            "agent_trace": calls, "usage": usage, "knowledge_gate": gate,
            "prior_revision": record.prior_revision, "new_revision": record.new_revision,
        }

    def applicable_examples(self) -> list[dict[str, Any]]:
        """Examples for this worldbook: same input mode, and package_match (if given) contained in the package name or path."""
        name = f"{self.package.manifest.model_dump(mode='json').get('name', '')} {self.package.root}"
        return [example for example in self.examples
                if example["mode"] == self.input_mode and (not example.get("package_match") or example["package_match"].lower() in name.lower())]

    def inspect(self, visitor: Visitor, action_type: str) -> dict[str, Any]:
        return clip(visitor.harness.inspector.inspect_action(action_type, state=visitor.harness.session.load()), 2000)

    def search(self, visitor: Visitor, query: str, limit: int = 10) -> list[dict[str, Any]]:
        return clip(visitor.harness.inspector.search_knowledge(query, limit=limit), 800)

    # ----- mutations ----------------------------------------------------------------------------------------
    def step(self, visitor: Visitor, action: dict[str, Any] | None, context: dict[str, Any] | None,
             text: str | None = None, duration: float | None = None) -> dict[str, Any]:
        normalized = parse_text(text, self.input_mode, duration) if text is not None else NormalizedAction.model_validate(action)
        self._claim(visitor)
        try:
            self.charge(visitor)
            request = StepRequest(action=normalized, metadata=self._metadata(context))
            before = visitor.harness.session.load()
            started = time.time()
            retried = False
            with self.steps:
                try:
                    try:
                        result = visitor.harness.step(request)
                    except AgentProtocolError:
                        # A protocol slip of the model (a batch of tool calls past the budget, a reply without a submission)
                        # is not a property of the environment: the failed attempt stays in the audit, state is unchanged,
                        # and the step is simulated once more before the error is shown. The retry marker reaches the agent
                        # as context and makes the brief differ, so the call cache cannot replay the same slip.
                        retried = True
                        request = StepRequest(action=normalized, metadata={**request.metadata, "retry_attempt": 1})
                        result = visitor.harness.step(request)
                except (InvalidAction, NoApplicableRule, AmbiguousRuleError, AgentProtocolError) as exc:
                    records = visitor.harness.session.audit_records()
                    failed = self._audit_view(len(records), records[-1]) if records and not records[-1].committed else None
                    return {"ok": False, "error": {"type": type(exc).__name__, "message": str(exc)}, "audit": failed,
                            "action": normalized.model_dump(mode="json"), "input": text, "retried": retried,
                            "state": before.model_dump(mode="json"), "elapsed_s": round(time.time() - started, 2)}
            records = visitor.harness.session.audit_records()
            audit = self._audit_view(len(records), records[-1])
            return {
                "ok": True, "observation": result.observation, "outcome": result.outcome.value, "route": result.route, "retried": retried,
                "action": normalized.model_dump(mode="json"), "input": text,
                "tool_calls": result.tool_calls, "citations": result.citations, "retrieved_artifacts": result.retrieved_artifacts,
                "plan": {"rule_ids": result.plan.rule_ids, "effects": [m.model_dump(mode="json") for m in result.plan.effects],
                         "uncertainty": result.plan.uncertainty, "rationale": result.plan.rationale,
                         "observation_template": result.plan.observation_template},
                "verification": result.verification.model_dump(mode="json"),
                "state": result.state.model_dump(mode="json"), "state_diff": state_diff(before, result.state),
                "audit": audit, "elapsed_s": round(time.time() - started, 2),
            }
        finally:
            self._release(visitor)

    def reset(self, visitor: Visitor) -> dict[str, Any]:
        self._claim(visitor)
        try:
            visitor.episode_dir = self._episode_dir("episode")
            visitor.harness = self._build(visitor.episode_dir, credentials=visitor.credentials)
            visitor.current_example = None
            return self.info(visitor)
        finally:
            self._release(visitor)

    def start_example(self, visitor: Visitor, example_id: str) -> dict[str, Any]:
        example = next((item for item in self.applicable_examples() if item["id"] == example_id), None)
        if example is None:
            raise KeyError(f"no example {example_id!r} for this worldbook")
        self._claim(visitor)
        try:
            initial = EnvironmentState.model_validate(example["initial_state"]) if example.get("initial_state") else None
            visitor.episode_dir = self._episode_dir(example_id)
            visitor.harness = self._build(visitor.episode_dir, initial, credentials=visitor.credentials)
            visitor.current_example = example
            return {"example": example, "info": self.info(visitor)}
        finally:
            self._release(visitor)

    def set_credentials(self, visitor: Visitor, model: str, api_key: str, base_url: str | None = None) -> dict[str, Any]:
        """Use the visitor's own key and model from now on, over the same episode. Held in memory only."""
        model = (model or "").strip() or self.args.model
        api_key = (api_key or "").strip()
        if not api_key or len(api_key) > 256 or not api_key.isprintable() or any(c.isspace() for c in api_key):
            raise ValueError("the API key looks malformed")
        if not _MODEL_NAME.match(model):
            raise ValueError("the model name may contain letters, digits, '.', ':', '/' and '-'")
        if base_url and not re.match(r"^https?://[\w.\-:]+(/[\w./\-]*)?$", base_url):
            raise ValueError("the base URL must be an http(s) URL")
        self._claim(visitor)
        try:
            credentials = {"model": model, "api_key": api_key, **({"base_url": base_url} if base_url else {})}
            visitor.harness = self._build(visitor.episode_dir, credentials=credentials)  # reopens the same episode
            visitor.credentials = credentials
            return self.info(visitor)
        finally:
            self._release(visitor)

    def forget_credentials(self, visitor: Visitor) -> dict[str, Any]:
        self._claim(visitor)
        try:
            visitor.credentials = None
            visitor.harness = self._build(visitor.episode_dir)
            return self.info(visitor)
        finally:
            self._release(visitor)


class Handler(BaseHTTPRequestHandler):
    playground: Playground  # set by make_server()
    server_version = "Trace2EnvPlayground/1.1"

    def log_message(self, fmt: str, *args: Any) -> None:  # one compact line per request; never a body
        sys.stderr.write("[playground] %s\n" % (fmt % args))

    def _visitor_id(self) -> tuple[str, bool]:
        """The visitor id from the X-T2E-Visitor header (the page keeps it in local storage, which survives where
        third-party cookies are blocked, e.g. inside the Hugging Face iframe), else from the cookie, else a fresh one
        (and whether it must be set on the response)."""
        header = (self.headers.get(VISITOR_HEADER) or "").strip()
        if header.isalnum() and 8 <= len(header) <= 64:
            return header, False
        cookie = SimpleCookie(self.headers.get("Cookie", ""))
        morsel = cookie.get(VISITOR_COOKIE)
        if morsel and morsel.value.isalnum() and 8 <= len(morsel.value) <= 64:
            return morsel.value, False
        return secrets.token_hex(16), True

    def _secure(self) -> bool:
        """Served over HTTPS (directly or behind a proxy such as the Hugging Face Space router)?"""
        return (self.headers.get("X-Forwarded-Proto") or "").lower() == "https"

    def _send(self, status: int, body: bytes, content_type: str, set_cookie: str | None = None) -> None:
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        if set_cookie:
            # Inside a third-party iframe (the Space page) a cookie needs SameSite=None and Secure; localhost keeps Lax.
            attributes = "SameSite=None; Secure" if self._secure() else "SameSite=Lax"
            self.send_header("Set-Cookie", f"{VISITOR_COOKIE}={set_cookie}; Path=/; Max-Age=31536000; {attributes}")
            self.send_header(VISITOR_HEADER, set_cookie)  # the page stores it and sends it back as a header
        self.end_headers()
        self.wfile.write(body)

    def _json(self, status: int, payload: Any, set_cookie: str | None = None) -> None:
        self._send(status, json.dumps(payload, ensure_ascii=False, default=str).encode("utf-8"),
                   "application/json; charset=utf-8", set_cookie)

    def _read_json(self) -> dict[str, Any]:
        length = int(self.headers.get("Content-Length") or 0)
        raw = self.rfile.read(length) if length else b""
        if not raw:
            return {}
        data = json.loads(raw.decode("utf-8"))
        if not isinstance(data, dict):
            raise ValueError("the request body must be a JSON object")
        return data

    def do_GET(self) -> None:  # noqa: N802
        path = self.path.split("?", 1)[0]
        try:
            if path in STATIC:
                name, content_type = STATIC[path]
                self._send(HTTPStatus.OK, (HERE / name).read_bytes(), content_type)
                return
            visitor_id, fresh = self._visitor_id()
            cookie = visitor_id if fresh else None
            if path == "/api/info":
                self._json(HTTPStatus.OK, self.playground.info(self.playground.visitor(visitor_id)), cookie)
            elif path == "/api/state":
                self._json(HTTPStatus.OK, self.playground.state(self.playground.visitor(visitor_id)), cookie)
            elif path == "/api/history":
                self._json(HTTPStatus.OK, self.playground.history(self.playground.visitor(visitor_id)), cookie)
            elif path == "/api/examples":
                self._json(HTTPStatus.OK, self.playground.applicable_examples(), cookie)
            else:
                self._json(HTTPStatus.NOT_FOUND, {"error": {"type": "NotFound", "message": path}})
        except Exception as exc:  # noqa: BLE001
            traceback.print_exc()
            self._json(HTTPStatus.INTERNAL_SERVER_ERROR, {"error": {"type": type(exc).__name__, "message": str(exc)}})

    def do_POST(self) -> None:  # noqa: N802
        path = self.path.split("?", 1)[0]
        try:
            body = self._read_json()
            visitor_id, fresh = self._visitor_id()
            cookie = visitor_id if fresh else None
            pg = self.playground
            if path == "/api/step":
                action, text = body.get("action"), body.get("text")
                if text is None and (not isinstance(action, dict) or not action.get("type")):
                    raise ValueError('the body must contain {"text": "..."} or {"action": {"type": ..., "arguments": {...}}}')
                duration = body.get("duration")
                visitor = pg.visitor(visitor_id)
                result = pg.step(visitor, action, body.get("context"), str(text) if text is not None else None,
                                 float(duration) if duration not in (None, "") else None)
                self._json(HTTPStatus.OK, {**result, "budget": pg.budget(visitor)}, cookie)
            elif path == "/api/parse":
                duration = body.get("duration")
                self._json(HTTPStatus.OK, parse_text(str(body.get("text", "")), pg.input_mode,
                                                     float(duration) if duration not in (None, "") else None).model_dump(mode="json"), cookie)
            elif path == "/api/inspect":
                self._json(HTTPStatus.OK, pg.inspect(pg.visitor(visitor_id), str(body.get("type", ""))), cookie)
            elif path == "/api/search":
                self._json(HTTPStatus.OK, pg.search(pg.visitor(visitor_id), str(body.get("query", "")), int(body.get("limit") or 10)), cookie)
            elif path == "/api/reset":
                self._json(HTTPStatus.OK, pg.reset(pg.visitor(visitor_id)), cookie)
            elif path == "/api/examples/start":
                self._json(HTTPStatus.OK, pg.start_example(pg.visitor(visitor_id), str(body.get("id", ""))), cookie)
            elif path == "/api/credentials":
                visitor = pg.visitor(visitor_id)
                if body.get("forget"):
                    self._json(HTTPStatus.OK, pg.forget_credentials(visitor), cookie)
                else:
                    self._json(HTTPStatus.OK, pg.set_credentials(visitor, str(body.get("model", "")), str(body.get("api_key", "")),
                                                                 str(body.get("base_url") or "") or None), cookie)
            else:
                self._json(HTTPStatus.NOT_FOUND, {"error": {"type": "NotFound", "message": path}})
        except Busy as exc:
            self._json(HTTPStatus.CONFLICT, {"error": {"type": "Busy", "message": str(exc)}})
        except BudgetExhausted as exc:
            self._json(HTTPStatus.FORBIDDEN, {"error": {"type": "BudgetExhausted", "message": str(exc)}})
        except (ValueError, KeyError) as exc:
            self._json(HTTPStatus.BAD_REQUEST, {"error": {"type": type(exc).__name__, "message": str(exc)}})
        except Exception as exc:  # noqa: BLE001
            traceback.print_exc()
            self._json(HTTPStatus.INTERNAL_SERVER_ERROR, {"error": {"type": type(exc).__name__, "message": str(exc)}})


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--package", required=True, help="The worldbook (package directory) to simulate")
    parser.add_argument("--session", required=True, help="Root directory for episodes: every visitor, play and new episode gets its own "
                                                          "session directory (state, memory, audit) under it")
    parser.add_argument("--allow-unvalidated", action="store_true", help="Simulate a candidate package")
    parser.add_argument("--initial-state", help="Initial EnvironmentState as JSON or @path (what the episode already knows; reconstructed "
                                                "packages otherwise start from an explicitly empty state)")
    parser.add_argument("--environment-prompt", help="Environment description for the agent and tracker, text or @path")
    parser.add_argument("--title", help=f"Page title (default: {DEFAULT_TITLE!r})")
    parser.add_argument("--tagline", help="One line under the title (default: derived from the worldbook)")
    parser.add_argument("--history-window", type=int, default=12, help="Recent turns shown in the world-model agent's brief")
    parser.add_argument("--knowledge-gate", action="store_true", help="Label every worldbook item shown to the agent with its applicability")
    parser.add_argument("--knowledge-names", choices=["paths", "pages", "screens"], default="paths")
    parser.add_argument("--effect-rejection", choices=["fail", "keep_observation"], default="keep_observation",
                        help="When the agent's effects fail verification: fail the step, or keep its observation and drop the effects")
    parser.add_argument("--environment-prompt-chars", type=int, default=24000, help="Bound on --environment-prompt (0 = whole)")
    parser.add_argument("--no-model", action="store_true", help="Rules only: no world-model agent on the server's side (visitors may still bring a key)")
    parser.add_argument("--examples", help="JSON file of suggested plays (default: docs/playground/examples.json)")
    parser.add_argument("--allow-missing-observation", action="store_true",
                        help="Accept a rule-free agent submission without observation text (rendered from the JSON fallback, as in "
                             "awb-run); by default the playground sends it back once as invalid")
    parser.add_argument("--max-actions-per-visitor", type=int, default=20,
                        help="Free actions each visitor (browser cookie) may send on the server's key, across plays and episodes; 0 = unlimited")
    parser.add_argument("--max-actions-own-key", type=int, default=0, help="Actions a visitor may send on their own key; 0 = unlimited")
    parser.add_argument("--max-actions-total", type=int, default=0,
                        help="Free actions the server accepts in total before refusing the server's key to everyone; 0 = unlimited")
    parser.add_argument("--max-concurrent-steps", type=int, default=4, help="Steps simulated at the same time across visitors")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8765)
    _add_model_options(parser)
    _add_agent_options(parser)
    # Playground defaults: DeepSeek V4.1-Flash through OpenRouter, pinned to the upstream provider that answered every
    # call cleanly in the reported runs, with the settings those runs used (reasoning effort low, 65k cap, native tool
    # calling: over the JSON-schema transport DeepSeek returned empty effect lists on 68 of 69 play steps).
    parser.set_defaults(provider="chat", base_url="https://openrouter.ai/api/v1", api_key_env="OPENROUTER_API_KEY",
                        model="deepseek/deepseek-v4.1-flash", chat_reasoning_effort="low", agent_max_output_tokens=65536,
                        chat_provider_json='{"order": ["CoreWeave"], "allow_fallbacks": false}', agent_transport="tools")
    return parser


def make_server(args: argparse.Namespace) -> ThreadingHTTPServer:
    playground = Playground(args)
    handler = type("PlaygroundHandler", (Handler,), {"playground": playground})
    server = ThreadingHTTPServer((args.host, args.port), handler)
    server.playground = playground  # type: ignore[attr-defined]
    return server


def main(argv: list[str] | None = None) -> int:
    args = build_parser().parse_args(argv)
    server = make_server(args)
    pg = server.playground  # type: ignore[attr-defined]
    host, port = server.server_address[:2]
    print(f"Trace2Env playground: http://{host}:{port}", file=sys.stderr)
    print(f"  {pg.title} | {pg.tagline}", file=sys.stderr)
    name = pg.package.manifest.model_dump(mode='json').get('name')
    meta = pg.package.manifest.metadata
    print(f"  package  {name} ({meta.get('validation_status')}, {meta.get('construction_kind')})", file=sys.stderr)
    print(f"  model    {'none (rules only)' if args.no_model else args.model}  transport={args.agent_transport}  "
          f"trust={args.trust_policy or 'schema_checked'}  fast_path={args.fast_path}", file=sys.stderr)
    print(f"  episodes under {pg.session_root}  initial state: {pg.initial_state_source}", file=sys.stderr)
    print(f"  budget   {args.max_actions_per_visitor or 'unlimited'} free actions per visitor, "
          f"{args.max_actions_total or 'unlimited'} in total; own keys: {args.max_actions_own_key or 'unlimited'}", file=sys.stderr)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        pass
    finally:
        server.server_close()
    return 0


if __name__ == "__main__":
    sys.exit(main())