154 lines
4.9 KiB
Python
154 lines
4.9 KiB
Python
"""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 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())
|