aios/tests/baseline_v1_1/test_loop.py

from __future__ import annotations
import sys
import time
from pathlib import Path
import pytest
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
from task_plane.store import TaskPlaneStore
from aios_gitops_job_controller import JobStore
from aios_baseline.authorization import AuthorizationService, sign_owner_delegation_grant
from aios_baseline.host import SystemdHost
from aios_baseline.lineage import LineageLedger
from aios_baseline.operations import OperationRegistry
from aios_baseline.service import BaselineService, Identity, ServiceError
from aios_baseline.worker_policy import WorkerProfile


def fixture_service(tmp_path):
    owner = Ed25519PrivateKey.generate()
    machine = Ed25519PrivateKey.generate()
    auth = AuthorizationService(tmp_path/'auth.db', {'owner':owner.public_key()}, machine, 'machine',
        owner_key_scopes={'owner':{'tenant_ids':['tenant-a'],'subject_ids':['user-a']}})
    auth.register_grant(sign_owner_delegation_grant(owner,'owner',grant_id='grant-a',generation=1,
        tenant_id='tenant-a',subject_id='user-a',mission_ids=('mission-a',),task_ids=('*',),
        operations=('service-status',),component_ids=('sample',),environments=('staging',),
        issued_at=int(time.time())-1,validity_mode='UNTIL_REVOKED',max_effects=10))
    store=TaskPlaneStore(str(tmp_path/'tasks.db'),high_assurance=True)
    worker=WorkerProfile.from_dict(dict(worker_id='worker-a',identity='machine-a',tenant_scope=['tenant-a'],
        subject_scope=['user-a'],capabilities=['service-status'],priority=10,fallback_executor=False,
        heartbeat_seconds=15,max_concurrency=1))
    program=tmp_path/'systemctl-fixture'
    program.write_text('#!'+sys.executable+'\nprint("Id=sample.service\\nActiveState=active\\nSubState=running\\nMainPID=123\\nInvocationID=abc")\n')
    program.chmod(0o700)
    host=SystemdHost({'tenant-a:staging:sample':'sample.service'},executable=str(program))
    operation=dict(operation_id='service-status',version=1,effect_class='READ_ONLY',handler_ref='host.service_status',
        input_schema='aios.operation.service-status.input.v1',required_permissions=['service-status'],timeout_seconds=10,
        readback={'mode':'NONE','required':False},retry={'mode':'SAFE','max_attempts':2},
        reconcile={'mode':'NONE'},rollback={'mode':'NONE'},owner_boundary={'mode':'WITHIN_GRANT'})
    contract=dict(contract_schema='aios.runtime-contract.v1',contract_revision=1,component_id='sample',
        environment='staging',tenant_scope='tenant-a',lifecycle='active',
        entrypoints={'host':{'component_id':'sample','environment':'staging','resolver_ref':'service://sample/status'}},
        bindings={},resource_refs={},credential_refs={},operations=['service-status'])
    service=BaselineService(store, tmp_path/'jobs', auth, [contract], OperationRegistry([operation]),
        {'worker-a':worker},host,LineageLedger(tmp_path/'lineage.db',store,JobStore(tmp_path/'jobs').read),machine)
    return service


USER=Identity('tenant-a','user-a','user')
WORKER=Identity('tenant-a','machine-a','worker','worker-a')
REQUEST=dict(mission_id='mission-a',grant_id='grant-a',operation='service-status',component_id='sample',
             environment='staging',inputs={},intent='检查 staging 服务并形成诊断结果')


def test_real_local_task_goes_through_claim_job_receipt_and_lineage(tmp_path):
    service=fixture_service(tmp_path)
    created=service.submit_goal(USER,REQUEST,idempotency_key='request-a')
    duplicate=service.submit_goal(USER,REQUEST,idempotency_key='request-a')
    assert duplicate['task_id']==created['task_id']
    service.register_worker(WORKER)
    claim=service.claim_next(WORKER)
    assert claim['task_id']==created['task_id']
    result=service.execute_claim(WORKER,claim['task_id'],claim['claim_id'])
    assert result['status']=='COMPLETED'
    assert result['receipt']['observation']['properties']['MainPID']=='123'
    assert service.get_task(USER,created['task_id'])['status']=='COMPLETED'
    assert service.explain(USER,'mission-a')['mission_id']=='mission-a'
    assert service.claim_next(WORKER) is None


def test_spoofed_tenant_cross_subject_and_worker_results_are_denied(tmp_path):
    service=fixture_service(tmp_path)
    with pytest.raises(ServiceError): service.submit_goal(USER,{**REQUEST,'tenant_id':'tenant-b'},idempotency_key='a')
    created=service.submit_goal(USER,REQUEST,idempotency_key='a')
    for outsider in (Identity('tenant-b','user-a','user'),Identity('tenant-a','user-b','user')):
        with pytest.raises(ServiceError): service.get_task(outsider,created['task_id'])
        with pytest.raises(ServiceError): service.cancel_task(outsider,created['task_id'])
    with pytest.raises(ServiceError): service.register_worker(Identity('tenant-b','machine-a','worker','worker-a'))
    with pytest.raises(ServiceError): service.execute_claim(USER,created['task_id'],'invented')


def test_changed_payload_with_same_request_key_cannot_change_task(tmp_path):
    service=fixture_service(tmp_path)
    service.submit_goal(USER,REQUEST,idempotency_key='a')
    with pytest.raises(ServiceError): service.submit_goal(USER,{**REQUEST,'intent':'different goal'},idempotency_key='a')


