aios/tests/test_task_principal_auto_recovery.py

from __future__ import annotations

from pathlib import Path
import sys

import pytest

ROOT = Path(__file__).resolve().parents[1]
APP = ROOT / "active-runtime" / "app"
if str(APP) not in sys.path:
    sys.path.insert(0, str(APP))

from aios_task_admission_bootstrap import (  # noqa: E402
    OwnerDelegationEpoch,
    TaskAdmissionError,
    TaskAdmissionGrant,
    TaskAdmissionStore,
)
from aios_task_admission_live_adapters import (  # noqa: E402
    TaskAdmissionPrincipalProvider,
)


def _grant(*, capabilities=None) -> TaskAdmissionGrant:
    return TaskAdmissionGrant(
        grant_id="grant-1",
        envelope_id="env-1",
        task_id="task-1",
        repository="aios/aios",
        authority_base_oid="a" * 40,
        target_branch="agent/task-1",
        generation=1,
        execution_digest="b" * 64,
        capabilities=tuple(
            capabilities
            or (
                "task.principal.issue",
                "writer.lease",
                "governed.write.file",
                "governed.write.files",
            )
        ),
        file_allowlist=("active-runtime/app/example.py",),
        risk="medium",
        trust_domain="aios.owner",
        policy_digest="c" * 64,
        authorization_kind="delegation_epoch",
        owner_identity="owner",
        parent_epoch_id="epoch-1",
        issued_at=100,
        expires_at=1000,
        one_time_use=True,
    )


def _epoch(*, capabilities=None) -> OwnerDelegationEpoch:
    return OwnerDelegationEpoch(
        epoch_id="epoch-1",
        owner_identity="owner",
        trust_domain="aios.owner",
        repository="aios/aios",
        allowed_branch_prefixes=("agent/",),
        allowed_capabilities=frozenset(
            capabilities
            or {
                "task.principal.issue",
                "writer.lease",
                "governed.write.file",
                "governed.write.files",
            }
        ),
        allowed_file_prefixes=("active-runtime/app/",),
        risk_ceiling="medium",
        issued_at=1,
        not_before=1,
        expires_at=2000,
        revocation_epoch=0,
        policy_digest="c" * 64,
    )


class Store:
    def __init__(self, grant):
        self.grant = grant

    def get_grant_by_task_id(self, task_id):
        return self.grant if self.grant and self.grant.task_id == task_id else None


class Authorization:
    def __init__(self, epoch=None, error=None):
        self.epoch = epoch or _epoch()
        self.error = error
        self.calls = 0

    def verify_delegation_epoch(self, epoch_id):
        self.calls += 1
        if self.error is not None:
            raise self.error
        assert epoch_id == self.epoch.epoch_id
        return self.epoch


class Port:
    def __init__(self, *, state="CLAIMED", job_id="job-1"):
        self.state = state
        self.job_id = job_id
        self.calls = 0

    def verify_active_task(self, grant):
        self.calls += 1
        if self.state not in {"CLAIMED", "RUNNING"}:
            raise TaskAdmissionError(
                "TASK_PRINCIPAL_TASK_NOT_CLAIMED", 409, details={"state": self.state}
            )
        return (
            {"task_id": grant.task_id, "state": self.state},
            {
                "task_id": grant.task_id,
                "repository": grant.repository,
                "target_branch": grant.target_branch,
                "generation": grant.generation,
                "external_job_id": self.job_id,
            },
        )


def _request(**updates):
    value = {
        "task_id": "task-1",
        "job_id": "job-1",
        "repository": "aios/aios",
        "base_oid": "a" * 40,
        "target_branch": "agent/task-1",
        "generation": 1,
    }
    value.update(updates)
    return value


def test_provider_recovers_principal_from_claimed_task_without_caller_bearer():
    auth = Authorization()
    provider = TaskAdmissionPrincipalProvider(
        store=Store(_grant()),
        authorization=auth,
        task_port=Port(state="CLAIMED"),
        epoch_id="epoch-1",
        clock=lambda: 200,
    )
    result = provider({"operation": "acquireWriterLease", "headers": {}}, _request())
    assert result["schema"] == "aios.task-principal-server-attestation.v1"
    assert result["task_id"] == "task-1"
    assert result["job_id"] == "job-1"
    assert result["capability"] == "writer.lease"
    assert result["task_state"] == "CLAIMED"
    assert auth.calls == 1
    assert provider.readiness()["ready"] is True


def test_provider_requires_claim_before_writer_lease():
    provider = TaskAdmissionPrincipalProvider(
        store=Store(_grant()),
        authorization=Authorization(),
        task_port=Port(state="QUEUED"),
        epoch_id="epoch-1",
        clock=lambda: 200,
    )
    with pytest.raises(TaskAdmissionError, match="TASK_PRINCIPAL_TASK_NOT_CLAIMED"):
        provider({"operation": "acquireWriterLease"}, _request())


@pytest.mark.parametrize(
    ("field", "value"),
    [
        ("job_id", "job-other"),
        ("repository", "other/repo"),
        ("target_branch", "agent/other"),
        ("generation", 2),
        ("base_oid", "d" * 40),
    ],
)
def test_provider_rejects_request_binding_drift(field, value):
    provider = TaskAdmissionPrincipalProvider(
        store=Store(_grant()),
        authorization=Authorization(),
        task_port=Port(),
        epoch_id="epoch-1",
        clock=lambda: 200,
    )
    with pytest.raises(TaskAdmissionError, match="TASK_PRINCIPAL_REQUEST_BINDING_MISMATCH"):
        provider(
            {"operation": "acquireWriterLease"},
            _request(**{field: value}),
        )


def test_provider_reverifies_revocation_on_every_use():
    provider = TaskAdmissionPrincipalProvider(
        store=Store(_grant()),
        authorization=Authorization(
            error=TaskAdmissionError("TASK_ADMISSION_EPOCH_REVOKED", 403)
        ),
        task_port=Port(),
        epoch_id="epoch-1",
        clock=lambda: 200,
    )
    with pytest.raises(TaskAdmissionError, match="TASK_ADMISSION_EPOCH_REVOKED"):
        provider({"operation": "acquireWriterLease"}, _request())


def test_provider_rejects_missing_execution_capability():
    provider = TaskAdmissionPrincipalProvider(
        store=Store(_grant(capabilities=("task.principal.issue",))),
        authorization=Authorization(),
        task_port=Port(),
        epoch_id="epoch-1",
        clock=lambda: 200,
    )
    with pytest.raises(TaskAdmissionError, match="TASK_PRINCIPAL_CAPABILITY_DENIED"):
        provider({"operation": "acquireWriterLease"}, _request())


def test_provider_does_not_treat_missing_admission_as_owner_reauthorization():
    provider = TaskAdmissionPrincipalProvider(
        store=Store(None),
        authorization=Authorization(),
        task_port=Port(),
        epoch_id="epoch-1",
        clock=lambda: 200,
    )
    with pytest.raises(TaskAdmissionError, match="TASK_PRINCIPAL_AUTHORIZATION_REQUIRED"):
        provider({"operation": "acquireWriterLease"}, _request())


def test_store_recovers_exact_grant_by_task_id(tmp_path):
    store = TaskAdmissionStore(tmp_path / "admission.sqlite3")
    grant = _grant()
    store.save_grant_and_receipt(
        envelope_id=grant.envelope_id,
        grant=grant,
        receipt={"schema": "fixture", "task_id": grant.task_id},
        created_at=200,
    )
    recovered = store.get_grant_by_task_id(grant.task_id)
    assert recovered == grant
    assert store.get_grant_by_task_id("task-missing") is None