aios/tests/baseline_v1_1/test_job_safety.py

from __future__ import annotations

import base64
from datetime import datetime, timedelta, timezone
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
import textwrap

import pytest
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey

import aios_gitops_job_executor as executor_module
from aios_gitops_job_controller import (
    GitOpsJobController,
    JobError,
    JobState,
    JobStore,
    ServiceUnavailable,
)


OLD = "bba675fab3b7e0c5fba7066720252b3c91b9f20d"
NEW = "3fbbc18fb2131bbbd9e511004a59ee95cf3e86b7"


class Clock:
    def __init__(self) -> None:
        self.value = datetime(2026, 9, 8, tzinfo=timezone.utc)

    def __call__(self) -> datetime:
        return self.value


class Authority:
    def __init__(self) -> None:
        self.branch = OLD
        self.main = OLD
        self.mutations = 0

    def read_ref(self, ref: str):
        return self.main if ref == "refs/heads/main" else self.branch

    def ancestry_check(self, older: str, newer: str) -> bool:
        return older == newer or (older, newer) == (OLD, NEW)

    def advance_task_branch_CAS(self, _branch: str, expected: str, new: str) -> str:
        assert self.branch == expected
        self.mutations += 1
        self.branch = new
        return new

    def ff_CAS_main_after_verified_protected_merge(self, expected: str, new: str) -> str:
        assert self.main == expected
        self.mutations += 1
        self.main = new
        return new

    def commit_file_in_isolated_worktree(self, **_kwargs) -> str:
        self.mutations += 1
        return NEW


class Broker:
    def __init__(self) -> None:
        self.remote = OLD
        self.mutations = 0

    def verify_identity(self):
        return {"login": "aios-autonomy-agent", "repository": executor_module.REPOSITORY}

    def readback(self):
        return self.remote

    def push_branch(self, new: str) -> str:
        self.mutations += 1
        self.remote = new
        return new

    def create_pr(self, *_args):
        self.mutations += 1
        return {"number": 1}

    def status(self, *_args):
        self.mutations += 1
        return {"status": "success"}

    def merge(self, *_args):
        self.mutations += 1
        return {"merged": True}


def campaign_request(operation: str = "SYNC_TASK_BRANCH_TO_GITEA", **updates):
    request = {
        "campaign_id": executor_module.CAMPAIGN_ID,
        "repository": executor_module.REPOSITORY,
        "operation": operation,
        "branch": executor_module.BRANCH,
        "expected_branch_oid": OLD if operation == "ADVANCE_EXISTING_CAMPAIGN_BRANCH" else NEW,
        "payload_digest": "0" * 64,
        "task_id": executor_module.TASK_ID,
        "new_branch_oid": NEW,
        "authorization_recovery_digest": executor_module.AUTHORIZATION_DIGEST,
    }
    if operation == "RECONCILE_AUTHORITY_MAIN":
        request["expected_main_oid"] = OLD
    request.update(updates)
    return request


def recovery():
    return {
        "digest": executor_module.AUTHORIZATION_DIGEST,
        "campaign_id": executor_module.CAMPAIGN_ID,
        "task_id": executor_module.TASK_ID,
        "repository": executor_module.REPOSITORY,
        "branch": executor_module.BRANCH,
        "expected_branch_oid": OLD,
    }


def generic_request(**updates):
    request = {
        "schema": "aios.operation-job.v1",
        "tenant_id": "tenant-a",
        "subject_id": "subject-a",
        "mission_id": "mission-a",
        "grant_id": "grant-a",
        "task_id": "task-a",
        "effect_id": "effect-a",
        "operation": "host.status.v1",
        "operation_version": "1",
        "resource_kind": "host",
        "resource_id": "tenant-a/staging/local-host",
        "input_digest": hashlib.sha256(b"{}").hexdigest(),
    }
    request.update(updates)
    return request


