Agent/tests/test_recursive_validation_core.py

import json
import tempfile
import unittest

from pydantic import ValidationError

from cyber_agent.core.validation import (
    LLMValidator,
    ValidationResult,
    build_validation_packet,
    parse_validation_result,
    validation_error,
)
from cyber_agent.trace.models import Message, Trace
from cyber_agent.trace.store import FileSystemTraceStore


class FakeLLM:
    def __init__(self, response=None, error=None):
        self.response = response
        self.error = error
        self.calls = []

    async def __call__(self, **kwargs):
        self.calls.append(kwargs)
        if self.error:
            raise self.error
        return self.response


def valid_response(scope="task"):
    return {
        "content": json.dumps({
            "outcome": "passed",
            "scope": scope,
            "reason": "The persisted evidence satisfies every criterion.",
            "issues": [],
            "retry_from": None,
        }),
        "tool_calls": None,
        "prompt_tokens": 120,
        "completion_tokens": 30,
        "cost": 0.012,
        "finish_reason": "stop",
    }


class ValidationModelTest(unittest.TestCase):
    def test_result_enforces_outcome_invariants(self):
        passed = ValidationResult(
            validator_trace_id="validator",
            evaluated_trace_id="child",
            outcome="passed",
            scope="task",
            reason="all criteria passed",
            issues=[],
            retry_from=None,
        )
        self.assertEqual("passed", passed.outcome)

        invalid = [
            {"outcome": "passed", "issues": ["unexpected"], "retry_from": None},
            {"outcome": "failed", "issues": [], "retry_from": "evidence"},
            {"outcome": "failed", "issues": ["gap"], "retry_from": None},
            {"outcome": "error", "issues": ["broken"], "retry_from": "output"},
        ]
        for fields in invalid:
            with self.subTest(fields=fields), self.assertRaises(ValidationError):
                ValidationResult(
                    validator_trace_id="validator",
                    evaluated_trace_id="child",
                    scope="task",
                    reason="invalid combination",
                    **fields,
                )

    def test_parser_injects_ids_and_rejects_model_owned_ids(self):
        parsed = parse_validation_result(
            valid_response()["content"],
            validator_trace_id="validator",
            evaluated_trace_id="child",
            expected_scope="task",
        )
        self.assertEqual("validator", parsed.validator_trace_id)
        self.assertEqual("child", parsed.evaluated_trace_id)

        forged = json.loads(valid_response()["content"])
        forged["evaluated_trace_id"] = "forged"
        with self.assertRaises(ValidationError):
            parse_validation_result(
                json.dumps(forged),
                validator_trace_id="validator",
                evaluated_trace_id="child",
                expected_scope="task",
            )
        with self.assertRaisesRegex(ValueError, "expected 'task'"):
            parse_validation_result(
                valid_response(scope="root")["content"],
                validator_trace_id="validator",
                evaluated_trace_id="child",
                expected_scope="task",
            )

    def test_error_result_is_deterministic_and_fail_closed(self):
        result = validation_error(
            validator_trace_id="validator",
            evaluated_trace_id="child",
            scope="root",
            reason=" invalid JSON ",
        )
        self.assertEqual("error", result.outcome)
        self.assertEqual(["invalid JSON"], result.issues)
        self.assertIsNone(result.retry_from)

    def test_packet_keeps_contract_and_newest_main_path_within_limit(self):
        trajectory = [
            {
                "sequence": index,
                "role": "tool",
                "name": "read_file",
                "content": f"marker-{index}-" + ("x" * 300),
                "reasoning_content": "must-not-leak",
            }
            for index in range(1, 8)
        ]
        trajectory.insert(6, {
            "sequence": 99,
            "role": "assistant",
            "content": "side-branch-secret",
            "branch_type": "compression",
        })
        packet = build_validation_packet(
            validation_scope="task",
            task_brief={
                "objective": "check result",
                "completion_criteria": ["criterion-kept"],
                "expected_outputs": ["output-kept"],
            },
            task_report={"summary": "report-kept"},
            trajectory=trajectory,
            max_chars=1_000,
        )
        self.assertLessEqual(len(packet), 1_000)
        self.assertIn("criterion-kept", packet)
        self.assertIn("output-kept", packet)
        self.assertIn("report-kept", packet)
        self.assertIn("marker-7", packet)
        self.assertNotIn("marker-1", packet)
        self.assertNotIn("reasoning_content", packet)
        self.assertNotIn("side-branch-secret", packet)

    def test_packet_rejects_contract_that_alone_exceeds_limit(self):
        with self.assertRaisesRegex(ValueError, "contract exceeds"):
            build_validation_packet(
                validation_scope="task",
                task_brief={"objective": "x" * 1_000},
                trajectory=[],
                max_chars=100,
            )

    def test_packet_removes_reasoning_from_persisted_assistant_message(self):
        message = Message.create(
            trace_id="evaluated",
            role="assistant",
            sequence=1,
            content={
                "text": "observable answer",
                "reasoning_content": "hidden chain of thought",
            },
        )
        packet = build_validation_packet(
            validation_scope="task",
            trajectory=[message],
        )
        self.assertIn("observable answer", packet)
        self.assertNotIn("hidden chain of thought", packet)
        self.assertNotIn("reasoning_content", packet)


