script-build-and-framework-tmp/agent/tests/test_orchestration_validation_policy.py

from __future__ import annotations

import asyncio
from types import SimpleNamespace

import pytest

from agent.orchestration.config import OrchestrationConfig
from agent.orchestration.evidence import (
    EvidenceBudgetExceeded,
    EvidenceOwnershipError,
    EvidenceProviderError,
    EvidenceProviderResult,
    EvidenceQuery,
)
from agent.orchestration.models import (
    AcceptanceCriterion,
    AgentRole,
    ArtifactRef,
    ArtifactSnapshot,
    AttemptSubmission,
    TaskAttempt,
    TaskRecord,
    TaskSpec,
    ValidationMode,
    ValidationPlan,
    ValidationVerdict,
)
from agent.orchestration.protocols import WorkerRunResult
from agent.orchestration.store import FileSystemArtifactStore, FileSystemTaskStore
from agent.orchestration.validation_policy import (
    DefaultValidationPolicy,
    DeterministicRuleResult,
    RuleBasedDeterministicValidator,
    ValidationContext,
)
from agent.orchestration.wiring import wire_orchestration

from test_coordinator_integration import (
    ROOT_TASK_SPEC,
    FakeExecutor,
    create_task,
    make_coordinator,
)


def make_context(root: str = "root") -> ValidationContext:
    task = TaskRecord(
        task_id="task-1",
        goal_id=None,
        parent_task_id=None,
        display_path="1",
        specs=[
            TaskSpec(
                version=1,
                objective="validate",
                acceptance_criteria=[AcceptanceCriterion("criterion", "validate result")],
            )
        ],
    )
    attempt = TaskAttempt(
        attempt_id="attempt-1",
        task_id=task.task_id,
        spec_version=1,
        worker_trace_id="worker-1",
        worker_preset="worker",
        execution_mode="new",
        snapshot_id="snapshot-1",
    )
    snapshot = ArtifactSnapshot(
        snapshot_id="snapshot-1",
        attempt_id=attempt.attempt_id,
        normalized_content={"summary": "done"},
        sha256="digest",
        artifact_refs=[],
        evidence_refs=[],
    )
    return ValidationContext(root, task, attempt, snapshot)


class StaticRule:
    def __init__(self, rule_id, verdict):
        self.rule_id = rule_id
        self.verdict = verdict

    def evaluate(self, context):
        assert context.snapshot.snapshot_id
        return DeterministicRuleResult(self.rule_id, self.verdict, "checked")


class AsyncRule(StaticRule):
    async def evaluate(self, context):
        await asyncio.sleep(0)
        return super().evaluate(context)


class ErrorRule:
    rule_id = "error"

    def evaluate(self, context):
        raise RuntimeError("broken rule")


class DeterministicPolicy:
    def __init__(self, verdict=ValidationVerdict.PASSED):
        self.verdict = verdict

    def plan(self, context):
        return ValidationPlan(
            mode=ValidationMode.DETERMINISTIC,
            validator_preset=None,
            rule_ids=[item.criterion_id for item in context.task.current_spec.acceptance_criteria],
        )


class FixedPolicy:
    def __init__(self, plan):
        self.fixed_plan = plan

    def plan(self, context):
        return self.fixed_plan


class InvalidBlankPolicy:
    def plan(self, context):
        plan = ValidationPlan()
        object.__setattr__(plan, "validator_preset", " ")
        return plan


