aios/scripts/smoke_task_identity_overlay.py

#!/usr/bin/env python3
"""End-to-end smoke test for the loopback task identity overlay."""
from __future__ import annotations

import json
import os
import sys
import tempfile
import threading
import urllib.request
from pathlib import Path


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

from aios_task_identity_overlay_api import OverlayServer  # noqa: E402
from owner_authorization_broker import ACTION  # noqa: E402
from task_identity_overlay import (  # noqa: E402
    IdentityOverlayConfig,
    TaskBinding,
    TaskIdentityOverlayService,
    TaskIdentityOverlayStore,
)


TASK_ID = "V2R1C6_TASK_SCOPED_WRITER_PROVISIONING"
BRANCH = "agent/v2r1c6-mainline-reconciliation-20260805"
BASE_OID = "875d41c3101093149414292067fd700dc214d4d3"


class FakeProvisioner:
    def provision(self, raw, owner):
        return {
            "status": "READY_FOR_CODE_AUTHORING",
            "approval_plan_id": "sha256:" + "a" * 64,
            "grant_id": "grant-smoke",
            "grant_digest_sha256": "b" * 64,
            "branch": BRANCH,
            "branch_oid": BASE_OID,
            "base_oid": BASE_OID,
            "grant_consumed": True,
            "writer_scope": {
                "workspace_role": "task-scoped-isolated-writer",
                "writes_enabled": True,
                "allowed_branch": BRANCH,
                "allowed_base_oid": BASE_OID,
                "allowed_routes": ["POST /api/commit-file"],
                "allow_main_write": False,
                "allow_tag_write": False,
                "allow_baseline_write": False,
                "allow_force_push": False,
            },
        }


def request_body():
    return {
        "action": ACTION,
        "task_id": TASK_ID,
        "repository": "aios/aios",
        "base_ref": "refs/heads/main",
        "base_oid": BASE_OID,
        "target_branch": BRANCH,
        "one_time_use": True,
        "expires_in_seconds": 900,
        "production_change_authorized": False,
        "tag_change_authorized": False,
        "baseline_change_authorized": False,
        "main_write_authorized": False,
        "allowed_operations": [
            "create_worktree",
            "create_target_branch_if_absent",
            "commit_to_target_branch",
            "push_target_branch",
        ],
        "denied_operations": [
            "direct_main_push",
            "force_push",
            "tag_create",
            "tag_update",
            "tag_delete",
            "branch_delete",
            "repository_admin",
            "baseline_write",
            "production_deploy",
            "production_marker_write",
            "history_rewrite",
        ],
    }


def main() -> int:
    os.environ["AIOS_GITOPS_IDENTITY_OVERLAY_ENABLED"] = "true"
    with tempfile.TemporaryDirectory() as temporary:
        root = Path(temporary)
        config = IdentityOverlayConfig(
            state_root=root / "state",
            secret_root=root / "run",
            broker_state_root=root / "broker",
            binding=TaskBinding(
                task_id=TASK_ID,
                repository="aios/aios",
                base_ref="refs/heads/main",
                base_oid=BASE_OID,
                target_branch=BRANCH,
            ),
            owner_identity="aiman",
            max_ttl_seconds=900,
            require_root_issue=False,
        )
        store = TaskIdentityOverlayStore(config)
        issued = store.issue(
            ttl_seconds=900,
            authorization_assertion_ref="owner-session:V2R1C6:smoke",
            operator_identity="smoke-agent",
        )
        secret_ref = Path(issued["secret_ref"])
        token = secret_ref.read_text().strip()

        server = OverlayServer(
            ("127.0.0.1", 0),
            TaskIdentityOverlayService(
                config, provisioner=FakeProvisioner()
            ),
        )
        thread = threading.Thread(target=server.serve_forever, daemon=True)
        thread.start()
        try:
            host, port = server.server_address
            health = json.loads(
                urllib.request.urlopen(
                    f"http://{host}:{port}/healthz", timeout=5
                ).read()
            )
            if health.get("status") != "active":
                raise RuntimeError("health failed")

            body = json.dumps(request_body()).encode("utf-8")
            http_request = urllib.request.Request(
                f"http://{host}:{port}/api/governance/owner-authorizations",
                data=body,
                method="POST",
                headers={
                    "Authorization": f"Bearer {token}",
                    "Content-Type": "application/json",
                },
            )
            result = json.loads(
                urllib.request.urlopen(http_request, timeout=10).read()
            )
            assert result["status"] == "READY_FOR_CODE_AUTHORING"
            assert result["identity_overlay_consumed"] is True
            assert result["identity_overlay_secret_removed"] is True
            assert not secret_ref.exists()
        finally:
            server.shutdown()
            server.server_close()
            thread.join(timeout=5)
    print("task identity overlay smoke: PASS")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())