src/trace2env/state_tracking.py

"""Run-time state reconstruction from observed interactions.

A benchmark record or a live caller supplies a trajectory prefix: the actions a task agent took
and the observations the real environment returned, but no explicit state. ``StateTracker``
advances an ``EnvironmentState`` through such turns. A turn explained by an applicable rule
whose rendered observation matches the observed one is applied deterministically; every other
turn is batched into a bounded ``track_state`` model call whose typed mutations are checked
against the state schema before they touch the state. Tracking follows observed reality: an
invariant contradicted by an observation is recorded as a diagnostic about the package, not
enforced against the evidence.
"""

from __future__ import annotations

import json
import re
from typing import Any

from trace2env.engine import (
    AmbiguousRuleError,
    apply_mutations,
    check_invariants,
    deterministic_plan,
    process_due_events,
    render_template,
    select_rule,
)
from trace2env.eligibility import rule_eligible
from trace2env.llm import ProviderRefusal, StructuredLLM, StructuredOutputError
from trace2env.memory import EpisodicMemory
from trace2env.models import (
    EnvironmentState,
    MemoryEntry,
    NormalizedAction,
    ObservedTurn,
    StateMutation,
    StateTrackingResult,
    StateTrackingStep,
    TransitionPlan,
)
from trace2env.package import EnvironmentPackage
from trace2env.prompts import STATE_TRACKER, STATE_TRACKER_EXACT, STATE_TRACKER_WAIT
from trace2env.transcript import initial_prompt_state
from trace2env.validation import validate_mutations, validate_state_types
from trace2env.workspace import clip_text, compact_action

# Same default bound as runtime.ENVIRONMENT_PROMPT_CHARS, kept local so neither module imports the other.
ENVIRONMENT_PROMPT_CHARS = 24000
_WHITESPACE = re.compile(r"\s+")


def normalize_observation(text: str) -> str:
    return _WHITESPACE.sub(" ", text).strip()


