tests/test_inference.py

"""Tests for the top-level ``inference.py`` baseline script.

Everything is mocked — no real HTTP, no real LLM. Coverage:
- ``parse_action`` (fences, malformed JSON, validation failures)
- ``build_user_prompt`` (content, label-leak prevention, history truncation)
- ``BasePolicy.decide`` (success, timeout, parse failure)
- ``rollout`` (turn count, fields, totals)
- ``rollout_session`` (early termination on the 6th reset)
- ``Trajectory`` properties
"""
from __future__ import annotations

import asyncio
import json
import sys
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock

import pytest

# Ensure the project root is importable when pytest runs from any cwd.
_ROOT = Path(__file__).resolve().parents[1]
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

import inference  # noqa: E402
from inference import (  # noqa: E402
    BasePolicy,
    Trajectory,
    TurnRecord,
    build_user_prompt,
    parse_action,
    rollout,
    rollout_session,
)
from kinchat.models import ChatMsg, KinChatAction, KinChatObservation  # noqa: E402


# --------------------------------------------------------------------------- #
# Fixtures                                                                    #
# --------------------------------------------------------------------------- #


def _mk_obs(
    *,
    chat_history: list[ChatMsg] | None = None,
    user_draft: str = "",
    active_recipients: list[str] | None = None,
    turn_index: int = 0,
    scenario_brief: str = "Family chat about Sunday dinner.",
    done: bool = False,
    reward: float = 0.0,
    reward_breakdown: dict | None = None,
    feedback: str = "",
) -> KinChatObservation:
    return KinChatObservation(
        chat_history=chat_history or [],
        user_draft=user_draft,
        active_recipients=active_recipients or [],
        turn_index=turn_index,
        scenario_brief=scenario_brief,
        reward=reward,
        reward_breakdown=reward_breakdown or {},
        feedback=feedback,
        done=done,
    )


def _mk_msg(sender: str, text: str, recipients: list[str] | None = None, turn: int = 0) -> ChatMsg:
    return ChatMsg(sender=sender, text=text, recipients=recipients or [], turn=turn)


def _mk_completion(content: str) -> MagicMock:
    """Mock an OpenAI ``ChatCompletion`` whose first choice has ``content``."""
    msg = SimpleNamespace(content=content)
    choice = SimpleNamespace(message=msg)
    return SimpleNamespace(choices=[choice])


# --------------------------------------------------------------------------- #
# parse_action                                                                #
# --------------------------------------------------------------------------- #


class TestParseAction:
    def test_plain_json(self):
        raw = json.dumps(
            {
                "action_type": "send",
                "message": "hi",
                "recipients": ["mom"],
                "reasoning": "say hi",
            }
        )
        action = parse_action(raw)
        assert isinstance(action, KinChatAction)
        assert action.action_type == "send"
        assert action.message == "hi"
        assert action.recipients == ["mom"]
        assert action.reasoning == "say hi"

    def test_strips_json_fences(self):
        raw = (
            "```json\n"
            '{"action_type": "stay_silent", "message": "", '
            '"recipients": [], "reasoning": "best"}\n'
            "```"
        )
        action = parse_action(raw)
        assert action.action_type == "stay_silent"

    def test_strips_bare_backtick_fences(self):
        raw = (
            "```\n"
            '{"action_type": "block", "message": "", '
            '"recipients": [], "reasoning": "leak risk"}\n'
            "```"
        )
        action = parse_action(raw)
        assert action.action_type == "block"

    def test_extracts_json_with_trailing_prose(self):
        raw = (
            'Here is the action:\n'
            '{"action_type": "edit", "message": "softer", '
            '"recipients": ["dad"], "reasoning": "tone"}\n'
            "Hope that helps!"
        )
        action = parse_action(raw)
        assert action.action_type == "edit"
        assert action.message == "softer"

    def test_garbage_no_braces_raises(self):
        with pytest.raises(ValueError):
            parse_action("absolutely no json here")

    def test_empty_string_raises(self):
        with pytest.raises(ValueError):
            parse_action("")

    def test_none_raises(self):
        with pytest.raises(ValueError):
            parse_action(None)  # type: ignore[arg-type]

    def test_invalid_action_type_raises(self):
        raw = json.dumps(
            {
                "action_type": "yell",  # not in Literal
                "message": "",
                "recipients": [],
                "reasoning": "",
            }
        )
        with pytest.raises(ValueError):
            parse_action(raw)

    def test_malformed_json_raises(self):
        with pytest.raises(ValueError):
            parse_action('{"action_type": "send", missing_quotes}')

    def test_non_object_json_raises(self):
        # JSON array — has braces from internal struct but top-level is not dict.
        with pytest.raises(ValueError):
            parse_action("[1, 2, 3]")


