heicode/scripts/agent_sub_mode_smoke.py

#!/usr/bin/env python3
"""
Smoke test for Heicode Manager ordinary sub-mode / Agent integration.

Default mode checks public health/status endpoints only. To validate an
authenticated deployment flow, provide:

  HEICODE_ACCESS_TOKEN=<manager access token>
  HEICODE_USER_ID=<user id>
  HEICODE_DEPLOYMENT_ID=<deployment id>

Optional:
  HEICODE_RUN_SIMULATE=true
  HEICODE_RUN_CALLBACK_SMOKE=true
  HEICODE_CALLBACK_TOKEN=<Manager callback service token>
  HEICODE_CALLBACK_SWARM_ID=<runtime swarm id, optional>
  HEICODE_MANAGER_BASE_URL=https://code.xinghanlab.com
  AGENT_MANAGER_HEALTH_URL=http://20.212.121.126/api/agent/health
"""

from __future__ import annotations

import json
import os
import sys
import urllib.error
import urllib.request


BASE = os.getenv("HEICODE_MANAGER_BASE_URL", "https://code.xinghanlab.com").rstrip("/")
AGENT_HEALTH = os.getenv(
    "AGENT_MANAGER_HEALTH_URL", "http://20.212.121.126/api/agent/health"
)
ACCESS_TOKEN = os.getenv("HEICODE_ACCESS_TOKEN", "").strip()
USER_ID = os.getenv("HEICODE_USER_ID", "").strip()
DEPLOYMENT_ID = os.getenv("HEICODE_DEPLOYMENT_ID", "").strip()
CALLBACK_TOKEN = os.getenv("HEICODE_CALLBACK_TOKEN", "").strip()
CALLBACK_SWARM_ID = os.getenv("HEICODE_CALLBACK_SWARM_ID", "").strip()
RUN_SIMULATE = os.getenv("HEICODE_RUN_SIMULATE", "").strip().lower() in {
    "1",
    "true",
    "yes",
}
RUN_CALLBACK_SMOKE = os.getenv("HEICODE_RUN_CALLBACK_SMOKE", "").strip().lower() in {
    "1",
    "true",
    "yes",
}
UA = "heicode-agent-sub-mode-smoke/1.0"
FAILS: list[str] = []


def request(
    method: str,
    url: str,
    headers: dict[str, str] | None = None,
    body: bytes | None = None,
) -> tuple[int, dict[str, str], str]:
    req = urllib.request.Request(url, method=method, data=body)
    req.add_header("User-Agent", UA)
    for key, value in (headers or {}).items():
        req.add_header(key, value)
    try:
        resp = urllib.request.urlopen(req, timeout=20)
        return resp.status, dict(resp.headers), resp.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as exc:
        return exc.code, dict(exc.headers), exc.read().decode("utf-8", "replace")


def manager_request(
    method: str, path: str, *, auth: bool = False, payload: dict | None = None
) -> tuple[int, dict[str, str], str]:
    headers = {"Content-Type": "application/json"}
    if auth:
        headers["Authorization"] = f"Bearer {ACCESS_TOKEN}"
        headers["X-User-ID"] = USER_ID
    body = json.dumps(payload).encode("utf-8") if payload is not None else None
    return request(method, BASE + path, headers=headers, body=body)


def check(label: str, ok: bool, detail: str = "") -> None:
    print(f"[{'OK' if ok else 'FAIL'}] {label}" + (f" -- {detail}" if detail else ""))
    if not ok:
        FAILS.append(label)


def parse_json(label: str, body: str) -> dict:
    try:
        parsed = json.loads(body)
        check(label, isinstance(parsed, dict))
        return parsed if isinstance(parsed, dict) else {}
    except json.JSONDecodeError as exc:
        check(label, False, str(exc))
        return {}


print(f"Manager base: {BASE}")
print(f"Agent health: {AGENT_HEALTH}\n")

code, _headers, body = manager_request("GET", "/api/status")
status = parse_json("Manager /api/status JSON parseable", body)
check("Manager /api/status success", code == 200 and status.get("success") is True, f"HTTP {code}")
version = status.get("data", {}).get("version") if isinstance(status.get("data"), dict) else ""
check("Manager status carries version", bool(version), f"version={version!r}")

