tests/test_mind.py

"""Tests for the GODSEED mind (Agent C) — mock backend only.

Covers: trace shape for diverse wishes, all-11-tool coverage across the
suite, emit event ordering, the malformed-JSON re-ask/skip path (via
deliberately broken fake backends), moderation default-deny, determinism,
and grammar-file sanity. No model weights, no network, no asyncio plugins
(tests drive coroutines with asyncio.run).
"""

from __future__ import annotations

import asyncio
import json
import pathlib
import sys

import pytest

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from mind import prompts  # noqa: E402
from mind.backends import MockBackend  # noqa: E402
from mind.moderation_judge import judge  # noqa: E402
from mind.planner import Planner  # noqa: E402
from mind.validate import (  # noqa: E402
    MODERATION_GRAMMAR,
    TURN_GRAMMAR,
    is_moderation_grammar,
    parse_turn,
)

TOOLS = set(prompts.tool_names())

WORLD_SUMMARY = "Epoch 3. A barren obsidian world; one monolith; a void sky."

DIVERSE_WISHES = {
    "melancholy": "i miss the sea we used to walk beside",
    "joyful": "a festival of flowers and dancing under a gold sky",
    "nonsense": "blorptangle the seventh hum until it squirbles",
    "long": (
        "i wish for a quiet village on a hill where every window stays lit "
        "and the kettle is always warm and nobody has to knock to come in"
    ),
}

# Wishes chosen so the union of their mock scripts spans all 11 tools.
COVERAGE_WISHES = [
    "i wish for a mountain to watch over the valley",   # raise, weather, inscribe
    "a quiet sea under patient stars",                  # lower, water, sky, life(birds)
    "an ancient forest full of fireflies",              # flora, weather, life(fireflies)
    "a lighthouse for everyone still at sea",           # raise, structure, sky
    "let mushrooms sing beneath the dark",              # flora, structure, weather, sky
    "a storm to end all storms",                        # sky, weather, flora
    "a bright city of towers and warehouses and cafes", # district, structure(new kinds), road, life(carts)
    "a new town district by the market",                # district, structure(bank/market/house), road, life
    "blorptangle the seventh hum",                      # flora, structure, inscribe
]


# --------------------------------------------------------------------------
# Harness
# --------------------------------------------------------------------------

def run_grant(wish: str, backend=None):
    """Drive Planner.grant to completion; return (trace, events, calls)."""
    backend = backend if backend is not None else MockBackend(delay=0.0)
    events: list[dict] = []
    calls: list[dict] = []

    async def act(call: dict) -> str:
        calls.append(call)
        return f"ok: {call['tool']} done"

    async def emit(event: dict) -> None:
        events.append(dict(event))

    planner = Planner(backend)
    trace = asyncio.run(planner.grant(wish, WORLD_SUMMARY, act, emit))
    return trace, events, calls


class FlakyTurnBackend:
    """Yields malformed JSON on every FIRST turn attempt; the re-ask (detected
    via prompts.REASK_NOTICE) is delegated to a real MockBackend. Exercises
    the planner's single-re-ask recovery without ever skipping a turn."""

    name = "flaky"
    model_id = "flaky-fake"

    def __init__(self):
        self.inner = MockBackend(delay=0.0)
        self.first_attempts = 0
        self.reasks = 0

    async def generate_stream(self, prompt, grammar=None, max_tokens=256):
        if grammar is None or is_moderation_grammar(grammar):
            async for chunk in self.inner.generate_stream(prompt, grammar, max_tokens):
                yield chunk
            return
        if prompts.REASK_NOTICE in prompt:
            self.reasks += 1
            async for chunk in self.inner.generate_stream(prompt, grammar, max_tokens):
                yield chunk
            return
        self.first_attempts += 1
        yield '{"thought": "the words come out wro'  # truncated, unparseable


class AlwaysBrokenBackend:
    """Free-text requests yield nothing; constrained requests yield garbage.
    The planner must survive, skip turns, and bail after two consecutive
    skips with a fallback reading and epitaph."""

    name = "broken"
    model_id = "broken-fake"

    async def generate_stream(self, prompt, grammar=None, max_tokens=256):
        if grammar is None:
            return  # empty reading
        yield "}{ this is not json at all"


class ExplodingBackend:
    """Raises mid-stream — the judge must treat this as deny."""

    name = "exploding"
    model_id = "exploding-fake"

    async def generate_stream(self, prompt, grammar=None, max_tokens=256):
        yield '{"allo'
        raise RuntimeError("gpu fell off")