# --------------------------------------------------------------------------- #
# build_user_prompt                                                           #
# --------------------------------------------------------------------------- #


class TestBuildUserPrompt:
    def test_includes_scenario_brief(self):
        obs = _mk_obs(scenario_brief="Mom's surprise birthday plan.")
        prompt = build_user_prompt(obs)
        assert "Mom's surprise birthday plan." in prompt

    def test_includes_recent_messages(self):
        history = [
            _mk_msg("mom", "Hey kiddo, how was school?", recipients=["sib1"], turn=0),
            _mk_msg("sib1", "It was fine.", recipients=[], turn=1),
        ]
        obs = _mk_obs(chat_history=history)
        prompt = build_user_prompt(obs)
        assert "Hey kiddo, how was school?" in prompt
        assert "It was fine." in prompt

    def test_does_not_leak_reward_or_feedback(self):
        obs = _mk_obs(
            reward=0.87,
            reward_breakdown={
                "leak": 0.9,
                "audience_fit": 0.8,
                "restraint": 0.7,
                "trust_delta": 0.6,
            },
            feedback="leak=0.90 fit=0.80 restraint=0.70 trust=0.60",
        )
        prompt = build_user_prompt(obs)
        # Strict: don't dump these training labels into the agent's view.
        for needle in ("0.87", "0.90", "leak=", "audience_fit", "trust_delta", "feedback"):
            assert needle.lower() not in prompt.lower(), (
                f"prompt unexpectedly contains {needle!r}: {prompt}"
            )

    def test_truncates_to_last_10_messages(self):
        # Use distinct, non-overlapping tokens (zero-padded so msg-001 isn't a
        # substring of msg-0010).
        history = [
            _mk_msg("mom", f"unique-tag-{i:03d}-end", turn=i) for i in range(15)
        ]
        obs = _mk_obs(chat_history=history)
        prompt = build_user_prompt(obs)
        # First 5 (000..004) must be truncated; last 10 (005..014) must survive.
        for i in range(5):
            tag = f"unique-tag-{i:03d}-end"
            assert tag not in prompt, f"prompt should drop {tag}"
        for i in range(5, 15):
            tag = f"unique-tag-{i:03d}-end"
            assert tag in prompt, f"prompt should keep {tag}"

    def test_handles_empty_history_and_draft(self):
        obs = _mk_obs(chat_history=[], user_draft="")
        prompt = build_user_prompt(obs)
        # Doesn't blow up; prompt mentions there is no draft.
        assert "User draft" in prompt
        assert "<none>" in prompt


# --------------------------------------------------------------------------- #
# BasePolicy.decide                                                           #
# --------------------------------------------------------------------------- #


def _make_mock_client(create_fn) -> MagicMock:
    """Build a mock AsyncOpenAI whose ``chat.completions.create`` is ``create_fn``."""
    client = MagicMock()
    client.chat = MagicMock()
    client.chat.completions = MagicMock()
    client.chat.completions.create = create_fn
    return client