code, _headers, body = manager_request("GET", "/api/agent/callbacks/swarm-events/schema")
schema = parse_json("callback schema JSON parseable", body)
schema_data = schema.get("data", {}) if isinstance(schema.get("data"), dict) else {}
schema_events = schema_data.get("events") or []
schema_event_types = {
    item.get("event_type")
    for item in schema_events
    if isinstance(item, dict)
}
check(
    "callback schema lists swarm and ordinary sub events",
    code == 200
    and schema.get("success") is True
    and {"task.claimed", "handoff.requested", "phase.changed", "artifact.created"}.issubset(schema_event_types),
    f"HTTP {code}",
)

code, _headers, body = request("GET", AGENT_HEALTH)
agent = parse_json("Agent Manager health JSON parseable", body)
agent_data = agent.get("data", {}) if isinstance(agent.get("data"), dict) else {}
agent_status = agent_data.get("status") or agent.get("status")
check(
    "Agent Manager health reachable",
    code == 200 and str(agent_status).lower() in {"healthy", "ok", "up"},
    f"HTTP {code}, status={agent_status!r}",
)

if not (ACCESS_TOKEN and USER_ID and DEPLOYMENT_ID):
    print("\nAuthenticated deployment checks skipped: set HEICODE_ACCESS_TOKEN, HEICODE_USER_ID and HEICODE_DEPLOYMENT_ID.")