def receipt_signer():
    private = Ed25519PrivateKey.generate()

    def sign(payload):
        raw = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
        return {
            "schema": "aios.operation-receipt.v1",
            "key_id": "test-machine",
            "payload": payload,
            "signature": base64.b64encode(private.sign(raw)).decode(),
        }

    def verify(envelope):
        try:
            raw = json.dumps(envelope["payload"], sort_keys=True, separators=(",", ":")).encode()
            private.public_key().verify(base64.b64decode(envelope["signature"], validate=True), raw)
            return envelope.get("key_id") == "test-machine"
        except (InvalidSignature, KeyError, TypeError, ValueError):
            return False

    return sign, verify


def signed_generic_receipt(sign, request, job_id, lease):
    return sign({
        **{key: request[key] for key in (
            "tenant_id", "subject_id", "mission_id", "grant_id", "task_id",
            "effect_id", "operation", "operation_version", "resource_kind",
            "resource_id", "input_digest",
        )},
        "job_id": job_id,
        "claim_id": lease["lease_id"],
        "generation": int(lease["generation"]),
        "observation": {"active": True},
        "verified": True,
    })


def make_executor(tmp_path: Path, monkeypatch, operation="SYNC_TASK_BRANCH_TO_GITEA"):
    monkeypatch.setattr(executor_module, "load_recovered_authorization", lambda _path: recovery())
    clock = Clock()
    controller = GitOpsJobController(JobStore(tmp_path / "jobs"), clock=clock)
    job = controller.submit(campaign_request(operation))
    authority, broker = Authority(), Broker()
    if operation != "ADVANCE_EXISTING_CAMPAIGN_BRANCH":
        authority.branch = NEW
    executor = executor_module.DurableJobExecutor(
        controller, authority, broker, tmp_path / "recovery.json", clock=clock
    )
    return clock, controller, job, authority, broker, executor


def test_lease_requires_authorized_live_job_and_bounded_future_deadline(tmp_path):
    clock = Clock()
    controller = GitOpsJobController(JobStore(tmp_path / "jobs"), clock=clock)
    job = controller.submit(campaign_request("ADVANCE_EXISTING_CAMPAIGN_BRANCH"))
    kwargs = dict(
        lease_id="lease-1",
        lineage_anchor_oid=OLD,
        current_head_oid=OLD,
        generation=1,
        expected_revision=job.revision,
    )
    with pytest.raises(JobError, match="LEASE_JOB_STATE_INVALID"):
        controller.lease_job(job.job_id, lease_expires_at=(clock.value + timedelta(seconds=30)).isoformat(), **kwargs)
    authorized = controller.advance(job.job_id, JobState.AUTHORIZED, side_effect_state="AUTHORIZED")
    kwargs["expected_revision"] = authorized.revision
    with pytest.raises(JobError, match="LEASE_EXPIRED"):
        controller.lease_job(job.job_id, lease_expires_at=(clock.value - timedelta(seconds=1)).isoformat(), **kwargs)
    with pytest.raises(JobError, match="LEASE_DURATION_EXCEEDED"):
        controller.lease_job(job.job_id, lease_expires_at=(clock.value + timedelta(hours=1)).isoformat(), **kwargs)

    terminal_controller = GitOpsJobController(JobStore(tmp_path / "terminal"), clock=clock)
    terminal = terminal_controller.submit(campaign_request("ADVANCE_EXISTING_CAMPAIGN_BRANCH"))
    terminal = terminal_controller.advance(
        terminal.job_id,
        JobState.MANUAL_OWNER_REVIEW_REQUIRED,
        side_effect_state="MANUAL_OWNER_REVIEW_REQUIRED",
    )
    with pytest.raises(JobError, match="LEASE_JOB_STATE_INVALID"):
        terminal_controller.lease_job(
            terminal.job_id,
            lease_expires_at=(clock.value + timedelta(seconds=30)).isoformat(),
            **{**kwargs, "expected_revision": terminal.revision},
        )


