aios/tests/test_v2r2c5_capability_contract.py

import pytest

from aios_capability_contract import (
    has_capabilities,
    normalize_required_capabilities,
    normalize_worker_capabilities,
)
from aios_agent_task_protocol import TaskProtocolError, TaskStore


def test_n01_unknown_required_token_is_explicit_and_denied():
    assert normalize_required_capabilities(["unknown-token"]) == frozenset({"exact:unknown-token"})
    assert has_capabilities(["testing"], ["unknown-token"]) is False


def test_n02_unknown_worker_token_is_never_granted():
    assert "exact:unknown-token" not in normalize_worker_capabilities(["unknown-token"])
    assert has_capabilities(["unknown-token"], ["testing"]) is False


def test_n03_admitted_capability_and_claim_match_use_same_contract():
    assert has_capabilities(["git.read"], ["repository_read"]) is True
    assert has_capabilities(["git.read"], ["testing"]) is False


def test_n01_task_store_rejects_unknown_before_queue_insert(tmp_path):
    store = TaskStore(tmp_path / "tasks.sqlite3")
    with pytest.raises(TaskProtocolError) as exc:
        store.create_task(
            task_id="task-unknown",
            goal="unknown",
            instructions="unknown",
            success_criteria=[],
            required_capabilities=["unknown-token"],
            idempotency_key="idempotency-unknown",
        )
    assert exc.value.code == "TASK_CAPABILITY_UNKNOWN"
    with store._connect() as connection:
        assert connection.execute("SELECT COUNT(*) FROM tasks").fetchone()[0] == 0


def test_n02_existing_unknown_requirement_is_denied_by_poll_and_claim(tmp_path):
    store = TaskStore(tmp_path / "tasks.sqlite3")
    store.register_agent(
        agent_id="agent-old-row",
        agent_type="test",
        agent_version="1",
        capabilities=["testing"],
    )
    task = store.create_task(
        task_id="task-old-row",
        goal="old",
        instructions="old",
        success_criteria=[],
        required_capabilities=["testing"],
        idempotency_key="idempotency-old-row",
    )
    with store._connect() as connection:
        connection.execute(
            "UPDATE tasks SET required_capabilities_json=? WHERE task_id=?",
            ('["unknown-token"]', task["task_id"]),
        )
    assert store.poll_tasks(agent_id="agent-old-row") == []
    with pytest.raises(TaskProtocolError) as exc:
        store.claim_task(task_id=task["task_id"], agent_id="agent-old-row")
    assert exc.value.code == "AGENT_CAPABILITY_MISMATCH"


def test_n03_atomic_claim_denies_admitted_requirement_worker_drift(tmp_path):
    store = TaskStore(tmp_path / "tasks.sqlite3")
    store.register_agent(
        agent_id="agent-drift",
        agent_type="test",
        agent_version="1",
        capabilities=["git.read"],
    )
    task = store.create_task(
        task_id="task-drift",
        goal="drift",
        instructions="drift",
        success_criteria=[],
        required_capabilities=["testing"],
        idempotency_key="idempotency-drift",
    )
    with pytest.raises(TaskProtocolError) as exc:
        store.claim_task(task_id=task["task_id"], agent_id="agent-drift")
    assert exc.value.code == "AGENT_CAPABILITY_MISMATCH"
    assert store.get_task(task["task_id"])["status"] == "QUEUED"


# ---------------------------------------------------------------------------
# C5 production.* capability family (V2R2C5 / formal-lineage-canonical-v17)
# ---------------------------------------------------------------------------
# The full 17-cap set of the provisioned C5 promotion worker
# (agent/c5-promotion-exec-20260929).
C5_WORKER17 = [
    "agent-task:control",
    "code.read",
    "code.write",
    "git.read",
    "git.test",
    "governed.task.create",
    "governed.write.branch",
    "governed.write.file",
    "governed.write.files",
    "production.plan.apply",
    "production.promote",
    "production.rollback",
    "task.principal.issue",
    "test.execute",
    "testing",
    "verification",
    "writer.lease",
]


def test_c5_01_production_required_normalization_recognized():
    # production.* must NOT be tagged as unknown (exact:prefix) after patch.
    required = normalize_required_capabilities(["production.promote"])
    assert required == frozenset({"production.promote"})
    unknown = [c for c in normalize_required_capabilities(C5_WORKER17) if c.startswith("exact:")]
    assert unknown == []


def test_c5_02_production_worker_normalization_recognized():
    effective = normalize_worker_capabilities(C5_WORKER17)
    assert {"production.plan.apply", "production.promote", "production.rollback"}.issubset(effective)
    assert len(effective) >= 17


def test_c5_03_has_capabilities_full17_is_true():
    # Previously raw claim returned 403 AGENT_CAPABILITY_MISMATCH; patched contract must be True.
    assert has_capabilities(C5_WORKER17, C5_WORKER17) is True


def test_c5_04_unknown_capability_still_fail_closed():
    assert has_capabilities(C5_WORKER17, ["production.promote", "unknown-cap"]) is False
    assert normalize_required_capabilities(["unknown-cap"]) == frozenset({"exact:unknown-cap"})


def test_c5_05_wildcard_still_rejected():
    assert has_capabilities(["production.*"], ["production.promote"]) is False
    assert normalize_required_capabilities(["production.*"]) == frozenset({"exact:production.*"})


def test_c5_06_missing_production_still_denied():
    # A worker WITHOUT production.* must not be able to claim a production task.
    worker_no_prod = [c for c in C5_WORKER17 if not c.startswith("production.")]
    assert has_capabilities(worker_no_prod, ["production.promote"]) is False