"""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())