@pytest.mark.asyncio
async def test_default_policy_and_validation_plan_roundtrip():
    plan = DefaultValidationPolicy("careful-validator").plan(make_context())
    assert plan.mode == ValidationMode.AGENT
    assert plan.validator_preset == "careful-validator"
    assert ValidationPlan.from_dict(plan.to_dict()) == plan

    deterministic = ValidationPlan(
        mode=ValidationMode.DETERMINISTIC,
        validator_preset=None,
        rule_ids=["schema", "checksum"],
        max_evidence_queries=2,
    )
    assert ValidationPlan.from_dict(deterministic.to_dict()) == deterministic
    assert ValidationPlan.from_dict({"mode": "deterministic", "rule_ids": []}).validator_preset is None
    with pytest.raises(ValueError, match="cannot contain"):
        ValidationPlan(rule_ids=["unexpected"])
    with pytest.raises(ValueError, match="cannot be blank"):
        ValidationPlan(validator_preset=" ")
    for kwargs in (
        {"max_evidence_queries": True},
        {"max_evidence_queries": 1.5},
        {"max_evidence_items_per_query": False},
        {"evidence_timeout_seconds": True},
    ):
        with pytest.raises(ValueError):
            ValidationPlan(**kwargs)
    with pytest.raises(ValueError, match="cannot select"):
        ValidationPlan(mode=ValidationMode.DETERMINISTIC)
    with pytest.raises(ValueError, match="unique"):
        ValidationPlan(
            mode=ValidationMode.DETERMINISTIC,
            validator_preset=None,
            rule_ids=["same", "same"],
        )


@pytest.mark.asyncio
async def test_deterministic_rules_aggregate_failures_and_errors():
    context = make_context()
    validator = RuleBasedDeterministicValidator(
        [
            StaticRule("pass", ValidationVerdict.PASSED),
            AsyncRule("fail", ValidationVerdict.FAILED),
            ErrorRule(),
        ]
    )
    result = await validator.validate(
        context,
        ValidationPlan(
            mode=ValidationMode.DETERMINISTIC,
            validator_preset=None,
            rule_ids=["pass", "missing", "error", "fail"],
        ),
    )
    assert result.verdict == ValidationVerdict.FAILED
    assert len(result.errors) == 2
    assert result.to_dict()["verdict"] == "failed"

    empty = await validator.validate(
        context,
        ValidationPlan(mode=ValidationMode.DETERMINISTIC, validator_preset=None),
    )
    assert empty.verdict == ValidationVerdict.INCONCLUSIVE
    with pytest.raises(ValueError, match="deterministic plan"):
        await validator.validate(context, ValidationPlan())
    with pytest.raises(ValueError, match="Duplicate"):
        RuleBasedDeterministicValidator(
            [
                StaticRule("same", ValidationVerdict.PASSED),
                StaticRule("same", ValidationVerdict.PASSED),
            ]
        )


@pytest.mark.asyncio
async def test_coordinator_freezes_default_agent_plan(tmp_path):
    executor = FakeExecutor([ValidationVerdict.PASSED])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    task_id = await create_task(coordinator, "default agent policy")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    report = (await store.load("root")).validations[result.validation_id]

    assert executor.validator_calls == 1
    assert report.validation_plan.mode == ValidationMode.AGENT
    assert report.validation_plan.validator_preset == "validator"
    assert ValidationPlan.from_dict(report.validation_plan) == report.validation_plan


@pytest.mark.asyncio
@pytest.mark.parametrize("verdict", [ValidationVerdict.PASSED, ValidationVerdict.FAILED])
async def test_coordinator_executes_deterministic_plan_without_agent_validator(tmp_path, verdict):
    executor = FakeExecutor([])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    coordinator.validation_policy = DeterministicPolicy()
    coordinator.deterministic_validator = RuleBasedDeterministicValidator([StaticRule("c1", verdict)])
    task_id = await create_task(coordinator, "deterministic")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    report = (await store.load("root")).validations[result.validation_id]

    assert executor.worker_calls == 1
    assert executor.validator_calls == 0
    assert report.validation_plan.mode == ValidationMode.DETERMINISTIC
    assert report.verdict == verdict
    assert report.criterion_results[0].criterion_id == "c1"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("preflight_verdict", "validator_calls", "final_verdict"),
    [
        (ValidationVerdict.FAILED, 0, ValidationVerdict.FAILED),
        (ValidationVerdict.PASSED, 1, ValidationVerdict.PASSED),
    ],
)
async def test_agent_validation_skips_model_on_hard_preflight_failure(
    tmp_path, preflight_verdict, validator_calls, final_verdict
):
    executor = FakeExecutor([ValidationVerdict.PASSED])
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    coordinator.validation_policy = FixedPolicy(ValidationPlan(preflight_rule_ids=("preflight",)))
    coordinator.deterministic_validator = RuleBasedDeterministicValidator([StaticRule("preflight", preflight_verdict)])
    task_id = await create_task(coordinator, "hybrid preflight")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    report = (await store.load("root")).validations[result.validation_id]

    assert executor.validator_calls == validator_calls
    assert report.verdict == final_verdict
    if preflight_verdict is ValidationVerdict.FAILED:
        assert "Deterministic preflight rejected" in report.summary
        assert report.execution_stats.total_tokens == 0


