nextcraft/apps/ai-service/tests/grading/test_engine.py

"""GradingEngine tests (Task 3-2-01, REQ-3-004, G-4 binding).

Contract under test:
    - G-4 gate FIRST: seq gaps OR integrity-flagged trace →
      verdict=UNGRADABLE_TRACE_INCOMPLETE with the gap list surfaced in
      scores; the LLM is NEVER called on those paths (asserted). Empty
      trace → UNGRADABLE_EMPTY_TRACE. All are first-class GradeRecords
      (returned + persisted), not exceptions.
    - Complete trace: prompt carries the digest JSON but NO raw trace
      material (D-028) — a distinctive marker planted in a command payload
      must be absent from every message the provider received.
    - Malformed LLM JSON → the D-020 retry path recovers (scripted bad
      first, good second) → graded.
    - Re-grade overwrites the stored record (GradeStore upsert).

Zero network: everything runs against the mock provider (conftest rule).
"""

from __future__ import annotations

import json
from datetime import UTC, datetime, timedelta
from pathlib import Path

import pytest
from pydantic import BaseModel, ValidationError

from ai_service.grading.engine import GradingEngine, RubricScore
from ai_service.grading.store import SQLiteGradeStore
from ai_service.llm.types import Message
from ai_service.telemetry.ingest import TraceIntegrityMap
from ai_service.telemetry.models import TelemetryEvent
from ai_service.telemetry.store import SQLiteTraceStore
from ai_service.variants.store import SQLiteVariantStore

T0 = datetime(2026, 9, 12, 1, 0, 0, tzinfo=UTC)
RAW_MARKER = "SECRET-COMMAND-MARKER-7f3a"


class RubricJSON:
    """Canonical well-formed rubric payload (matches engine.RubricScore)."""

    PAYLOAD: dict = {
        "criteria": {
            "process_quality": 3,
            "correctness": 4,
            "debugging_discipline": 3,
            "test_usage": 4,
        },
        "strengths": ["tight edit-test loops throughout"],
        "gaps": ["final commit discipline loose"],
        "verdict": "mastered",
    }


class RecordingProvider:
    """Mock LLMProvider: scripted replies + captured requests (no network).

    Subclass-composes ai_service.llm.mock.MockProvider so the conftest
    cloud-free guard philosophy holds — but records every message list
    and counts calls, which the stock mocks don't do.
    """

    def __init__(self) -> None:
        from ai_service.llm.mock import MockProvider

        self._mock = MockProvider()
        self.replies: list[str] = []  # consumed left-to-right; last repeats
        self.requests: list[list[Message]] = []
        self.calls = 0

    async def chat(self, messages, *, model, temperature=0.7, response_format=None):
        self.calls += 1
        self.requests.append([Message(role=m.role, content=m.content) for m in messages])
        if self.replies:
            reply = self.replies.pop(0)
        else:
            reply = json.dumps(RubricJSON.PAYLOAD)
        return reply

    async def stream_chat(self, messages, *, model, temperature=0.7, response_format=None):
        yield await self.chat(
            messages, model=model, temperature=temperature, response_format=response_format
        )


def _event(
    seq: int,
    kind: str,
    payload: dict | None = None,
    offset_s: float = 0.0,
    *,
    learner: str = "engine-learner",
    task: str = "engine-task",
) -> TelemetryEvent:
    return TelemetryEvent(
        learner_id=learner,
        task_id=task,
        seq=seq,
        kind=kind,
        payload=payload or {},
        ts=T0 + timedelta(seconds=offset_s),
        sandbox_id="sbx-engine",
    )


def _complete_trace(marker_command: bool = True) -> list[TelemetryEvent]:
    """A healthy iterative trace; seq 0..N contiguous (no gaps).

    `marker_command` plants a distinctive string inside a command payload
    — the D-028 assertion target: raw trace material must never reach the
    provider, so the marker must not appear in any captured request.
    """
    if marker_command:
        cmd_payload = {"cmd": f"echo {RAW_MARKER} && pytest -q"}
    else:
        cmd_payload = {"cmd": "pytest -q"}
    return [
        _event(0, "activity", {"state": "starting"}, 0.0),
        _event(1, "file_diff", {"path": "a.py", "added": 12}, 10.0),
        _event(2, "command", cmd_payload, 20.0),
        _event(3, "test_result", {"passed": False, "exit_code": 1}, 25.0),
        _event(4, "file_diff", {"path": "a.py", "added": 4, "removed": 2}, 40.0),
        _event(5, "command", {"cmd": "pytest -q"}, 60.0),
        _event(6, "test_result", {"passed": True, "exit_code": 0}, 65.0),
    ]


@pytest.fixture
def trace_store(tmp_path: Path) -> SQLiteTraceStore:
    store = SQLiteTraceStore(db_path=tmp_path / "traces.db")
    yield store
    store.close()


