Agentswarm/scripts/test-runtime-contract.py

"""Lightweight contract checks for the HeiCode Agent Manager runtime bridge."""
import asyncio
import os
import sys
import types
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
if "redis.asyncio" not in sys.modules:
    redis_module = types.ModuleType("redis")
    redis_asyncio_module = types.ModuleType("redis.asyncio")
    redis_asyncio_module.Redis = object
    redis_module.asyncio = redis_asyncio_module
    sys.modules["redis"] = redis_module
    sys.modules["redis.asyncio"] = redis_asyncio_module
if "httpx" not in sys.modules:
    httpx_module = types.ModuleType("httpx")

    class AsyncClient:
        def __init__(self, *args, **kwargs):
            pass

        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb):
            return False

        async def post(self, *args, **kwargs):
            return types.SimpleNamespace(status_code=204, text="")

    httpx_module.AsyncClient = AsyncClient
    sys.modules["httpx"] = httpx_module

from orchestrator.redis_client import redis_client
from orchestrator.swarm_runtime import RuntimeValidationError, SwarmRuntime


class FakeRedis:
    def __init__(self):
        self.values = {}
        self.lists = {}

    async def set(self, key, value, ex=None):
        self.values[key] = value

    async def get(self, key):
        return self.values.get(key)

    async def incr(self, key):
        self.values[key] = int(self.values.get(key, 0)) + 1
        return self.values[key]

    async def keys(self, pattern):
        prefix = pattern.rstrip("*")
        return [key for key in self.values if key.startswith(prefix)]

    async def rpush(self, key, *values):
        self.lists.setdefault(key, []).extend(values)

    async def lrange(self, key, start, end):
        items = self.lists.get(key, [])
        if end == -1:
            return items[start:]
        return items[start : end + 1]


def valid_body():
    return {
        "orchestration_plan": {
            "objective": "Add a multiply helper and push the branch",
            "sub_mode": "code",
            "risk_level": "low",
            "budget": {"duration_seconds": 3600, "token_limit": 20000},
            "agents": [
                {
                    "task_id": "backend-1",
                    "role": "backend",
                    "title": "Backend change",
                    "description": "Implement the helper",
                    "depends_on": [],
                    "resource_grants": [
                        {"secret_ref": "azkv://heicode/git-write-token"}
                    ],
                }
            ],
        },
        "callback": {
            "url": "http://manager.local/api/agnet/callbacks/swarm-events",
            "subscribed_events": ["task.created"],
        },
        "metadata": {
            "manager_deployment_id": "dep_manager_123",
            "correlation_id": "corr_123",
        },
        "billing_context": {"secret_ref": "azkv://heicode/billing"},
        "resource_grants": [{"ref": "azkv://heicode/repo-main"}],
    }


async def main():
    fake = FakeRedis()
    redis_client.client = fake
    runtime = SwarmRuntime()
    os.environ["ENABLE_SUBTASK_HANDOFF"] = "true"

    body = valid_body()
    runtime.validate_create_request(body)

    invalid = valid_body()
    invalid["resource_grants"][0]["ref"] = "plain-token"
    try:
        runtime.validate_create_request(invalid)
        raise AssertionError("plain resource ref was accepted")
    except RuntimeValidationError:
        pass

    invalid = valid_body()
    invalid["metadata"]["api_key"] = "secret-value"
    try:
        runtime.validate_create_request(invalid)
        raise AssertionError("plaintext secret metadata was accepted")
    except RuntimeValidationError:
        pass

    run, created = await runtime.get_or_create_run(body, "idem-1", "corr_123")
    assert created is True
    same_run, created_again = await runtime.get_or_create_run(body, "idem-1", "corr_123")
    assert created_again is False
    assert same_run.swarm_id == run.swarm_id

    tasks = runtime.build_task_descriptions(body)
    assert tasks[0]["task_id"] == "backend-1"
    assert tasks[0]["depends_on"] == []
    assert tasks[0]["workflow_mode"] == "multi_agent"

    await runtime.emit_event(run, "timeline.updated", payload={"summary": "ok"})
    logs = await runtime.list_events(run.swarm_id)
    assert logs["events"][0]["event_type"] == "deployment.status_changed"
    assert logs["events"][-1]["event_type"] == "timeline.updated"

    approval_body = valid_body()
    approval_body["orchestration_plan"]["risk_level"] = "high"
    approval_run, _ = await runtime.get_or_create_run(approval_body, "idem-2", "corr_123")
    approval_id = next(iter(approval_run.approvals))
    rejected = await runtime.record_approval_decision(
        approval_run.swarm_id,
        approval_id,
        {"decision": "rejected", "reason": "user rejected"},
    )
    assert rejected.status == "blocked"

    print("runtime contract checks passed")


if __name__ == "__main__":
    asyncio.run(main())