def test_wire_orchestration_preserves_legacy_positional_defaults(tmp_path):
    runner = SimpleNamespace(trace_store=None, task_coordinator=None)
    config = OrchestrationConfig(max_parallel_tasks=2)
    event_sink = object()

    coordinator = wire_orchestration(
        runner,
        FileSystemTaskStore(str(tmp_path)),
        FileSystemArtifactStore(str(tmp_path)),
        config,
        event_sink,
    )

    assert runner.task_coordinator is coordinator
    assert coordinator.config is config
    assert coordinator.event_sink is event_sink
    assert isinstance(coordinator.validation_policy, DefaultValidationPolicy)
    assert coordinator.deterministic_validator is None
    assert coordinator.evidence_provider is None


@pytest.mark.asyncio
async def test_wire_orchestration_applies_custom_validation_dependencies(tmp_path):
    runner = SimpleNamespace(trace_store=None, task_coordinator=None)
    policy = DeterministicPolicy()
    validator = RuleBasedDeterministicValidator([StaticRule("c1", ValidationVerdict.PASSED)])
    provider = EvidenceProvider()
    coordinator = wire_orchestration(
        runner,
        FileSystemTaskStore(str(tmp_path)),
        FileSystemArtifactStore(str(tmp_path)),
        validation_policy=policy,
        deterministic_validator=validator,
        evidence_provider=provider,
    )
    executor = FakeExecutor([])
    executor.coordinator = coordinator
    coordinator.set_executor(executor)
    await coordinator.ensure_ledger("root", ROOT_TASK_SPEC)
    task_id = await create_task(coordinator, "wired deterministic validation")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]
    report = (await coordinator.task_store.load("root")).validations[result.validation_id]

    assert runner.task_coordinator is coordinator
    assert coordinator.validation_policy is policy
    assert coordinator.deterministic_validator is validator
    assert coordinator.evidence_provider is provider
    assert executor.validator_calls == 0
    assert report.validation_plan.mode == ValidationMode.DETERMINISTIC
    assert report.verdict == ValidationVerdict.PASSED


@pytest.mark.asyncio
@pytest.mark.parametrize("rule_ids", [[], ["unknown"]])
async def test_deterministic_plan_must_cover_hard_criteria_without_unknowns(tmp_path, rule_ids):
    executor = FakeExecutor([])
    coordinator, _, _ = await make_coordinator(tmp_path, executor)
    coordinator.validation_policy = FixedPolicy(
        ValidationPlan(
            mode=ValidationMode.DETERMINISTIC,
            validator_preset=None,
            rule_ids=rule_ids,
        )
    )
    coordinator.deterministic_validator = RuleBasedDeterministicValidator([])
    task_id = await create_task(coordinator, "invalid deterministic mapping")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]

    assert "criteria" in (result.error or "")
    assert executor.validator_calls == 0


@pytest.mark.asyncio
async def test_coordinator_rejects_custom_policy_with_blank_agent_preset(tmp_path):
    executor = FakeExecutor([])
    coordinator, _, _ = await make_coordinator(tmp_path, executor)
    coordinator.validation_policy = InvalidBlankPolicy()
    task_id = await create_task(coordinator, "invalid preset")

    result = (await coordinator.dispatch_tasks("root", [task_id]))[0]

    assert "cannot be blank" in (result.error or "")
    assert executor.validator_calls == 0


class EvidencePolicy:
    def __init__(self, max_queries=1, timeout=0.1):
        self.max_queries = max_queries
        self.timeout = timeout

    def plan(self, context):
        return ValidationPlan(
            max_evidence_queries=self.max_queries,
            max_evidence_items_per_query=1,
            evidence_timeout_seconds=self.timeout,
        )