@pytest.fixture
def grade_store(tmp_path: Path) -> SQLiteGradeStore:
    store = SQLiteGradeStore(db_path=tmp_path / "grades.db")
    yield store
    store.close()


@pytest.fixture
def integrity() -> TraceIntegrityMap:
    return TraceIntegrityMap()


@pytest.fixture
def provider() -> RecordingProvider:
    return RecordingProvider()


@pytest.fixture
def engine(trace_store, grade_store, integrity, provider) -> GradingEngine:
    return GradingEngine(
        trace_store, grade_store, integrity, provider, model="gemma4:31b"
    )


def _ingest(store: SQLiteTraceStore, events: list[TelemetryEvent]) -> None:
    for e in events:
        store.append(e)


class TestRubricScoreModel:
    """The D-020 schema contract — engine-side validation rules."""

    def _rubric(self, **overrides) -> dict:
        payload = json.loads(json.dumps(RubricJSON.PAYLOAD))  # deep copy
        payload.update(overrides)
        return payload

    def test_well_formed_payload_validates(self):
        score = RubricScore.model_validate(RubricJSON.PAYLOAD)
        assert score.criteria["process_quality"] == 3
        assert score.verdict == "mastered"

    def test_unknown_criterion_rejected(self):
        bad = self._rubric()
        bad["criteria"]["extra_criterion"] = 2
        with pytest.raises(ValidationError, match="criteria keys must be exactly"):
            RubricScore.model_validate(bad)

    def test_missing_criterion_rejected(self):
        bad = self._rubric()
        del bad["criteria"]["test_usage"]
        with pytest.raises(ValidationError, match="criteria keys must be exactly"):
            RubricScore.model_validate(bad)

    def test_out_of_range_score_rejected(self):
        bad = self._rubric()
        bad["criteria"]["process_quality"] = 5
        with pytest.raises(ValidationError, match="must be within 0-4"):
            RubricScore.model_validate(bad)

    def test_unknown_verdict_rejected(self):
        bad = self._rubric()
        bad["verdict"] = "excellent"
        with pytest.raises(ValidationError, match="verdict must be one of"):
            RubricScore.model_validate(bad)


class TestCompleteTraceGrading:
    async def test_graded_returns_validated_rubric_persisted_and_returned(
        self, engine, trace_store, grade_store, provider
    ):
        _ingest(trace_store, _complete_trace())

        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "GRADED"
        assert record.model == "gemma4:31b"
        assert record.scores == RubricJSON.PAYLOAD
        assert record.scores["criteria"]["process_quality"] == 3
        assert record.digest.get("error_fix_cycles") == 1
        assert record.digest.get("final_test_status") == "pass"
        # digest persisted for auditability (D-028 reproducible input)
        assert record.digest.get("event_count") == 7

        stored = grade_store.get("engine-learner", "engine-task")
        assert stored is not None
        assert stored.scores == record.scores
        assert stored.verdict == "GRADED"

    async def test_prompt_contains_digest_but_no_raw_trace(self, engine, trace_store, provider):
        """D-028: the provider saw the digest JSON but NEVER the raw trace.

        The marker was planted inside a command payload (seq 2); it must be
        absent from every message of every request, while the digest marker
        line + a JSON object with the digest's field names must be present.
        """
        _ingest(trace_store, _complete_trace(marker_command=True))

        await engine.grade("engine-learner", "engine-task")

        assert provider.calls == 1  # happy path: exactly one LLM call
        all_messages = [m for req in provider.requests for m in req]
        prompt_text = "\n".join(m.content for m in all_messages)

        assert RAW_MARKER not in prompt_text, (
            "raw trace material reached the LLM prompt (D-028 violation)"
        )
        # digest presence: marker line + a JSON object carrying digest fields
        assert "PROCESS TRACE DIGEST (JSON):" in prompt_text
        assert '"error_fix_cycles"' in prompt_text
        assert '"final_test_status"' in prompt_text
        # learner anonymity: identity strings never reach the LLM either
        assert "engine-learner" not in prompt_text
        assert "engine-task" not in prompt_text

    async def test_system_prompt_has_rubric_and_a4_advisory(
        self, engine, trace_store, provider
    ):
        _ingest(trace_store, _complete_trace())

        await engine.grade("engine-learner", "engine-task")

        system = provider.requests[0][0]
        assert system.role == "system"
        assert "process_quality" in system.content
        assert "correctness" in system.content
        assert "debugging_discipline" in system.content
        assert "test_usage" in system.content
        # a-4: churn-without-test-progress is a process-quality NEGATIVE
        assert "churn is not" in system.content
        # 0-4 level anchors present for each criterion
        assert "ADVISORY" in system.content


