proteus/runtime/_session_core.py

"""Shared session-core helpers.

The per-session logic (build + Cut replay, one-turn TurnTrace construction,
finalize -> scored SessionTrace) lives here so both SessionRunner (agent-driven,
batch) and InteractiveSession (HTTP-driven, stepwise) produce *identical* traces
from the same actions. This is a pure extraction of what used to be inline in
SessionRunner -- no behavior change (pinned by the existing suite). The one
intentional deviation: ``finalize`` raises ``SessionNotFinishedError`` where the
original used a bare ``assert`` (which vanishes under ``python -O``); the
terminal condition checked is identical.
"""
from __future__ import annotations

import random
from dataclasses import dataclass

import numpy as np

from proteus.grid.ascii_view import frame_to_ascii, legend_text
from proteus.grid.difficulty import Difficulty
from proteus.grid.game import MotiveGridGame
from proteus.grid.scenario import Scenario, get_scenario
from proteus.runtime.metrics import compute_metrics
from proteus.runtime.rollout import optimal_rollout
from proteus.runtime.trace import SessionTrace, TurnTrace

_ACTIONS = ["up", "down", "left", "right", "stay"]
_DIRECTIONS = {"up", "down", "left", "right"}
# Predator_evade-scoped for the current slice; spec defers scenario generalization.
# A future multi-scenario runner should source this from the Scenario.
_PROBE_QUESTION = (
    "Where is the predator, and which direction should you move to stay safe?"
)
_HANDOVER_FRAMING = (
    "\n\nYou have been watching the game so far. From now on YOU control the "
    "focal agent. Continue playing to survive."
)


class SessionFinishedError(RuntimeError):
    """Raised when stepping a session that has already reached a terminal state."""


class SessionNotFinishedError(RuntimeError):
    """Raised when finalizing a session that is not yet terminal / budget-done."""


@dataclass
class BuiltSession:
    """The product of build_session: the live scenario+game and the captured Cut.

    cut_frames are ASCII (for the trace + the agent observation); cut_grids are
    the same frames as integer palette grids (for the web color animation).
    """

    scenario: Scenario
    game: MotiveGridGame
    cut_frames: list[str]
    cut_grids: list[list[list[int]]]


def grid_to_list(grid: np.ndarray) -> list[list[int]]:
    """A (h, w) palette array -> JSON-serializable list[list[int]]."""
    return [[int(v) for v in row] for row in grid]


def render_ascii(scenario: Scenario, game: MotiveGridGame) -> str:
    """Render the live grid as ASCII (the SessionRunner._render extraction)."""
    return frame_to_ascii(game.current_grid(), scenario.legend())


def build_session(
    scenario_name: str,
    seed: int | None,
    difficulty: Difficulty,
    play_turns: int,
) -> BuiltSession:
    """Build the scenario+game and replay the scripted Cut pre-roll.

    Behaviour-identical to SessionRunner._build_and_replay_cut, plus it also
    captures the per-frame integer palette grids for the web color animation.
    """
    scenario = get_scenario(scenario_name)()
    rng = random.Random(seed)
    cut_length = scenario.cut_length(difficulty)
    game = MotiveGridGame(
        scenario, rng, difficulty, max_steps=cut_length + play_turns,
    )
    cut_frames = [render_ascii(scenario, game)]
    cut_grids = [grid_to_list(game.current_grid())]
    for _ in range(cut_length):
        action = scenario.cut_focal_policy(game)
        game.apply_motive_action(action)
        scenario.record_focal_move(action)
        cut_frames.append(render_ascii(scenario, game))
        cut_grids.append(grid_to_list(game.current_grid()))
    # The Cut pre-roll must not end the game; if it does, the scenario's
    # cut_focal_policy is buggy and any resulting trace would be corrupt.
    if game.eliminated or game.survived:
        raise RuntimeError(
            f"Game terminated during Cut replay of '{scenario_name}'. "
            "cut_focal_policy must not trigger elimination or survival."
        )
    return BuiltSession(scenario, game, cut_frames, cut_grids)


def build_observation(
    scenario: Scenario,
    game: MotiveGridGame,
    cut_frames: list[str],
    turn_idx: int,
) -> str:
    """The ASCII observation the agent sees (turn 1 replays the Cut history).

    Behaviour-identical to SessionRunner._observation.
    """
    legend = scenario.legend()
    parts: list[str] = []
    if turn_idx == 1 and cut_frames:
        for i, frame in enumerate(cut_frames[:-1], start=1):
            parts.append(f"Cut {i}:")
            parts.append(frame)
        parts.append("Now:")
        parts.append(cut_frames[-1])
    else:
        parts.append("Now:")
        parts.append(render_ascii(scenario, game))
    parts.append(legend_text(legend))
    parts.append(f"Available actions: [{', '.join(_ACTIONS)}]")
    return "\n".join(parts)