def test_resource_job_profile_and_lease_are_typed_and_fenced(tmp_path):
    clock = Clock()
    controller = GitOpsJobController(JobStore(tmp_path / "jobs"), clock=clock)
    request = {
        "schema": "aios.operation-job.v1",
        "tenant_id": "tenant-a",
        "subject_id": "subject-a",
        "mission_id": "mission-a",
        "grant_id": "grant-a",
        "task_id": "task-a",
        "effect_id": "effect-a",
        "operation": "host.status.v1",
        "operation_version": "1",
        "resource_kind": "host",
        "resource_id": "local-host",
        "input_digest": hashlib.sha256(b"{}").hexdigest(),
    }
    submitted = controller.submit(request)
    authorized = controller.advance(submitted.job_id, JobState.AUTHORIZED, side_effect_state="AUTHORIZED")
    leased = controller.lease_resource_job(
        submitted.job_id,
        lease_id="resource-lease-1",
        resource_kind="host",
        resource_id="local-host",
        generation=None,
        lease_expires_at=(clock.value + timedelta(seconds=30)).isoformat(),
        expected_revision=authorized.revision,
    )
    assert leased.state is JobState.LEASED
    first_lease = controller.get_execution_lease(submitted.job_id)
    assert first_lease["generation"] == 1
    controller.assert_execution_lease(submitted.job_id, lease_id="resource-lease-1", generation=1)
    with pytest.raises(JobError, match="LEASE_FENCE_REJECTED"):
        controller.assert_execution_lease(submitted.job_id, lease_id="resource-lease-1", generation=0)

    clock.value += timedelta(seconds=31)
    second_request = {**request, "task_id": "task-b", "effect_id": "effect-b"}
    second = controller.submit(second_request)
    second = controller.advance(second.job_id, JobState.AUTHORIZED, side_effect_state="AUTHORIZED")
    controller.lease_resource_job(
        second.job_id,
        lease_id="resource-lease-2",
        resource_kind="host",
        resource_id="local-host",
        generation=None,
        lease_expires_at=(clock.value + timedelta(seconds=30)).isoformat(),
        expected_revision=second.revision,
    )
    assert controller.get_execution_lease(second.job_id)["generation"] == 2
    with pytest.raises(JobError, match="LEASE_FENCE_REJECTED"):
        controller.assert_execution_lease(submitted.job_id, lease_id="resource-lease-1", generation=1)


def test_generic_job_cannot_skip_intent_readback_or_trusted_receipt(tmp_path):
    clock = Clock()
    sign, verify = receipt_signer()
    controller = GitOpsJobController(
        JobStore(tmp_path / "jobs"), clock=clock, generic_receipt_verifier=verify
    )
    request = generic_request()
    submitted = controller.submit(request)
    invented = sign({"verified": True})
    with pytest.raises(JobError, match="GENERIC_JOB_TRANSITION_INVALID"):
        controller.advance(
            submitted.job_id, JobState.COMPLETED,
            side_effect_state="VERIFIED", receipt=invented,
        )
    assert controller.store.read(submitted.job_id)["state"] == JobState.REQUESTED.value

    authorized = controller.advance(
        submitted.job_id, JobState.AUTHORIZED, side_effect_state="AUTHORIZED"
    )
    leased = controller.lease_resource_job(
        submitted.job_id, lease_id="claim-a", resource_kind=request["resource_kind"],
        resource_id=request["resource_id"], generation=None,
        lease_expires_at=(clock.value + timedelta(seconds=30)).isoformat(),
        expected_revision=authorized.revision,
    )
    controller.advance(
        submitted.job_id, JobState.EXECUTING,
        side_effect_state="INTENT_COMMITTED", expected_revision=leased.revision,
    )
    lease = controller.get_execution_lease(submitted.job_id)
    valid_receipt = signed_generic_receipt(sign, request, submitted.job_id, lease)
    tampered = {**valid_receipt, "payload": {**valid_receipt["payload"], "input_digest": "f" * 64}}
    with pytest.raises(JobError, match="GENERIC_RECEIPT_BINDING_MISMATCH"):
        controller.advance(
            submitted.job_id, JobState.LOCAL_VERIFIED,
            side_effect_state="READBACK_VERIFIED", receipt=tampered,
        )
    local = controller.advance(
        submitted.job_id, JobState.LOCAL_VERIFIED,
        side_effect_state="READBACK_VERIFIED", receipt=valid_receipt,
    )
    completed = controller.advance(
        submitted.job_id, JobState.COMPLETED,
        side_effect_state="READBACK_VERIFIED", receipt=valid_receipt,
        expected_revision=local.revision,
    )
    assert completed.state is JobState.COMPLETED