class EvidenceProvider:
    def __init__(self, *, delay=0, error=None, item_count=1, items=None):
        self.delay = delay
        self.error = error
        self.item_count = item_count
        self.items = items
        self.requests = []

    async def query(self, request):
        self.requests.append(request)
        if self.delay:
            await asyncio.sleep(self.delay)
        if self.error:
            raise self.error
        return EvidenceProviderResult(
            items=(self.items if self.items is not None else [{"index": index} for index in range(self.item_count)]),
            evidence_refs=[ArtifactRef(uri="memory://evidence", version="1")],
        )


class BlockingValidatorExecutor:
    def __init__(self):
        self.coordinator = None
        self.context_ready = asyncio.Event()
        self.context = None

    async def run_worker(self, context):
        await self.coordinator.submit_attempt(
            {
                **context,
                "role": AgentRole.WORKER.value,
                "trace_id": context["worker_trace_id"],
                "tool_call_id": f"submit-{context['attempt_id']}",
            },
            AttemptSubmission("done", [ArtifactRef(uri="memory://work", version="1")]),
        )
        return WorkerRunResult(context["worker_trace_id"], "completed")

    async def run_validator(self, context):
        self.context = context
        self.context_ready.set()
        await asyncio.Event().wait()

    async def stop(self, trace_id):
        return True


async def running_evidence_validation(tmp_path, provider, *, max_queries=1, timeout=0.1):
    executor = BlockingValidatorExecutor()
    coordinator, store, _ = await make_coordinator(tmp_path, executor)
    coordinator.config = OrchestrationConfig(stop_grace_seconds=0)
    coordinator.validation_policy = EvidencePolicy(max_queries, timeout)
    coordinator.evidence_provider = provider
    task_id = await create_task(coordinator, "evidence")
    operation = await coordinator.start_operation("root", "dispatch", task_ids=[task_id])
    await asyncio.wait_for(executor.context_ready.wait(), timeout=1)
    context = executor.context
    actor = {
        **context,
        "role": AgentRole.VALIDATOR.value,
        "trace_id": context["validator_trace_id"],
        "tool_call_id": "evidence-1",
    }
    request = EvidenceQuery(
        root_trace_id="root",
        task_id=task_id,
        attempt_id=context["attempt_id"],
        snapshot_id=context["snapshot_id"],
        validation_id=context["validation_id"],
        query="facts",
        limit=100,
    )
    return coordinator, store, operation, actor, request


@pytest.mark.asyncio
async def test_wire_orchestration_applies_custom_evidence_provider(tmp_path):
    runner = SimpleNamespace(trace_store=None, task_coordinator=None)
    provider = EvidenceProvider()
    policy = EvidencePolicy(max_queries=1, timeout=0.1)
    coordinator = wire_orchestration(
        runner,
        FileSystemTaskStore(str(tmp_path)),
        FileSystemArtifactStore(str(tmp_path)),
        OrchestrationConfig(stop_grace_seconds=0),
        validation_policy=policy,
        evidence_provider=provider,
    )
    executor = BlockingValidatorExecutor()
    executor.coordinator = coordinator
    coordinator.set_executor(executor)
    await coordinator.ensure_ledger("root", ROOT_TASK_SPEC)
    task_id = await create_task(coordinator, "wired evidence provider")
    operation = await coordinator.start_operation(
        "root",
        "dispatch",
        task_ids=[task_id],
    )
    try:
        await asyncio.wait_for(executor.context_ready.wait(), timeout=1)
        context = executor.context
        actor = {
            **context,
            "role": AgentRole.VALIDATOR.value,
            "trace_id": context["validator_trace_id"],
            "tool_call_id": "wired-evidence-1",
        }
        request = EvidenceQuery(
            root_trace_id="root",
            task_id=task_id,
            attempt_id=context["attempt_id"],
            snapshot_id=context["snapshot_id"],
            validation_id=context["validation_id"],
            query="wired facts",
            limit=100,
        )
        response = await coordinator.query_evidence(actor, request)
        report = (await coordinator.task_store.load("root")).validations[request.validation_id]

        assert coordinator.validation_policy is policy
        assert coordinator.evidence_provider is provider
        assert len(provider.requests) == 1
        assert provider.requests[0].limit == 1
        assert response.items == [{"index": 0}]
        assert report.evidence_queries_used == 1
    finally:
        await coordinator.stop_operation("root", operation.operation_id)