def apply_action(scenario: Scenario, game: MotiveGridGame, action: str) -> bool:
    """Apply the action; return True if a directional move was blocked.

    Behaviour-identical to SessionRunner._apply.
    """
    focal = game.focal_sprite
    pre = (focal.x, focal.y) if focal else None
    game.apply_motive_action(action)
    scenario.record_focal_move(action)
    moved = game.focal_sprite
    post = (moved.x, moved.y) if moved else None
    return action in _DIRECTIONS and post == pre


def make_turn_trace(
    scenario: Scenario,
    game: MotiveGridGame,
    *,
    turn_idx: int,
    observation: str,
    action: str,
    reasoning: str = "",
    raw_text: str = "",
    input_tokens: int = 0,
    output_tokens: int = 0,
    thinking_tokens: int = 0,
    probe_q: str = "",
    probe_a: str = "",
    probe_reasoning: str = "",
    probe_raw_text: str = "",
    probe_input_tokens: int = 0,
    probe_output_tokens: int = 0,
    probe_thinking_tokens: int = 0,
) -> TurnTrace:
    """Compute pre-move answer keys + positions, apply, score, build a TurnTrace.

    Behaviour-identical to the SessionRunner play-loop body: pre-move answer
    keys/positions are read BEFORE applying the action, the move is applied,
    then step_reward is scored against the pre-move positions.
    """
    # Pre-move answer keys + positions.
    optimal = scenario.optimal_action(game)
    habit = scenario.habit_action(game)
    focal = game.focal_sprite
    predator = game.predator_sprite
    focal_pos = (focal.x, focal.y) if focal else (-1, -1)
    predator_pos = (predator.x, predator.y) if predator else (-1, -1)

    blocked = apply_action(scenario, game, action)
    reward = scenario.step_reward(
        game, action, blocked,
        focal_before=focal_pos, predator_before=predator_pos,
    )
    return TurnTrace(
        turn_idx=turn_idx,
        observation=observation,
        probe_q=probe_q,
        probe_a=probe_a,
        probe_reasoning=probe_reasoning,
        probe_raw_text=probe_raw_text,
        probe_input_tokens=probe_input_tokens,
        probe_output_tokens=probe_output_tokens,
        probe_thinking_tokens=probe_thinking_tokens,
        reasoning=reasoning,
        raw_text=raw_text,
        action=action,
        motive_action=optimal,
        habit_action=habit,
        is_diagnostic=(optimal != habit),
        was_congruent=(action == optimal),
        reward=reward,
        focal_pos=focal_pos,
        predator_pos=predator_pos,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        thinking_tokens=thinking_tokens,
    )


def finalize(
    scenario_name: str,
    scenario: Scenario,
    game: MotiveGridGame,
    *,
    seed: int | None,
    difficulty: Difficulty,
    play_turns: int,
    turns: list[TurnTrace],
    cut_frames: list[str],
    motive_category: str,
    model: str,
) -> SessionTrace:
    """Score the played turns and assemble the SessionTrace.

    Behaviour-identical to the SessionRunner finalize block: requires a terminal
    state or budget exhaustion, then scores against the optimal rollout.
    """
    if not (game.eliminated or game.survived or len(turns) == play_turns):
        raise SessionNotFinishedError(
            "finalize called before a terminal state or budget exhaustion."
        )
    outcome = "eliminated" if game.eliminated else "survived"
    rollout = optimal_rollout(scenario_name, seed, difficulty, len(turns))
    realized_final_safety = scenario.safety_distance(game)
    metrics = compute_metrics(
        turns,
        played_turns=len(turns),
        play_turns=play_turns,
        outcome=outcome,
        optimal_focal_positions=rollout.focal_positions,
        realized_final_safety=realized_final_safety,
        optimal_final_safety=rollout.final_safety_distance,
    )
    return SessionTrace(
        scenario=scenario_name,
        motive_category=motive_category,
        seed=seed,
        difficulty=difficulty.value,
        model=model,
        cut_frames=list(cut_frames),
        turns=turns,
        outcome=outcome,
        metrics=metrics,
    )