def test_generic_complete_requires_configured_verifier_and_recovery_accepts_bound_receipt(tmp_path):
    clock = Clock()
    sign, verify = receipt_signer()
    request = generic_request()
    controller = GitOpsJobController(
        JobStore(tmp_path / "jobs"), clock=clock, generic_receipt_verifier=verify,
    )
    submitted = controller.submit(request)
    authorized = controller.advance(submitted.job_id, JobState.AUTHORIZED, side_effect_state="AUTHORIZED")
    leased = controller.lease_resource_job(
        submitted.job_id, lease_id="claim-a", resource_kind=request["resource_kind"],
        resource_id=request["resource_id"], generation=None,
        lease_expires_at=(clock.value + timedelta(seconds=30)).isoformat(),
        expected_revision=authorized.revision,
    )
    executing = controller.advance(
        submitted.job_id, JobState.EXECUTING,
        side_effect_state="INTENT_COMMITTED", expected_revision=leased.revision,
    )
    receipt = signed_generic_receipt(
        sign, request, submitted.job_id, controller.get_execution_lease(submitted.job_id)
    )
    pending = controller.advance(
        submitted.job_id, JobState.RECONCILIATION_REQUIRED,
        side_effect_state="UNKNOWN", expected_revision=executing.revision,
    )
    clock.value += timedelta(seconds=31)
    completed = controller.reconcile(
        submitted.job_id,
        lambda _row: {
            "state": JobState.COMPLETED.value,
            "side_effect_state": "READBACK_VERIFIED",
            "receipt": receipt,
        },
    )
    assert completed.state is JobState.COMPLETED

    untrusted = GitOpsJobController(JobStore(tmp_path / "untrusted"), clock=Clock())
    other_request = generic_request(effect_id="effect-b", task_id="task-b")
    other = untrusted.submit(other_request)
    other = untrusted.advance(other.job_id, JobState.AUTHORIZED, side_effect_state="AUTHORIZED")
    other = untrusted.lease_resource_job(
        other.job_id, lease_id="claim-b", resource_kind=other_request["resource_kind"],
        resource_id=other_request["resource_id"], generation=None,
        lease_expires_at=(Clock().value + timedelta(seconds=30)).isoformat(),
        expected_revision=other.revision,
    )
    untrusted.advance(other.job_id, JobState.EXECUTING, side_effect_state="INTENT_COMMITTED")
    other_receipt = signed_generic_receipt(
        sign, other_request, other.job_id, untrusted.get_execution_lease(other.job_id)
    )
    with pytest.raises(JobError, match="GENERIC_RECEIPT_VERIFIER_REQUIRED"):
        untrusted.advance(
            other.job_id, JobState.LOCAL_VERIFIED,
            side_effect_state="READBACK_VERIFIED", receipt=other_receipt,
        )


def test_reconcile_crash_can_be_taken_over_and_stale_observer_cannot_commit(tmp_path):
    clock = Clock()
    controller = GitOpsJobController(
        JobStore(tmp_path / "jobs"), clock=clock, reconcile_lease_seconds=10
    )
    job = controller.submit(campaign_request("ADVANCE_EXISTING_CAMPAIGN_BRANCH"))
    controller.advance(job.job_id, JobState.RECOVERABLE_FAILURE, side_effect_state="UNKNOWN")
    with pytest.raises(KeyboardInterrupt):
        controller.reconcile(job.job_id, lambda _row: (_ for _ in ()).throw(KeyboardInterrupt()))
    assert controller.store.read(job.job_id)["state"] == JobState.RECONCILING.value
    with pytest.raises(JobError, match="RECONCILIATION_CLAIM_ACTIVE"):
        controller.reconcile(job.job_id, lambda _row: {"state": "REQUESTED"})

    clock.value += timedelta(seconds=11)

    def stale_observer(_row):
        clock.value += timedelta(seconds=11)
        winner = controller.reconcile(
            job.job_id,
            lambda _new: {"state": "REQUESTED", "side_effect_state": "RECONCILED_ABSENT"},
        )
        assert winner.state is JobState.REQUESTED
        return {"state": "COMPLETED", "side_effect_state": "VERIFIED"}

    with pytest.raises(JobError, match="RECONCILIATION_CLAIM_LOST"):
        controller.reconcile(job.job_id, stale_observer)
    assert controller.store.read(job.job_id)["state"] == JobState.REQUESTED.value