else:
    if RUN_CALLBACK_SMOKE:
        if not CALLBACK_TOKEN:
            check("callback smoke has HEICODE_CALLBACK_TOKEN", False)
        else:
            callback_headers = {
                "Content-Type": "application/json",
                "X-Agent-Service-Token": CALLBACK_TOKEN,
                "X-Correlation-ID": f"smoke-corr-{DEPLOYMENT_ID}",
            }
            swarm_id = CALLBACK_SWARM_ID or f"smoke-swarm-{DEPLOYMENT_ID}"
            callback_events = [
                {
                    "event_id": f"smoke-task-created-{DEPLOYMENT_ID}",
                    "event_type": "task.created",
                    "deployment_id": DEPLOYMENT_ID,
                    "swarm_id": swarm_id,
                    "task_id": f"smoke-task-{DEPLOYMENT_ID}",
                    "source": "smoke-runtime",
                    "payload": {
                        "title": "Smoke task created",
                        "agent_role": "backend",
                    },
                },
                {
                    "event_id": f"smoke-task-claimed-{DEPLOYMENT_ID}",
                    "event_type": "task.claimed",
                    "deployment_id": DEPLOYMENT_ID,
                    "swarm_id": swarm_id,
                    "task_id": f"smoke-task-{DEPLOYMENT_ID}",
                    "source": "smoke-runtime",
                    "payload": {
                        "agent_role": "backend",
                    },
                },
                {
                    "event_id": f"smoke-handoff-{DEPLOYMENT_ID}",
                    "event_type": "handoff.requested",
                    "deployment_id": DEPLOYMENT_ID,
                    "swarm_id": swarm_id,
                    "task_id": f"smoke-task-{DEPLOYMENT_ID}",
                    "source": "smoke-runtime",
                    "payload": {
                        "from_role": "backend",
                        "to_role": "reviewer",
                        "reason": "Smoke handoff validation",
                    },
                },
                {
                    "event_id": f"smoke-artifact-{DEPLOYMENT_ID}",
                    "event_type": "artifact.created",
                    "deployment_id": DEPLOYMENT_ID,
                    "swarm_id": swarm_id,
                    "task_id": f"smoke-task-{DEPLOYMENT_ID}",
                    "source": "smoke-runtime",
                    "artifact": {
                        "artifact_id": f"smoke-artifact-{DEPLOYMENT_ID}",
                        "artifact_type": "test_report",
                        "title": "Smoke artifact",
                        "summary": "Callback smoke wrote this artifact through the Manager callback endpoint.",
                        "uri": f"artifact://smoke/{DEPLOYMENT_ID}/report",
                    },
                    "payload": {
                        "stage": "testing",
                        "checkpoint": "artifact_ready",
                    },
                },
            ]
            for event in callback_events:
                body_bytes = json.dumps(event).encode("utf-8")
                callback_headers["X-Agent-Event-Id"] = event["event_id"]
                code, _headers, body = request(
                    "POST",
                    BASE + "/api/agent/callbacks/swarm-events",
                    headers=callback_headers,
                    body=body_bytes,
                )
                callback_response = parse_json(
                    f"callback {event['event_type']} JSON parseable", body
                )
                check(
                    f"callback {event['event_type']} accepted",
                    code == 200 and callback_response.get("success") is True,
                    f"HTTP {code}",
                )

    if RUN_SIMULATE:
        code, _headers, body = manager_request(
            "POST",
            f"/api/agent/user/deployments/{DEPLOYMENT_ID}/simulate-events",
            auth=True,
            payload={},
        )
        simulated = parse_json("simulate-events JSON parseable", body)
        check(
            "simulate-events accepted",
            code == 200 and simulated.get("success") is True,
            f"HTTP {code}",
        )

    code, _headers, body = manager_request(
        "GET", f"/api/agent/user/deployments/{DEPLOYMENT_ID}", auth=True
    )
    detail = parse_json("deployment detail JSON parseable", body)
    detail_data = detail.get("data", {}) if isinstance(detail.get("data"), dict) else {}
    check(
        "deployment detail success",
        code == 200 and detail.get("success") is True and detail_data.get("deployment_id") == DEPLOYMENT_ID,
        f"HTTP {code}",
    )

    code, _headers, body = manager_request(
        "GET", f"/api/agent/user/deployments/{DEPLOYMENT_ID}/artifacts", auth=True
    )
    artifacts = parse_json("artifacts JSON parseable", body)
    artifacts_data = artifacts.get("data", {}) if isinstance(artifacts.get("data"), dict) else {}
    artifact_items = artifacts_data.get("artifacts") or artifacts_data.get("items") or []
    check(
        "artifacts endpoint success",
        code == 200 and artifacts.get("success") is True and isinstance(artifact_items, list),
        f"HTTP {code}, count={len(artifact_items) if isinstance(artifact_items, list) else 'n/a'}",
    )

    code, _headers, body = manager_request(
        "GET", f"/api/agent/user/deployments/{DEPLOYMENT_ID}/timeline", auth=True
    )
    timeline = parse_json("timeline JSON parseable", body)
    timeline_data = timeline.get("data", {}) if isinstance(timeline.get("data"), dict) else {}
    callbacks = timeline_data.get("callbacks") or []
    merged = timeline_data.get("timeline") or []
    event_types = {
        item.get("event_type") or item.get("event")
        for item in merged
        if isinstance(item, dict)
    }
    sources = {
        item.get("source")
        for item in merged
        if isinstance(item, dict) and item.get("source")
    }
    check(
        "timeline endpoint success",
        code == 200 and timeline.get("success") is True and isinstance(merged, list),
        f"HTTP {code}, callbacks={len(callbacks) if isinstance(callbacks, list) else 'n/a'}",
    )
    check(
        "timeline includes ordinary sub-mode task flow when simulated/runtime callbacks exist",
        any(str(event).startswith("task.") or str(event).startswith("handoff.") for event in event_types),
        "events=" + ",".join(sorted(str(event) for event in event_types if event)[:10]),
    )
    check(
        "timeline includes artifact or approval flow",
        "artifact.created" in event_types or "approval.requested" in event_types,
        "events=" + ",".join(sorted(str(event) for event in event_types if event)[:10]),
    )
    if RUN_CALLBACK_SMOKE:
        check(
            "timeline includes smoke callback source",
            "smoke-runtime" in sources,
            "sources=" + ",".join(sorted(str(source) for source in sources)[:10]),
        )
        check(
            "callback smoke artifact is queryable",
            any(
                isinstance(item, dict)
                and item.get("artifact_id") == f"smoke-artifact-{DEPLOYMENT_ID}"
                for item in artifact_items
            ),
            f"artifact_count={len(artifact_items) if isinstance(artifact_items, list) else 'n/a'}",
        )

if FAILS:
    print(f"\n{len(FAILS)} check(s) failed:")
    for item in FAILS:
        print(f"  - {item}")
    sys.exit(1)

print("\nALL AGENT SUB-MODE SMOKE CHECKS PASSED")