class GarbageVerdictBackend:
    name = "garbage"
    model_id = "garbage-fake"

    async def generate_stream(self, prompt, grammar=None, max_tokens=256):
        yield "the spirits say probably fine"


# --------------------------------------------------------------------------
# Trace shape & content
# --------------------------------------------------------------------------

@pytest.mark.parametrize("label", sorted(DIVERSE_WISHES))
def test_trace_shape(label):
    trace, _events, _calls = run_grant(DIVERSE_WISHES[label])

    assert set(trace) == {"reading", "turns", "epitaph", "ms_total", "model", "backend"}
    assert isinstance(trace["reading"], str) and trace["reading"].strip()
    assert isinstance(trace["turns"], list) and 1 <= len(trace["turns"]) <= 7
    assert isinstance(trace["epitaph"], str) and 0 < len(trace["epitaph"]) <= 120
    assert isinstance(trace["ms_total"], int) and trace["ms_total"] >= 0
    assert trace["backend"] == "mock"
    assert trace["model"] == "godseed-mock-scripts"

    # Final turn is the done turn: no call, no observation.
    last = trace["turns"][-1]
    assert last["call"] is None and last["observation"] is None

    # Trace must be JSON-serializable (it becomes a JSONL line).
    json.dumps(trace)


@pytest.mark.parametrize("label", sorted(DIVERSE_WISHES))
def test_call_turns_valid(label):
    trace, _events, calls = run_grant(DIVERSE_WISHES[label])
    call_turns = [t for t in trace["turns"] if t["call"] is not None]
    assert call_turns, "every script should execute at least one tool call"
    for turn in trace["turns"]:
        assert isinstance(turn["thought"], str) and len(turn["thought"]) <= 160
    for turn in call_turns:
        assert turn["call"]["tool"] in TOOLS
        assert isinstance(turn["call"]["args"], dict)
        assert isinstance(turn["observation"], str)
        assert turn["observation"].startswith("ok:")
    assert [t["call"] for t in call_turns] == calls


def test_all_eleven_tools_covered_across_suite():
    used: set[str] = set()
    for wish in COVERAGE_WISHES:
        trace, _events, _calls = run_grant(wish)
        used |= {t["call"]["tool"] for t in trace["turns"] if t["call"]}
    assert used == TOOLS, f"missing tools: {TOOLS - used}"
    # The town-mode tools must be among them.
    assert {"build_district", "place_road"} <= used


def test_city_script_uses_new_kinds_and_spawn_life():
    """The city script exercises the new building kinds and the 9th tool, and
    every call turn re-parses cleanly through parse_turn (grammar's JSON shape)."""
    trace, _events, _calls = run_grant("a bright city of towers and warehouses and cafes")
    call_turns = [t for t in trace["turns"] if t["call"] is not None]

    structure_kinds = {
        t["call"]["args"]["kind"]
        for t in call_turns
        if t["call"]["tool"] == "place_structure"
    }
    assert {"tower", "warehouse", "cafe"} <= structure_kinds

    life_calls = [t["call"] for t in call_turns if t["call"]["tool"] == "spawn_life"]
    assert life_calls, "city script should call spawn_life"
    for call in life_calls:
        args = call["args"]
        assert args["kind"] in ("carts", "birds", "fireflies")
        assert isinstance(args["count"], int) and 1 <= args["count"] <= 12
        assert set(args) == {"lat", "lon", "radius_deg", "kind", "count", "hue"}

    # Every emitted call turn is valid against the turn parser/grammar shape.
    for turn in call_turns:
        rendered = json.dumps({"thought": turn["thought"], "call": turn["call"]})
        obj, err = parse_turn(rendered)
        assert err is None and obj["call"]["tool"] == turn["call"]["tool"]