class StateTracker:
    def __init__(
        self,
        package: EnvironmentPackage,
        llm: StructuredLLM | None = None,
        *,
        window: int = 8,
        observation_chars: int = 4000,
        environment_prompt: str = "",
        memory: EpisodicMemory | None = None,
        features: set[str] | None = None,
        disabled: bool = False,
        environment_prompt_chars: int | None = ENVIRONMENT_PROMPT_CHARS,
    ):
        self.package = package
        self.llm = llm
        self.memory = memory
        # ``disabled``: the trace2env_no_state control — every observed turn is remembered (history stays
        # available) but no rule and no model ever changes the state.
        self.disabled = disabled
        self.window = max(1, window)
        self.observation_chars = observation_chars
        # Caller-owned environment context, bounded by default (the setting of every reported AgentWorldBench
        # run). It may end with tool definitions or examples; ``environment_prompt_chars=None`` keeps it whole
        # (the long-horizon runs).
        self.environment_prompt = environment_prompt if environment_prompt_chars is None else environment_prompt[:environment_prompt_chars]
        # "wait": a rule without effects and without a template explains nothing and never replaces
        # model tracking; "compact": static schema/environment text goes into the (cacheable) system
        # message as a compact field list; "unknown": exact observed values are kept verbatim.
        self.features = set(features or ())
        self.aliases = {
            alias: spec.name for spec in package.action_schema.actions for alias in [spec.name, *spec.aliases]
        }
        self.minimum_confidence = float(package.manifest.metadata.get("min_rule_confidence", 0.0))
        self._carried: list[str] = []

    def canonical(self, action: NormalizedAction) -> NormalizedAction:
        return action.model_copy(update={"type": self.aliases.get(action.type, action.type)})

    def advance(
        self, state: EnvironmentState, turns: list[ObservedTurn]
    ) -> tuple[EnvironmentState, list[StateTrackingStep]]:
        steps: list[StateTrackingStep] = []
        pending: list[ObservedTurn] = []
        if self.disabled:
            for observed in turns:
                self._remember(observed.model_copy(update={"action": self.canonical(observed.action)}), [], [])
            return state, [StateTrackingStep(turns=[t.turn for t in turns], route="skipped",
                                             issues=["state tracking disabled; turns recorded in memory only"])]
        for observed in turns:
            turn = observed.model_copy(update={"action": self.canonical(observed.action)})
            rule = self._applicable_rule(state, turn)
            if rule is not None and pending:
                # Pending model-explained turns precede this one; refresh the state before trusting the rule.
                state, step = self._flush(state, pending)
                steps.append(step)
                pending = []
                rule = self._applicable_rule(state, turn)
            if rule is not None:
                outcome = self._apply_rule(state, rule, turn)
                if outcome is not None:
                    state, step = outcome
                    steps.append(step)
                    continue
            pending.append(turn)
            if len(pending) >= self.window:
                state, step = self._flush(state, pending)
                steps.append(step)
                pending = []
        if pending:
            state, step = self._flush(state, pending)
            steps.append(step)
        return state, steps

    def _eligible_rules(self, state: EnvironmentState) -> list:
        return [
            rule for rule in self.package.rules
            if rule_eligible(rule, self.package.scope, state, minimum_confidence=self.minimum_confidence)
        ]

    def _applicable_rule(self, state: EnvironmentState, turn: ObservedTurn):
        working, _ = process_due_events(state, state.step + 1)
        try:
            return select_rule(self._eligible_rules(working), working, turn.action, self.package.scope)
        except AmbiguousRuleError as exc:
            self._carried.append(f"turn {turn.turn}: {exc}")
            return None

    def _render_context(self, before: EnvironmentState, after: EnvironmentState, plan: TransitionPlan) -> dict[str, Any]:
        return {
            "action": plan.action.model_dump(mode="python"),
            "state_before": before.model_dump(mode="python"),
            "state_after": after.model_dump(mode="python"),
            "outcome": plan.outcome.value,
        }

    def _apply_rule(
        self, state: EnvironmentState, rule, turn: ObservedTurn
    ) -> tuple[EnvironmentState, StateTrackingStep] | None:
        working, _ = process_due_events(state, state.step + 1)
        working.step = state.step + 1
        if "wait" in self.features and not rule.effects:
            contract_template = next((item.template for item in self.package.renderers if item.id == rule.renderer), None)
            if not (rule.observation_template or contract_template):
                # Nothing to apply and nothing to check against the observation: the observed turn still
                # has to be explained by the model, or its effects (a finished process, new files) are lost.
                self._carried.append(f"turn {turn.turn}: rule {rule.id} has no effects and no template; tracked by the model instead")
                return None
        plan = deterministic_plan(rule, turn.action)
        spec = next((item for item in self.package.action_schema.actions if item.name == turn.action.type), None)
        try:
            validate_mutations(plan.effects, self.package.state_schema, spec)
            candidate = apply_mutations(working, plan.effects, action=turn.action)
            validate_state_types(candidate, self.package.state_schema)
        except (ValueError, KeyError) as exc:
            self._carried.append(f"turn {turn.turn}: rule {rule.id} could not be applied: {exc}")
            return None
        contract = next((item for item in self.package.renderers if item.id == plan.renderer_id), None)
        template = plan.observation_template or (contract.template if contract else None)
        if template:
            try:
                rendered = render_template(template, self._render_context(state, candidate, plan))
            except KeyError as exc:
                self._carried.append(f"turn {turn.turn}: rule {rule.id} template failed: {exc}")
                return None
            if normalize_observation(rendered) != normalize_observation(turn.observation):
                # The rule disagrees with what the environment actually returned; do not trust it here.
                self._carried.append(f"turn {turn.turn}: rule {rule.id} predicted a different observation than observed")
                return None
        issues, self._carried = self._carried, []
        verification = check_invariants(self.package.invariants, candidate, turn.action)
        issues += [f"{issue.code}: {issue.message}" for issue in verification.issues if issue.severity == "error"]
        self._remember(turn, plan.effects, [f"rule:{rule.id}"])
        step = StateTrackingStep(turns=[turn.turn], route="rule", rule_id=rule.id, applied=len(plan.effects), issues=issues)
        return candidate, step

    def _remember(self, turn: ObservedTurn, effects: list[StateMutation], citations: list[str]) -> None:
        if self.memory is not None:
            self.memory.record(MemoryEntry(turn=turn.turn, kind="observed", action=turn.action,
                                           observation=turn.observation, effects=effects, citations=citations))

    def system_prompt(self) -> str:
        """The tracker's system text; with ``compact`` it also carries the static schema and environment description."""
        system = STATE_TRACKER
        if "wait" in self.features:
            system += "\n\n" + STATE_TRACKER_WAIT
        if "unknown" in self.features:
            system += "\n\n" + STATE_TRACKER_EXACT
        if "compact" in self.features:
            lines = [f"- {field.path} ({field.type}{'' if field.mutable else ', immutable'}): {field.description[:120]}"
                     for field in self.package.state_schema.fields]
            system += "\n\n# Declared state fields\n" + "\n".join(lines)
            if self.environment_prompt:
                system += "\n\n# Environment description (context only)\n" + self.environment_prompt
        return system

    def _tracking_request(self, working: EnvironmentState, pending: list[ObservedTurn]) -> tuple[str, dict[str, Any]]:
        turns = [
            {
                "turn": turn.turn,
                "action": compact_action(turn.action) if "compact" in self.features else turn.action.model_dump(mode="json"),
                "observation": clip_text(turn.observation, self.observation_chars),
            }
            for turn in pending
        ]
        instruction = ("Return the typed mutations, on declared state paths only, that reflect what these "
                       "observations establish about the environment after the last listed turn.")
        if "compact" in self.features:
            return self.system_prompt(), {"current_state": working.model_dump(mode="json"), "observed_turns": turns,
                                          "instruction": instruction}
        return self.system_prompt(), {
            "environment": self.environment_prompt,
            "state_schema": [field.model_dump(mode="json") for field in self.package.state_schema.fields],
            "current_state": working.model_dump(mode="json"),
            "observed_turns": turns,
            "instruction": instruction,
        }

    def track_initial_prompt(self, state: EnvironmentState, observation: str, action: NormalizedAction) -> tuple[EnvironmentState, StateTrackingStep] | None:
        """Track an initial screen that is just an idle prompt without a model call (``compact``).

        The working directory and the idle terminal mode are set when the schema declares them;
        the prompt line itself is kept in episodic memory, where the agent and the transcript
        scaffold read the user/host from.
        """
        if self.disabled:
            return None
        if "compact" not in self.features:
            return None
        parsed = initial_prompt_state(observation)
        if parsed is None:
            return None
        declared = {field.path: field for field in self.package.state_schema.fields}
        mutations = []
        if "session.cwd" in declared and declared["session.cwd"].type in {"string", "any"}:
            mutations.append(StateMutation(op="set", path="session.cwd", value=parsed["cwd"]))
        if "surface.mode" in declared and declared["surface.mode"].type in {"string", "any"}:
            mutations.append(StateMutation(op="set", path="surface.mode", value="prompt"))
        working, _ = process_due_events(state, state.step + 1)
        working.step = state.step + 1
        try:
            validate_mutations(mutations, self.package.state_schema, None)
            working = apply_mutations(working, mutations)
            validate_state_types(working, self.package.state_schema)
        except (ValueError, KeyError):
            return None
        turn = ObservedTurn(turn=0, action=action, observation=observation)
        self._remember(turn, mutations, ["prompt-line"])
        return working, StateTrackingStep(turns=[0], route="rule", rule_id="initial-prompt-line", applied=len(mutations))

    def _flush(self, state: EnvironmentState, pending: list[ObservedTurn]) -> tuple[EnvironmentState, StateTrackingStep]:
        issues, self._carried = self._carried, []
        turns = [turn.turn for turn in pending]
        working, _ = process_due_events(state, state.step + len(pending))
        working.step = state.step + len(pending)
        if self.llm is None:
            issues.append("no tracking model configured; state left unchanged for these turns")
            for turn in pending:
                self._remember(turn, [], [])
            return working, StateTrackingStep(turns=turns, route="skipped", issues=issues)
        system, payload = self._tracking_request(working, pending)
        try:
            result = self.llm.complete(
                system=system,
                user=json.dumps(payload, ensure_ascii=False, default=str),
                response_model=StateTrackingResult,
                role="track_state",
            )
        except RuntimeError as exc:
            # One handler for every model-side failure, because ProviderRefusal and StructuredOutputError are
            # RuntimeErrors too and a re-raise from a sibling clause would escape the whole try statement.
            truncated = "cut off at the output limit" in str(exc)
            if isinstance(exc, ProviderRefusal) or (isinstance(exc, StructuredOutputError) and not truncated):
                # A content-policy refusal, or an answer that never met the contract, is deterministic for this
                # window: the turns stay in memory (history is unaffected), the state is left unchanged, and the
                # gap is recorded instead of ending the trajectory.
                reason = "refused by the provider" if isinstance(exc, ProviderRefusal) else "answer invalid after repair attempts"
                issues.append(f"turns {turns[0]}-{turns[-1]}: track_state {reason} ({str(exc)[:120]}); state left unchanged")
                for turn in pending:
                    self._remember(turn, [], [])
                return working, StateTrackingStep(turns=turns, route="skipped", issues=issues)
            if not truncated:
                raise
            # The model tried to write more mutations than fit (typically whole file contents). Track
            # the window in halves so the rest of the trajectory keeps a state; a single turn that
            # still overflows is skipped with the state left unchanged, and the gap is recorded.
            if len(pending) > 1:
                middle = len(pending) // 2
                self._carried = issues + [f"turns {turns[0]}-{turns[-1]}: track_state output exceeded the limit; tracked in halves"]
                state, first = self._flush(state, pending[:middle])
                state, second = self._flush(state, pending[middle:])
                return state, StateTrackingStep(
                    turns=turns, route="model", applied=first.applied + second.applied,
                    dropped=first.dropped + second.dropped, issues=first.issues + second.issues)
            issues.append(f"turn {turns[0]}: track_state output exceeded the limit; state left unchanged for this turn")
            for turn in pending:
                self._remember(turn, [], [])
            return working, StateTrackingStep(turns=turns, route="skipped", issues=issues)
        accepted: list[StateMutation] = []
        dropped: list[str] = []
        for mutation in result.mutations:
            try:
                # No action spec: tracker mutations must carry literal values on declared paths.
                validate_mutations([mutation], self.package.state_schema, None)
                candidate = apply_mutations(working, [mutation])
                validate_state_types(candidate, self.package.state_schema)
            except (ValueError, KeyError) as exc:
                dropped.append(f"{mutation.op} {mutation.path}: {exc}")
                continue
            working = candidate
            accepted.append(mutation)
        issues += [f"uncertain: {path}" for path in result.uncertain_paths]
        verification = check_invariants(self.package.invariants, working, pending[-1].action)
        issues += [f"{issue.code}: {issue.message}" for issue in verification.issues if issue.severity == "error"]
        for turn in pending:
            # The window's mutations are attributed to its last observed turn.
            self._remember(turn, accepted if turn is pending[-1] else [], [])
        return working, StateTrackingStep(turns=turns, route="model", applied=len(accepted), dropped=dropped, issues=issues)