@pytest.mark.asyncio
async def test_coordinator_evidence_ownership_and_persistent_budget(tmp_path):
    provider = EvidenceProvider()
    coordinator, store, operation, actor, request = await running_evidence_validation(tmp_path, provider)
    response = await coordinator.query_evidence(actor, request)
    assert response.queries_used == 1
    assert response.queries_remaining == 0
    assert provider.requests[0].limit == 1
    report = (await store.load("root")).validations[request.validation_id]
    assert report.evidence_queries_used == 1
    assert next(iter(report.evidence_query_results.values()))["status"] == "completed"

    replay = await coordinator.query_evidence(actor, request)
    assert replay.queries_used == 1
    assert len(provider.requests) == 1
    with pytest.raises(EvidenceBudgetExceeded):
        await coordinator.query_evidence(
            {**actor, "tool_call_id": "evidence-2"},
            EvidenceQuery(**{**request.__dict__, "query": "more"}),
        )
    with pytest.raises(EvidenceOwnershipError):
        await coordinator.query_evidence(
            {**actor, "tool_call_id": "foreign"},
            EvidenceQuery(**{**request.__dict__, "snapshot_id": "foreign"}),
        )
    with pytest.raises(EvidenceOwnershipError):
        await coordinator.query_evidence(
            {**actor, "trace_id": "foreign", "tool_call_id": "foreign-trace"},
            request,
        )
    await coordinator.stop_operation("root", operation.operation_id)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("provider", "message"),
    [
        (EvidenceProvider(delay=0.05), "timed out"),
        (EvidenceProvider(error=RuntimeError("offline")), "offline"),
        (EvidenceProvider(item_count=2), "more than"),
        (EvidenceProvider(items=[{"bad": object()}]), "non-JSON"),
    ],
)
async def test_evidence_provider_failures_are_validation_errors(tmp_path, provider, message):
    timeout = 0.001 if provider.delay else 0.1
    coordinator, store, operation, actor, request = await running_evidence_validation(
        tmp_path,
        provider,
        timeout=timeout,
    )
    with pytest.raises(EvidenceProviderError, match=message):
        await coordinator.query_evidence(actor, request)
    assert (await store.load("root")).validations[request.validation_id].evidence_queries_used == 1
    with pytest.raises(EvidenceProviderError, match=message):
        await coordinator.query_evidence(actor, request)
    assert len(provider.requests) == 1
    await coordinator.stop_operation("root", operation.operation_id)


@pytest.mark.asyncio
async def test_concurrent_evidence_replay_calls_provider_once(tmp_path):
    provider = EvidenceProvider(delay=0.02)
    coordinator, _, operation, actor, request = await running_evidence_validation(
        tmp_path,
        provider,
    )
    responses = await asyncio.gather(
        coordinator.query_evidence(actor, request),
        coordinator.query_evidence(actor, request),
    )
    assert [item.queries_used for item in responses] == [1, 1]
    assert len(provider.requests) == 1
    await coordinator.stop_operation("root", operation.operation_id)


def test_evidence_query_rejects_empty_scope_and_limit():
    values = {
        "root_trace_id": "root",
        "task_id": "task",
        "attempt_id": "attempt",
        "snapshot_id": "snapshot",
        "validation_id": "validation",
        "query": "facts",
    }
    with pytest.raises(ValueError, match="query"):
        EvidenceQuery(**{**values, "query": " "})
    with pytest.raises(ValueError, match="limit"):
        EvidenceQuery(**values, limit=0)
    with pytest.raises(ValueError, match="limit"):
        EvidenceQuery(**values, limit=True)
    with pytest.raises(ValueError, match="root_trace_id"):
        EvidenceQuery(**{**values, "root_trace_id": 123})