@pytest.mark.parametrize(
    "operation",
    [
        "ADVANCE_EXISTING_CAMPAIGN_BRANCH",
        "SYNC_TASK_BRANCH_TO_GITEA",
        "CREATE_PR",
        "PUBLISH_CHECK",
        "PROTECTED_MERGE",
        "RECONCILE_AUTHORITY_MAIN",
        "COMMIT_FILE",
    ],
)
def test_dry_run_never_calls_a_mutating_adapter(tmp_path, monkeypatch, operation):
    _clock, controller, job, authority, broker, executor = make_executor(
        tmp_path, monkeypatch, operation
    )
    if operation == "COMMIT_FILE":
        payload = b'{"path":"x.txt","content_base64":"eA=="}'
        controller = GitOpsJobController(JobStore(tmp_path / "commit-jobs"), clock=Clock())
        job = controller.submit(
            campaign_request(operation, payload_digest=hashlib.sha256(payload).hexdigest()),
            payload=payload,
        )
        executor = executor_module.DurableJobExecutor(
            controller, authority, broker, tmp_path / "recovery.json", clock=Clock()
        )
    result = executor.execute(job.job_id, dry_run=True)
    assert result["mode"] == "DRY_RUN"
    assert authority.mutations == broker.mutations == 0
    assert controller.store.read(job.job_id)["state"] == JobState.REQUESTED.value


def test_execution_does_not_call_effect_when_intent_commit_is_refused(tmp_path, monkeypatch):
    _clock, controller, job, authority, broker, executor = make_executor(tmp_path, monkeypatch)
    original = controller.advance

    def refuse(job_id, state, **kwargs):
        if state is JobState.EXECUTING:
            raise ServiceUnavailable("JOB_DB_UNAVAILABLE")
        return original(job_id, state, **kwargs)

    monkeypatch.setattr(controller, "advance", refuse)
    with pytest.raises(ServiceUnavailable, match="JOB_DB_UNAVAILABLE"):
        executor.execute(job.job_id)
    assert authority.mutations == broker.mutations == 0
    with pytest.raises(executor_module.JobExecutionError, match="PERSISTENCE_FAILURE_LATCHED"):
        executor.execute(job.job_id)


def test_external_success_receipt_failure_latches_then_readback_recovers_without_repeat(tmp_path, monkeypatch):
    _clock, controller, job, _authority, broker, executor = make_executor(tmp_path, monkeypatch)
    original = controller.advance
    fail = {"once": True}

    def refuse_receipt(job_id, state, **kwargs):
        if state is JobState.COMPLETED and fail["once"]:
            fail["once"] = False
            raise ServiceUnavailable("JOB_DB_UNAVAILABLE")
        return original(job_id, state, **kwargs)

    monkeypatch.setattr(controller, "advance", refuse_receipt)
    with pytest.raises(ServiceUnavailable, match="JOB_DB_UNAVAILABLE"):
        executor.execute(job.job_id)
    assert broker.mutations == 1
    with pytest.raises(executor_module.JobExecutionError, match="PERSISTENCE_FAILURE_LATCHED"):
        executor.execute(job.job_id)
    next_job = controller.submit(campaign_request("CREATE_PR", idempotency_key="after-receipt-failure"))
    with pytest.raises(executor_module.JobExecutionError, match="PERSISTENCE_FAILURE_LATCHED"):
        executor.execute(next_job.job_id)
    assert broker.mutations == 1

    restarted = executor_module.DurableJobExecutor(
        controller, executor.authority, broker, tmp_path / "recovery.json", clock=executor.clock
    )
    with pytest.raises(executor_module.JobExecutionError, match="PERSISTENCE_FAILURE_LATCHED"):
        restarted.execute(job.job_id)

    recovered = executor.reconcile(job.job_id)
    assert recovered["state"] == JobState.COMPLETED.value
    assert broker.mutations == 1
    assert executor.execute(job.job_id)["state"] == JobState.COMPLETED.value