def test_town_script_grows_one_clustered_town():
    """The town-growth script lays a district, civic buildings (bank/market/
    house), a road, and carts — all clustered near one anchor — and every call
    turn re-parses cleanly through parse_turn."""
    trace, _events, _calls = run_grant("a new town district by the market")
    call_turns = [t for t in trace["turns"] if t["call"] is not None]
    tools_used = {t["call"]["tool"] for t in call_turns}
    assert "build_district" in tools_used
    assert "place_road" in tools_used

    structure_kinds = {
        t["call"]["args"]["kind"]
        for t in call_turns
        if t["call"]["tool"] == "place_structure"
    }
    assert {"bank", "market", "house"} <= structure_kinds

    # A district call carries exactly its spec'd args.
    district = next(t["call"]["args"] for t in call_turns
                    if t["call"]["tool"] == "build_district")
    assert set(district) == {"lat", "lon", "radius_deg", "density", "hue"}

    # A road is a path of 2..6 [lat,lon] waypoints.
    road = next(t["call"]["args"] for t in call_turns
                if t["call"]["tool"] == "place_road")
    assert set(road) == {"path"}
    assert 2 <= len(road["path"]) <= 6
    for pt in road["path"]:
        assert isinstance(pt, list) and len(pt) == 2

    # Everything clusters: all anchored coords sit within a few degrees of the
    # town's center (a town you build up close, not specks on a sphere).
    coords = []
    for t in call_turns:
        args = t["call"]["args"]
        if "lat" in args and "lon" in args:
            coords.append((args["lat"], args["lon"]))
    coords.extend(tuple(pt) for pt in road["path"])
    assert coords, "town script should anchor calls at coordinates"
    clat = sum(c[0] for c in coords) / len(coords)
    clon = sum(c[1] for c in coords) / len(coords)
    for lat, lon in coords:
        assert abs(lat - clat) <= 6 and abs(lon - clon) <= 6, (lat, lon)

    # Every emitted call turn is valid against the turn parser/grammar shape.
    for turn in call_turns:
        rendered = json.dumps({"thought": turn["thought"], "call": turn["call"]})
        obj, err = parse_turn(rendered)
        assert err is None and obj["call"]["tool"] == turn["call"]["tool"]


def test_spawn_life_scripts_parse_cleanly():
    """Each life kind appears in some mock script and round-trips through
    parse_turn without error."""
    life_kinds: set[str] = set()
    for wish in (
        "a bright city of towers and warehouses and cafes",  # carts
        "an ancient forest full of fireflies",               # fireflies
        "more stars and another moon over the night sky",    # birds
    ):
        trace, _events, _calls = run_grant(wish)
        for turn in trace["turns"]:
            call = turn["call"]
            if call and call["tool"] == "spawn_life":
                life_kinds.add(call["args"]["kind"])
                rendered = json.dumps({"thought": turn["thought"], "call": call})
                obj, err = parse_turn(rendered)
                assert err is None, f"spawn_life turn failed to parse: {err}"
    assert life_kinds == {"carts", "birds", "fireflies"}


def test_inscriptions_clamped_for_long_wish():
    trace, _events, _calls = run_grant(DIVERSE_WISHES["long"])
    inscriptions = [
        t["call"]["args"]
        for t in trace["turns"]
        if t["call"] and t["call"]["tool"] == "inscribe_wish"
    ]
    assert inscriptions, "village script should inscribe the wish"
    for args in inscriptions:
        assert len(args["text"]) <= 90
        assert args["style"] in ("orbit", "stone")


def test_determinism():
    wish = "a lighthouse for everyone still at sea"
    trace_a, _ev_a, _ca = run_grant(wish)
    trace_b, _ev_b, _cb = run_grant(wish)
    assert trace_a["reading"] == trace_b["reading"]
    assert trace_a["turns"] == trace_b["turns"]
    assert trace_a["epitaph"] == trace_b["epitaph"]


# --------------------------------------------------------------------------
# Emit event ordering
# --------------------------------------------------------------------------

def test_emit_ordering():
    trace, events, calls = run_grant("an ancient forest full of fireflies")

    assert events, "grant must emit events"
    types = [e["type"] for e in events]
    # The planner emits ONLY thought_token; tool_call/world_delta belong to the
    # queue worker (which holds the wish_id and the engine's canonical args).
    assert set(types) == {"thought_token"}

    # Reading streams before anything else.
    assert types[0] == "thought_token"

    # Every call turn reached the act callback, in order.
    call_turns = [t["call"] for t in trace["turns"] if t["call"] is not None]
    assert call_turns == calls
    assert len(calls) > 0

    # Each executed call's turn carries the thought that streamed before it.
    for turn in trace["turns"]:
        if turn["call"] is not None:
            assert turn["thought"]


def test_reading_streamed_through_emit():
    trace, events, _calls = run_grant("a quiet sea under patient stars")
    streamed = "".join(e["text"] for e in events if e["type"] == "thought_token")
    assert trace["reading"] in streamed
    # Turn thoughts stream too.
    for turn in trace["turns"]:
        if turn["thought"]:
            assert turn["thought"] in streamed


# --------------------------------------------------------------------------
# Malformed-JSON recovery
# --------------------------------------------------------------------------