class TestGateFirst:
    """G-4 binding: the gate fires BEFORE any grading work."""

    async def test_gapped_trace_ungradable_llm_not_called(
        self, engine, trace_store, provider, grade_store
    ):
        # seqs 0,1,3 stored — seq 2 missing (mirrors the task spec example)
        events = _complete_trace()
        gapped = [e for e in events if e.seq != 2]
        _ingest(trace_store, gapped)

        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "UNGRADABLE_TRACE_INCOMPLETE"
        assert provider.calls == 0, "LLM was called despite a gapped trace (G-4)"
        # gap list surfaced in scores
        assert record.scores["missing_seqs"] == [2]
        assert record.scores["integrity_flag"] is None
        # first-class record: persisted, retrievable, empty digest
        stored = grade_store.get("engine-learner", "engine-task")
        assert stored is not None
        assert stored.verdict == "UNGRADABLE_TRACE_INCOMPLETE"
        assert stored.scores["missing_seqs"] == [2]
        assert stored.digest == {}
        assert stored.model == "none"

    async def test_integrity_flagged_trace_ungradable_llm_not_called(
        self, engine, trace_store, integrity, provider
    ):
        _ingest(trace_store, _complete_trace())
        integrity.mark("engine-learner", "engine-task", "INCOMPLETE_FLOODED")

        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "UNGRADABLE_TRACE_INCOMPLETE"
        assert provider.calls == 0, "LLM was called despite integrity flag (G-4)"
        assert record.scores["integrity_flag"] == "INCOMPLETE_FLOODED"
        assert record.scores["missing_seqs"] == []  # complete rows, but untrusted

    async def test_flag_wins_even_with_gaps(self, engine, trace_store, integrity, provider):
        """Both gate signals set: the flag reason is surfaced (both listed)."""
        events = _complete_trace()
        _ingest(trace_store, [e for e in events if e.seq != 2])
        integrity.mark("engine-learner", "engine-task", "INCOMPLETE_FLOODED")

        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "UNGRADABLE_TRACE_INCOMPLETE"
        assert provider.calls == 0
        assert record.scores["integrity_flag"] == "INCOMPLETE_FLOODED"
        assert record.scores["missing_seqs"] == [2]

    async def test_empty_trace_ungradable(self, engine, provider, grade_store):
        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "UNGRADABLE_EMPTY_TRACE"
        assert provider.calls == 0, "LLM was called for an empty trace"
        assert record.scores == {"integrity_flag": None, "missing_seqs": []}
        assert record.digest == {}
        stored = grade_store.get("engine-learner", "engine-task")
        assert stored is not None
        assert stored.verdict == "UNGRADABLE_EMPTY_TRACE"


class TestD020RetryPath:
    async def test_malformed_first_response_recovers_via_retry(
        self, engine, trace_store, provider
    ):
        """First reply invalid (wrong shape, fenced), second valid → graded.

        This exercises the D-020 layer-4 path INSIDE the engine through the
        real agents/structured.py — the engine does not reimplement it.
        """
        _ingest(trace_store, _complete_trace())
        provider.replies = [
            '```json\n{"summary": "wrong shape"}\n```',  # layer 3 rejects
            json.dumps(RubricJSON.PAYLOAD),               # retry recovers
        ]

        record = await engine.grade("engine-learner", "engine-task")

        assert provider.calls == 2  # bounded to exactly one retry (layer 4)
        assert record.verdict == "GRADED"
        assert record.scores == RubricJSON.PAYLOAD
        # the retry feedback must carry the validation error back
        retry_request = provider.requests[1]
        retry_text = " ".join(m.content for m in retry_request)
        assert "previous response was invalid" in retry_text

    async def test_persistently_malformed_raises_after_bounded_retry(
        self, engine, trace_store, provider
    ):
        from ai_service.agents.structured import StructuredOutputError

        _ingest(trace_store, _complete_trace())
        provider.replies = ["not json at all", "still not json"]

        with pytest.raises(StructuredOutputError):
            await engine.grade("engine-learner", "engine-task")
        assert provider.calls == 2  # bounded: never more than one retry