def test_unsupported_exact_readback_reports_unknown_without_fake_success_or_effect(tmp_path, monkeypatch):
    _clock, controller, job, _authority, broker, executor = make_executor(
        tmp_path, monkeypatch, "PUBLISH_CHECK"
    )
    result = executor.execute(job.job_id)
    assert result["state"] == JobState.RECONCILIATION_REQUIRED.value
    assert result["receipt"]["verification"] == "UNKNOWN"
    assert result["receipt"]["blocker"] == "EXACT_READBACK_UNSUPPORTED"
    assert broker.mutations == 0
    again = executor.execute(job.job_id)
    assert again["state"] == JobState.RECONCILIATION_REQUIRED.value
    assert broker.mutations == 0


def test_subprocess_crash_after_durable_sink_success_recovers_without_repeat(tmp_path):
    """AT-22: process death after effect durability cannot repeat the effect."""
    job_root = tmp_path / "jobs"
    sink = tmp_path / "effect-sink.json"
    recovery_path = tmp_path / "recovery.json"
    script = textwrap.dedent(
        f"""
        import json, os
        from pathlib import Path
        import aios_gitops_job_executor as m
        from aios_gitops_job_controller import GitOpsJobController, JobStore

        OLD = {OLD!r}; NEW = {NEW!r}
        sink = Path({str(sink)!r})
        root = Path({str(job_root)!r})
        m.load_recovered_authorization = lambda _path: {{
            'digest': m.AUTHORIZATION_DIGEST, 'campaign_id': m.CAMPAIGN_ID,
            'task_id': m.TASK_ID, 'repository': m.REPOSITORY,
            'branch': m.BRANCH, 'expected_branch_oid': OLD,
        }}

        class Authority:
            def read_ref(self, _ref): return NEW
            def ancestry_check(self, older, newer): return older == newer or (older == OLD and newer == NEW)

        class Broker:
            def verify_identity(self): return {{'login':'aios-autonomy-agent','repository':m.REPOSITORY}}
            def readback(self):
                if sink.exists():
                    if os.environ.get('CRASH_AFTER_EFFECT') == '1': os._exit(77)
                    return json.loads(sink.read_text())['oid']
                return OLD
            def push_branch(self, oid):
                previous = json.loads(sink.read_text()) if sink.exists() else {{'count': 0}}
                payload = {{'count': previous['count'] + 1, 'oid': oid}}
                with sink.open('w') as stream:
                    json.dump(payload, stream); stream.flush(); os.fsync(stream.fileno())
                return oid

        controller = GitOpsJobController(JobStore(root))
        request = {campaign_request()!r}
        job = controller.submit(request)
        executor = m.DurableJobExecutor(controller, Authority(), Broker(), Path({str(recovery_path)!r}))
        if os.environ['ACTION'] == 'execute':
            executor.execute(job.job_id)
        else:
            print(json.dumps(executor.reconcile(job.job_id), sort_keys=True))
        """
    )
    environment = os.environ.copy()
    environment["PYTHONDONTWRITEBYTECODE"] = "1"
    environment["ACTION"] = "execute"
    environment["CRASH_AFTER_EFFECT"] = "1"
    _repo_root = Path(__file__).resolve().parents[2]
    environment["PYTHONPATH"] = str(_repo_root / "active-runtime" / "app") + os.pathsep + environment.get("PYTHONPATH", "")
    crashed = subprocess.run([sys.executable, "-c", script], env=environment, check=False)
    assert crashed.returncode == 77
    assert json.loads(sink.read_text()) == {"count": 1, "oid": NEW}

    environment["ACTION"] = "recover"
    environment["CRASH_AFTER_EFFECT"] = "0"
    recovered = subprocess.run(
        [sys.executable, "-c", script], env=environment, check=True,
        capture_output=True, text=True,
    )
    assert json.loads(recovered.stdout)["state"] == JobState.COMPLETED.value
    assert json.loads(sink.read_text()) == {"count": 1, "oid": NEW}