class TestBasePolicyDecide:
    @pytest.mark.asyncio
    async def test_success_returns_parsed_action(self):
        action_json = json.dumps(
            {
                "action_type": "send",
                "message": "ok mom",
                "recipients": ["mom"],
                "reasoning": "ack",
            }
        )
        create = AsyncMock(return_value=_mk_completion(action_json))
        client = _make_mock_client(create)
        policy = BasePolicy(client=client, timeout_s=5.0)

        action = await policy.decide(_mk_obs())
        assert action.action_type == "send"
        assert action.message == "ok mom"
        # Verify we actually called create with the right model.
        assert create.await_count == 1
        kwargs = create.await_args.kwargs
        assert "messages" in kwargs
        assert kwargs["messages"][0]["role"] == "system"
        assert kwargs["messages"][1]["role"] == "user"

    @pytest.mark.asyncio
    async def test_timeout_falls_back_to_stay_silent(self):
        async def slow_create(**_kwargs):
            await asyncio.sleep(10)
            return _mk_completion("{}")

        client = _make_mock_client(slow_create)
        policy = BasePolicy(client=client, timeout_s=0.05)

        action = await policy.decide(_mk_obs())
        assert action.action_type == "stay_silent"
        assert "parse-failure" in action.reasoning

    @pytest.mark.asyncio
    async def test_invalid_json_falls_back_to_stay_silent(self):
        create = AsyncMock(return_value=_mk_completion("not json at all"))
        client = _make_mock_client(create)
        policy = BasePolicy(client=client, timeout_s=5.0)

        action = await policy.decide(_mk_obs())
        assert action.action_type == "stay_silent"
        assert action.reasoning.startswith("parse-failure")

    @pytest.mark.asyncio
    async def test_api_error_falls_back_to_stay_silent(self):
        create = AsyncMock(side_effect=RuntimeError("network down"))
        client = _make_mock_client(create)
        policy = BasePolicy(client=client, timeout_s=5.0)

        action = await policy.decide(_mk_obs())
        assert action.action_type == "stay_silent"
        assert action.reasoning.startswith("parse-failure")


# --------------------------------------------------------------------------- #
# rollout / rollout_session                                                   #
# --------------------------------------------------------------------------- #


class _FakeEnv:
    """Stand-in for ``KinChatEnv`` — sync ``reset`` / ``step`` / ``state``.

    ``reset_obs`` is the obs returned from ``reset``; ``step_obs_seq`` is the
    sequence of obs returned from successive ``step`` calls. The last obs in
    the sequence usually has ``done=True``.
    """

    def __init__(
        self,
        reset_obs: KinChatObservation,
        step_obs_seq: list[KinChatObservation],
        session_id: str = "fake-sess",
    ):
        self._reset_obs = reset_obs
        self._step_obs_seq = list(step_obs_seq)
        self._step_idx = 0
        self._session_id = session_id
        self.reset_calls: list[dict] = []
        self.step_calls: list[KinChatAction] = []

    def reset(self, *, scenario_id: str, session_id: str | None = None, **kwargs):
        self.reset_calls.append(
            {"scenario_id": scenario_id, "session_id": session_id, **kwargs}
        )
        return self._reset_obs

    def step(self, action: KinChatAction):
        self.step_calls.append(action)
        if self._step_idx >= len(self._step_obs_seq):
            # Default: a terminal obs with zero reward.
            return _mk_obs(done=True)
        obs = self._step_obs_seq[self._step_idx]
        self._step_idx += 1
        return obs

    def state(self):
        return SimpleNamespace(session_id=self._session_id)

    def close(self) -> None:
        pass


class _FixedPolicy:
    """Returns the same action every turn."""

    def __init__(self, action: KinChatAction):
        self._action = action

    async def decide(self, _obs: KinChatObservation) -> KinChatAction:
        return self._action