def test_cancelled_expired_claim_is_not_resurrected_by_recovery(tmp_path):
    service=fixture_service(tmp_path)
    created=service.submit_goal(USER,REQUEST,idempotency_key='cancel')
    service.register_worker(WORKER)
    claim=service.claim_next(WORKER)
    service.cancel_task(USER,created['task_id'])
    with service.store._transaction():
        service.store.conn.execute('UPDATE claims SET lease_until=? WHERE claim_id=?',(time.time()-1,claim['claim_id']))
    service.recover()
    assert service.get_task(USER,created['task_id'])['status']=='CANCELLED'
    assert service.claim_next(WORKER) is None


def test_multiple_tasks_share_monotonic_resource_fence_without_sharing_state(tmp_path):
    service=fixture_service(tmp_path)
    service.register_worker(WORKER)
    generations=[]
    for request_key in ('first','second'):
        service.submit_goal(USER,REQUEST,idempotency_key=request_key)
        claim=service.claim_next(WORKER)
        result=service.execute_claim(WORKER,claim['task_id'],claim['claim_id'])
        assert result['status']=='COMPLETED'
        generations.append(result['receipt']['generation'])
    assert generations==[1,2]


def test_persisted_receipt_closes_expired_task_without_repeating_host_diagnostic(tmp_path,monkeypatch):
    service=fixture_service(tmp_path)
    service.submit_goal(USER,REQUEST,idempotency_key='recover')
    service.register_worker(WORKER)
    claim=service.claim_next(WORKER)
    original=service._close_task
    def lose_power(*args): raise OSError('simulated process exit before Task acknowledgment')
    monkeypatch.setattr(service,'_close_task',lose_power)
    with pytest.raises(ServiceError): service.execute_claim(WORKER,claim['task_id'],claim['claim_id'])
    with service.store._transaction():
        service.store.conn.execute('UPDATE claims SET lease_until=? WHERE claim_id=?',(time.time()-1,claim['claim_id']))
    monkeypatch.setattr(service,'_close_task',original)
    def must_not_execute(*args,**kwargs): raise AssertionError('must use durable receipt')
    monkeypatch.setattr(service.host,'status',must_not_execute)
    service.recover()
    assert service.get_task(USER,claim['task_id'])['status']=='COMPLETED'


def test_wrong_signed_lineage_cannot_be_attached_to_real_task(tmp_path):
    service=fixture_service(tmp_path)
    service.submit_goal(USER,REQUEST,idempotency_key='scope')
    service.register_worker(WORKER)
    claim=service.claim_next(WORKER)
    result=service.execute_claim(WORKER,claim['task_id'],claim['claim_id'])
    task=service.store.get_task(claim['task_id'])
    assert service._receipt_valid(task,service.signer.sign(result['receipt']))
    for field in ('job_id','mission_id','grant_id','claim_id','generation','operation','input_digest','resource_id'):
        forged=service.signer.sign({**result['receipt'],field:'wrong'})
        assert not service._receipt_valid(task,forged),field


def test_receipt_key_rotation_preserves_old_evidence_only_with_explicit_trust(tmp_path):
    from aios_baseline.evidence import ReceiptSigner
    service=fixture_service(tmp_path)
    service.submit_goal(USER,REQUEST,idempotency_key='rotation')
    service.register_worker(WORKER)
    claim=service.claim_next(WORKER)
    result=service.execute_claim(WORKER,claim['task_id'],claim['claim_id'])
    old=service.signer
    envelope=old.sign(result['receipt'])
    new_key=Ed25519PrivateKey.generate()
    assert not ReceiptSigner(new_key,key_id='next').verify(envelope)
    rotated=ReceiptSigner(new_key,key_id='next',trusted_public_keys={old.key_id:old.public_key})
    service.signer=rotated
    assert service._receipt_valid(service.store.get_task(claim['task_id']),envelope)


def test_fallback_is_normal_worker_eligibility_with_priority_and_same_claim_protocol(tmp_path):
    service=fixture_service(tmp_path)
    task=service.submit_goal(USER,REQUEST,idempotency_key='fallback')
    service.register_worker(WORKER)
    profile=service.worker_profiles['worker-a']
    fallback=WorkerProfile.from_dict({**profile.__dict__,'worker_id':'worker-fallback','identity':'machine-fallback',
        'fallback_executor':True,'priority':100})
    service.worker_profiles[fallback.worker_id]=fallback
    identity=Identity('tenant-a','machine-fallback','worker','worker-fallback')
    service.register_worker(identity)
    assert service.claim_next(identity) is None
    with service.store._transaction():
        service.store.conn.execute('UPDATE baseline_task_runtime SET ready_at=? WHERE task_id=?',(time.time()-601,task['task_id']))
    assert service.claim_next(WORKER) is None
    claim=service.claim_next(identity)
    assert claim['task_id']==task['task_id']
    assert service.execute_claim(identity,claim['task_id'],claim['claim_id'])['status']=='COMPLETED'


def test_deadline_expiring_while_waiting_for_executor_never_starts_host(tmp_path,monkeypatch):
    service=fixture_service(tmp_path)
    service.submit_goal(USER,REQUEST,idempotency_key='deadline')
    service.register_worker(WORKER)
    claim=service.claim_next(WORKER)
    def unexpected(*args,**kwargs): raise AssertionError('expired call executed')
    monkeypatch.setattr(service.host,'status',unexpected)
    service._execute_lock.acquire()
    try:
        with pytest.raises(ServiceError,match='DEADLINE'):
            service.execute_claim(WORKER,claim['task_id'],claim['claim_id'],deadline=time.time()+0.05)
    finally: service._execute_lock.release()
    assert service.get_task(USER,claim['task_id'])['job_id'] is None