class LLMValidatorTest(unittest.IsolatedAsyncioTestCase):
    async def asyncSetUp(self):
        self.temp_dir = tempfile.TemporaryDirectory()
        self.store = FileSystemTraceStore(self.temp_dir.name)
        self.evaluated = Trace(
            trace_id="root@delegate-child",
            mode="agent",
            task="child task",
            uid="user-1",
            model="fake-model",
            current_goal_id="1.1",
            context={
                "agent_mode": "recursive",
                "agent_mode_revision": 2,
                "root_trace_id": "root",
                "agent_depth": 2,
            },
        )
        await self.store.create_trace(self.evaluated)
        self.trajectory = [
            Message.create(
                trace_id=self.evaluated.trace_id,
                role="user",
                sequence=1,
                content="produce checked evidence",
            ),
            Message.create(
                trace_id=self.evaluated.trace_id,
                role="tool",
                sequence=2,
                parent_sequence=1,
                content={"tool_name": "read_file", "result": "actual evidence"},
            ),
        ]

    async def asyncTearDown(self):
        self.temp_dir.cleanup()

    async def test_single_tool_free_call_persists_validator_trace_and_usage(self):
        llm = FakeLLM(valid_response())
        validator = LLMValidator(llm_call=llm, trace_store=self.store)
        run = await validator.validate(
            evaluated_trace=self.evaluated,
            trajectory=self.trajectory,
            scope="task",
            task_brief={
                "objective": "check result",
                "completion_criteria": ["must cite evidence"],
                "expected_outputs": ["one conclusion"],
            },
            task_report={"summary": "done", "evidence": [{"value": "actual"}]},
            validator_trace_id="validator-1",
        )

        self.assertEqual(1, len(llm.calls))
        call = llm.calls[0]
        self.assertEqual([], call["tools"])
        self.assertEqual(0, call["temperature"])
        self.assertEqual("fake-model", call["model"])
        self.assertEqual(2, len(call["messages"]))
        self.assertEqual("passed", run.result.outcome)
        self.assertEqual(120, run.prompt_tokens)
        self.assertEqual(30, run.completion_tokens)
        self.assertEqual(0.012, run.cost)

        trace = await self.store.get_trace("validator-1")
        self.assertEqual("completed", trace.status)
        self.assertEqual("validator", trace.agent_type)
        self.assertEqual(self.evaluated.trace_id, trace.parent_trace_id)
        self.assertEqual("validator", trace.context["created_by_tool"])
        self.assertEqual("root", trace.context["root_trace_id"])
        self.assertEqual(2, trace.context["agent_depth"])
        self.assertEqual([], trace.tools)
        messages = await self.store.get_main_path_messages("validator-1", 3)
        self.assertEqual(["system", "user", "assistant"], [m.role for m in messages])
        self.assertIn("actual evidence", messages[1].content)

    async def test_invalid_json_returns_error_without_correction_call(self):
        response = valid_response()
        response["content"] = "```json\n{}\n```"
        llm = FakeLLM(response)
        validator = LLMValidator(llm_call=llm, trace_store=self.store)
        run = await validator.validate(
            evaluated_trace=self.evaluated,
            trajectory=self.trajectory,
            scope="task",
            validator_trace_id="validator-invalid",
        )
        self.assertEqual(1, len(llm.calls))
        self.assertEqual("error", run.result.outcome)
        trace = await self.store.get_trace(run.trace_id)
        self.assertEqual("failed", trace.status)
        self.assertIn("Validator failed", trace.error_message)

    async def test_tool_call_is_rejected_without_dispatch(self):
        response = valid_response()
        response["tool_calls"] = [{
            "id": "forged",
            "type": "function",
            "function": {"name": "agent", "arguments": "{}"},
        }]
        llm = FakeLLM(response)
        validator = LLMValidator(llm_call=llm, trace_store=self.store)
        run = await validator.validate(
            evaluated_trace=self.evaluated,
            trajectory=self.trajectory,
            scope="task",
        )
        self.assertEqual("error", run.result.outcome)
        self.assertIn("attempted to call tools", run.result.reason)
        self.assertEqual(1, len(llm.calls))

    async def test_model_exception_is_recorded_once(self):
        llm = FakeLLM(error=RuntimeError("provider unavailable"))
        validator = LLMValidator(llm_call=llm, trace_store=self.store)
        run = await validator.validate(
            evaluated_trace=self.evaluated,
            trajectory=self.trajectory,
            scope="root",
            completion_criteria=["root done"],
            candidate_output="candidate",
        )
        self.assertEqual(1, len(llm.calls))
        self.assertEqual("error", run.result.outcome)
        self.assertIn("provider unavailable", run.result.reason)

    async def test_non_success_path_creates_trace_without_calling_llm(self):
        llm = FakeLLM(valid_response())
        validator = LLMValidator(llm_call=llm, trace_store=self.store)
        run = await validator.record_non_success(
            evaluated_trace=self.evaluated,
            scope="task",
            outcome="failed",
            reason="child stopped",
            issues=["execution stopped before evidence was produced"],
            retry_from="evidence",
            validator_trace_id="validator-stopped",
        )
        self.assertEqual([], llm.calls)
        self.assertEqual("failed", run.result.outcome)
        trace = await self.store.get_trace(run.trace_id)
        self.assertEqual("failed", trace.status)
        self.assertEqual(0, trace.total_tokens)
        self.assertEqual(1, trace.total_messages)

    async def test_oversized_fixed_input_fails_before_llm(self):
        llm = FakeLLM(valid_response())
        validator = LLMValidator(
            llm_call=llm,
            trace_store=self.store,
            max_input_chars=2_000,
        )
        run = await validator.validate(
            evaluated_trace=self.evaluated,
            trajectory=[],
            scope="task",
            task_brief={"objective": "x" * 3_000},
        )
        self.assertEqual([], llm.calls)
        self.assertEqual("error", run.result.outcome)
        self.assertIn("input could not be built", run.result.reason)


if __name__ == "__main__":
    unittest.main()