def test_reask_recovers_malformed_turns():
    backend = FlakyTurnBackend()
    trace, _events, calls = run_grant("a mountain to watch over the valley", backend)

    assert backend.first_attempts > 0
    assert backend.reasks == backend.first_attempts  # one re-ask per failure
    # Every turn recovered — no skip records.
    skips = [
        t for t in trace["turns"]
        if t["call"] is None and (t["observation"] or "").startswith("error:")
    ]
    assert not skips
    assert calls, "recovered turns still execute their calls"
    assert trace["epitaph"]

    # Recovered plan matches the clean mock plan exactly (determinism holds
    # through the re-ask path).
    clean, _e, _c = run_grant("a mountain to watch over the valley")
    assert trace["turns"] == clean["turns"]


def test_always_broken_backend_fails_safe():
    trace, events, calls = run_grant("anything at all", AlwaysBrokenBackend())

    assert calls == []
    assert all(e["type"] != "tool_call" for e in events)
    # Fallback reading is still streamed so watchers see signs of life.
    assert any(e["type"] == "thought_token" and e["text"].strip() for e in events)
    assert trace["reading"]  # fallback reading
    # Two consecutive skips, then the loop bails.
    assert len(trace["turns"]) == 2
    for turn in trace["turns"]:
        assert turn["call"] is None
        assert turn["observation"].startswith("error:")
    assert 0 < len(trace["epitaph"]) <= 120  # fallback epitaph
    json.dumps(trace)


def test_parse_turn_unit():
    good_call = '{"thought": "lift", "call": {"tool": "raise_terrain", "args": {"lat": 1}}}'
    obj, err = parse_turn(good_call)
    assert err is None and obj["call"]["tool"] == "raise_terrain"

    good_done = '{"thought": "rest", "done": true, "epitaph": "made."}'
    obj, err = parse_turn(good_done)
    assert err is None and obj["done"] is True and obj["epitaph"] == "made."

    for bad in (
        "",
        "not json",
        '{"thought": "x"}',                                            # neither call nor done
        '{"thought": "x", "done": true}',                              # no epitaph
        '{"thought": "x", "call": {"tool": "summon_kraken", "args": {}}}',  # unknown tool
        '{"thought": "x", "call": {"tool": "set_sky"}}',               # no args
        '{"call": {"tool": "set_sky", "args": {}}}',                   # no thought
        '{"thought": "x", "done": true, "epitaph": "e", "call": {"tool": "set_sky", "args": {}}}',
    ):
        obj, err = parse_turn(bad)
        assert obj is None and err, f"should reject: {bad!r}"

    # Length clamps.
    long_thought = json.dumps({"thought": "x" * 500, "done": True, "epitaph": "y" * 500})
    obj, err = parse_turn(long_thought)
    assert err is None
    assert len(obj["thought"]) == 160 and len(obj["epitaph"]) == 120


# --------------------------------------------------------------------------
# Moderation judge — default-deny
# --------------------------------------------------------------------------

def test_judge_allows_benign():
    verdict = asyncio.run(judge("a meadow of flowers for my mother", MockBackend(delay=0.0)))
    assert verdict == {"allowed": True, "category": None}


def test_judge_denies_keyword():
    verdict = asyncio.run(judge("i want to kill everyone here", MockBackend(delay=0.0)))
    assert verdict["allowed"] is False
    assert verdict["category"] == "violence"


def test_judge_default_denies_on_exception():
    verdict = asyncio.run(judge("a meadow of flowers", ExplodingBackend()))
    assert verdict["allowed"] is False
    assert verdict["category"]


def test_judge_default_denies_on_garbage():
    verdict = asyncio.run(judge("a meadow of flowers", GarbageVerdictBackend()))
    assert verdict["allowed"] is False
    assert verdict["category"]


# --------------------------------------------------------------------------
# Grammar file sanity
# --------------------------------------------------------------------------

def test_grammar_sections():
    assert TURN_GRAMMAR.strip() and MODERATION_GRAMMAR.strip()
    assert "root ::=" in TURN_GRAMMAR and "root ::=" in MODERATION_GRAMMAR
    for tool in TOOLS:
        assert f'"\\"{tool}\\""' in TURN_GRAMMAR, f"tool {tool} not enum-locked"
    # The moderation key must appear ONLY in the moderation grammar — backends
    # without grammar enforcement rely on it to tell the two apart.
    assert '"allowed"' in MODERATION_GRAMMAR
    assert '"allowed"' not in TURN_GRAMMAR


def test_mood_inference():
    assert prompts.infer_mood("i miss the sea")[0] == "melancholy"
    assert prompts.infer_mood("a festival of dancing lights")[0] == "joyful"
    assert prompts.infer_mood("a haunted hollow where shadows whisper")[0] == "eerie"
    assert prompts.infer_mood("a planet-sized sock")[0] == "absurd"
    assert prompts.infer_mood("a small house by a road")[0] == "still"