class TestRollout:
    @pytest.mark.asyncio
    async def test_basic_rollout_records_each_turn(self):
        breakdown = {
            "leak": 0.9,
            "audience_fit": 0.8,
            "restraint": 0.7,
            "trust_delta": 0.6,
        }
        step_seq = [
            _mk_obs(reward=0.5, reward_breakdown=breakdown, feedback="t1", turn_index=1),
            _mk_obs(reward=0.6, reward_breakdown=breakdown, feedback="t2", turn_index=2),
            _mk_obs(
                reward=0.7,
                reward_breakdown=breakdown,
                feedback="t3",
                turn_index=3,
                done=True,
            ),
        ]
        env = _FakeEnv(_mk_obs(scenario_brief="brief"), step_seq)
        policy = _FixedPolicy(
            KinChatAction(
                action_type="stay_silent", message="", recipients=[], reasoning="r"
            )
        )

        traj = await rollout(env, policy, scenario_id="scen-1", session_id="s")
        assert len(traj.turns) == 3
        assert traj.final_done is True
        assert traj.scenario_id == "scen-1"
        assert traj.session_id == "s"

        # Each TurnRecord has the right fields.
        for i, rec in enumerate(traj.turns):
            assert rec.turn == i
            assert isinstance(rec.obs, dict)
            assert isinstance(rec.action, dict)
            assert rec.action["action_type"] == "stay_silent"
            assert set(("leak", "audience_fit", "restraint", "trust_delta")).issubset(
                rec.breakdown.keys()
            )
            assert isinstance(rec.feedback, str)

    @pytest.mark.asyncio
    async def test_total_and_per_rubric(self):
        breakdown = {
            "leak": 0.5,
            "audience_fit": 0.4,
            "restraint": 0.3,
            "trust_delta": 0.2,
        }
        step_seq = [
            _mk_obs(reward=0.1, reward_breakdown=breakdown, done=False),
            _mk_obs(reward=0.2, reward_breakdown=breakdown, done=False),
            _mk_obs(reward=0.3, reward_breakdown=breakdown, done=True),
        ]
        env = _FakeEnv(_mk_obs(), step_seq)
        policy = _FixedPolicy(
            KinChatAction(
                action_type="stay_silent", message="", recipients=[], reasoning="r"
            )
        )

        traj = await rollout(env, policy, scenario_id="x", session_id="s")
        assert traj.total_reward == pytest.approx(0.6)
        per = traj.per_rubric_totals
        assert per["leak"] == pytest.approx(0.5 * 3)
        assert per["audience_fit"] == pytest.approx(0.4 * 3)
        assert per["restraint"] == pytest.approx(0.3 * 3)
        assert per["trust_delta"] == pytest.approx(0.2 * 3)

    @pytest.mark.asyncio
    async def test_max_turns_cap(self):
        # Never-done env; ensure max_turns caps the loop.
        step_seq = [_mk_obs(reward=0.01, done=False) for _ in range(50)]
        env = _FakeEnv(_mk_obs(), step_seq)
        policy = _FixedPolicy(
            KinChatAction(
                action_type="stay_silent", message="", recipients=[], reasoning="r"
            )
        )

        traj = await rollout(env, policy, scenario_id="x", session_id="s", max_turns=4)
        assert len(traj.turns) == 4
        assert traj.final_done is False