class TestRegrade:
    async def test_regrade_overwrites_stored_record(
        self, engine, trace_store, grade_store, provider
    ):
        _ingest(trace_store, _complete_trace())

        first = await engine.grade("engine-learner", "engine-task")
        assert first.verdict == "GRADED"

        # second grade: the provider now scripts a DIFFERENT rubric outcome
        updated = json.loads(json.dumps(RubricJSON.PAYLOAD))
        updated["criteria"]["process_quality"] = 1
        updated["verdict"] = "not_yet"
        provider.replies = [json.dumps(updated)]

        second = await engine.grade("engine-learner", "engine-task")

        assert second.scores == updated
        assert second.created_at >= first.created_at
        # upsert, not append: exactly one row, holding the LATEST grade
        grades = grade_store.list_for_learner("engine-learner")
        assert len(grades) == 1
        assert grades[0].scores == updated
        assert grades[0].verdict == "GRADED"

    async def test_regrade_after_gate_outcome_replaces_it(
        self, engine, trace_store, grade_store, provider, integrity
    ):
        """A later successful regrade wholesale-replaces a gate record —
        the GradeStore upsert contract applied across verdict kinds."""
        # first: gapped → gate record persisted
        events = _complete_trace()
        _ingest(trace_store, [e for e in events if e.seq != 2])
        gated = await engine.grade("engine-learner", "engine-task")
        assert gated.verdict == "UNGRADABLE_TRACE_INCOMPLETE"

        # gap healed (late delivery): same pair regrades cleanly
        trace_store.append(next(e for e in events if e.seq == 2))
        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "GRADED"
        stored = grade_store.get("engine-learner", "engine-task")
        assert stored is not None
        assert stored.verdict == "GRADED"
        assert stored.scores == RubricJSON.PAYLOAD
        assert len(grade_store.list_for_learner("engine-learner")) == 1


class TestBoundary:
    def test_engine_module_imports_no_fastapi(self):
        """D-027: the engine module must not import fastapi (transitively
        beyond the sanctioned agents/structured + llm + telemetry + prompts)."""
        import sys

        import ai_service.grading.engine as engine_mod

        code = open(engine_mod.__file__).read()
        assert "fastapi" not in code
        assert "from ..api" not in code and "from ..api." not in code
        # the one sanctioned agents import is the module-direct structured defense
        assert "from ..agents.structured import" in code
        assert "from ..agents import" not in code
        # sanity: the module really is loaded (no import cycle surprises)
        assert engine_mod.__name__ in sys.modules


class TestModelValidationSmoke:
    """RubricScore as a plain pydantic model (D-020 layer-3 target type)."""

    def test_rubric_score_is_frozen_shape_for_scores_dict(self):
        score = RubricScore.model_validate(RubricJSON.PAYLOAD)
        dumped = score.model_dump()
        assert set(dumped["criteria"]) == {
            "process_quality",
            "correctness",
            "debugging_discipline",
            "test_usage",
        }
        assert isinstance(score, BaseModel)


class TestVariantAnchorsShipment:
    """Phase 4 MH#4: variant anchors + seed ship into grading (a-5 same bar)."""

    @pytest.fixture
    def variant_engine(self, trace_store, grade_store, integrity, provider, tmp_path):
        store = SQLiteVariantStore(db_path=tmp_path / "variants.db")
        yield GradingEngine(
            trace_store,
            grade_store,
            integrity,
            provider,
            model="gemma4:31b",
            variant_store=store,
        ), store
        store.close()

    @pytest.fixture
    def plain_engine(self, trace_store, grade_store, integrity, provider):
        return GradingEngine(
            trace_store, grade_store, integrity, provider, model="gemma4:31b"
        )

    async def test_variant_task_grades_with_anchors_and_seed(
        self, variant_engine, trace_store, provider
    ):
        engine, vstore = variant_engine
        _ingest(trace_store, _complete_trace())

        # A stored variant whose task_id matches the graded trace.
        from datetime import UTC
        from datetime import datetime as dt

        from ai_service.variants.store import VariantRecord

        vstore.save(
            VariantRecord(
                learner_id="engine-learner",
                task_id="engine-task",
                template_id="tpl-llm-judge",
                seed="cafe" * 16,
                params={"domain": "tutoring"},
                statement="Scripted statement long enough to be legal.",
                starter_files={"README.md": "x"},
                created_at=dt.now(UTC),
            )
        )

        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "GRADED"
        assert record.variant_seed == "cafe" * 16  # D-029 stamped
        user_msgs = [m.content for req in provider.requests for m in req if m.role == "user"]
        assert any("Expected effort envelope" in u for u in user_msgs)
        assert any("tpl-llm-judge" in u for u in user_msgs)
        assert any("expected_min_test_runs" in u for u in user_msgs)

    async def test_non_variant_task_has_no_anchors(
        self, variant_engine, trace_store, provider
    ):
        engine, _ = variant_engine
        _ingest(trace_store, _complete_trace())

        record = await engine.grade("engine-learner", "engine-task")

        assert record.verdict == "GRADED"
        assert record.variant_seed is None
        user_msgs = [m.content for req in provider.requests for m in req if m.role == "user"]
        assert not any("Expected effort envelope" in u for u in user_msgs)

    async def test_plain_engine_stays_variant_blind(
        self, plain_engine, trace_store
    ):
        _ingest(trace_store, _complete_trace())
        record = await plain_engine.grade("engine-learner", "engine-task")
        assert record.verdict == "GRADED"
        assert record.variant_seed is None