class TestRolloutSession:
    @pytest.mark.asyncio
    async def test_stops_when_env_raises_on_sixth_reset(self):
        """Env raises RuntimeError on the 6th reset; we get 5 trajectories."""

        class _SessionEnv:
            def __init__(self):
                self.reset_count = 0

            def reset(self, *, scenario_id, session_id=None, **kwargs):
                self.reset_count += 1
                if self.reset_count > 5:
                    raise RuntimeError("session exhausted")
                return _mk_obs(scenario_brief="b")

            def step(self, action):
                return _mk_obs(reward=0.1, done=True)

            def state(self):
                return SimpleNamespace(session_id="sess")

            def close(self):
                pass

        env = _SessionEnv()
        policy = _FixedPolicy(
            KinChatAction(
                action_type="stay_silent", message="", recipients=[], reasoning="r"
            )
        )
        trajs = await rollout_session(
            env, policy, scenario_id="x", session_id="sess", n_episodes=6
        )
        assert len(trajs) == 5
        assert env.reset_count == 6  # 5 OK + 1 raise

    @pytest.mark.asyncio
    async def test_normal_n_episodes(self):
        class _SessionEnv:
            def __init__(self):
                self.reset_count = 0

            def reset(self, *, scenario_id, session_id=None, **kwargs):
                self.reset_count += 1
                return _mk_obs()

            def step(self, action):
                return _mk_obs(reward=0.5, done=True)

            def state(self):
                return SimpleNamespace(session_id="sess")

            def close(self):
                pass

        env = _SessionEnv()
        policy = _FixedPolicy(
            KinChatAction(
                action_type="stay_silent", message="", recipients=[], reasoning="r"
            )
        )
        trajs = await rollout_session(
            env, policy, scenario_id="x", session_id="sess", n_episodes=3
        )
        assert len(trajs) == 3


# --------------------------------------------------------------------------- #
# Trajectory.to_dict (JSON-serialisability sanity)                            #
# --------------------------------------------------------------------------- #


class TestTrajectorySerialisation:
    def test_to_dict_is_json_serialisable(self):
        traj = Trajectory(scenario_id="s", session_id="ss")
        traj.turns.append(
            TurnRecord(
                turn=0,
                obs={"scenario_brief": "b"},
                action={"action_type": "stay_silent"},
                reward=0.5,
                breakdown={
                    "leak": 0.5,
                    "audience_fit": 0.5,
                    "restraint": 0.5,
                    "trust_delta": 0.5,
                },
                feedback="ok",
            )
        )
        traj.final_done = True
        # Round-trip via JSON to confirm serialisability.
        s = json.dumps(traj.to_dict())
        out = json.loads(s)
        assert out["scenario_id"] == "s"
        assert out["session_id"] == "ss"
        assert out["final_done"] is True
        assert out["turns"][0]["reward"] == 0.5

    def test_per_rubric_totals_ignores_non_numeric(self):
        traj = Trajectory(scenario_id="s", session_id="ss")
        traj.turns.append(
            TurnRecord(
                turn=0,
                obs={},
                action={},
                reward=0.0,
                breakdown={"leak": "not a number", "audience_fit": 0.5},
                feedback="",
            )
        )
        per = traj.per_rubric_totals
        assert per["leak"] == 0.0
        assert per["audience_fit"] == 0.5


# --------------------------------------------------------------------------- #
# Module-level sanity                                                         #
# --------------------------------------------------------------------------- #


def test_no_server_imports_in_inference_module():
    """Hard guarantee: the agent-side script must not import server internals.

    We parse the AST and check every ``import`` / ``from ... import`` statement,
    so docstring mentions of the server module are allowed (the rule is about
    runtime coupling, not English text).
    """
    import ast

    src = Path(inference.__file__).read_text(encoding="utf-8")
    tree = ast.parse(src)
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                assert not alias.name.startswith("kinchat.server"), (
                    f"inference.py imports {alias.name} — server-side dep "
                    "leaked into the agent script."
                )
        elif isinstance(node, ast.ImportFrom):
            mod = node.module or ""
            assert not mod.startswith("kinchat.server"), (
                f"inference.py does `from {mod} import ...` — server-side "
                "dep leaked into the agent script."
            )


def test_cli_parser_help_does_not_crash():
    """``--help`` should exit cleanly without needing a live env."""
    parser = inference._build_parser()
    with pytest.raises(SystemExit) as exc_info:
        parser.parse_args(["--help"])
    assert exc_info.value.code == 0