Agentswarm/scripts/test-agent-launcher.py

"""Swarm-side agent launcher tests (agent_swarm#16: Swarm launches agents + sets the limit).

Covers the pure/limit/env/resolve/backend-selection logic of orchestrator/agent_launcher:
  * launch_count + plan_launch_specs cap the pool at MAX_AGENTS_PER_USER (never push the user over);
  * each launch spec composes the agent env the runtime reads (ORCHESTRATOR_URL / AGENT_ID /
    AGENT_CAPABILITIES / OPENAI_API_BASE / OPENAI_API_KEY / HEICODE_USER_ID);
  * the model key is resolved server-side (override → azkv secret_ref dev-map → OPENAI_API_KEY),
    never fabricated;
  * `command` backend builds argv from the template; `none` backend is a no-op (no spawn).

Hermetic: no Redis / model / subprocess (backend forced to none/command-build only).

Run from agent_swarm_v6:
    python scripts/test-agent-launcher.py
"""
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from orchestrator import agent_launcher as al

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


class FakeRun:
    def __init__(self, swarm_id="swarm-abc"):
        self.swarm_id = swarm_id


def test_launch_count():
    check("count = min(pool, limit-connected)", al.launch_count(pool_size=3, limit=10, connected_user_agents=0) == 3)
    check("count respects per-user cap", al.launch_count(pool_size=5, limit=10, connected_user_agents=8) == 2)
    check("count never negative (already over cap)", al.launch_count(pool_size=3, limit=10, connected_user_agents=10) == 0)
    check("count clamps to pool when cap is high", al.launch_count(pool_size=3, limit=100, connected_user_agents=0) == 3)


def test_plan_specs():
    run = FakeRun()
    body = {"orchestration_plan": {"objective": "x"}, "billing_context": {"default_model_id": "gpt-x"}}
    specs = al.plan_launch_specs(run, body, connected_user_agents=8, limit=10, pool_size=3,
                                 model_key="sk-test", orchestrator_url="ws://orch:8000", user_id="u-1")
    check("plan caps at limit (8 connected, cap 10 -> launch 2)", len(specs) == 2)
    s = specs[0]
    check("spec env has ORCHESTRATOR_URL", s.env.get("ORCHESTRATOR_URL") == "ws://orch:8000")
    check("spec env has model key (server-side injected)", s.env.get("OPENAI_API_KEY") == "sk-test")
    check("spec env has model id from billing_context", s.env.get("OPENAI_MODEL") == "gpt-x")
    check("spec env has HEICODE_USER_ID for per-user cap", s.env.get("HEICODE_USER_ID") == "u-1")
    check("spec has AGENT_ID + capabilities", bool(s.agent_id) and bool(s.env.get("AGENT_CAPABILITIES")))
    check("agent ids unique", len({sp.agent_id for sp in specs}) == len(specs))
    # No key -> OPENAI_API_KEY omitted (not fabricated), no user -> HEICODE_USER_ID omitted.
    specs2 = al.plan_launch_specs(run, body, connected_user_agents=0, limit=10, pool_size=1,
                                  model_key=None, orchestrator_url="ws://orch", user_id=None)
    check("no model key -> OPENAI_API_KEY omitted", "OPENAI_API_KEY" not in specs2[0].env)
    check("no user -> HEICODE_USER_ID omitted", "HEICODE_USER_ID" not in specs2[0].env)


def test_resolve_model_key():
    for k in ("AGENT_LAUNCH_MODEL_KEY", "OPENAI_API_KEY", "HEICODE_SECRET_res_model_1"):
        os.environ.pop(k, None)
    # override wins
    os.environ["AGENT_LAUNCH_MODEL_KEY"] = "sk-override"
    check("override key wins", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/res_model_1"}}) == "sk-override")
    os.environ.pop("AGENT_LAUNCH_MODEL_KEY")
    # azkv secret_ref -> dev env map
    os.environ["HEICODE_SECRET_res_model_1"] = "sk-from-kv"
    check("azkv secret_ref resolved via dev map", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/res_model_1"}}) == "sk-from-kv")
    os.environ.pop("HEICODE_SECRET_res_model_1")
    # fallback to orchestrator OPENAI_API_KEY
    os.environ["OPENAI_API_KEY"] = "sk-orch"
    check("fallback to orchestrator OPENAI_API_KEY", al.resolve_model_key({}) == "sk-orch")
    os.environ.pop("OPENAI_API_KEY")
    check("unresolved -> None (never fabricated)", al.resolve_model_key({"billing_context": {"secret_ref": "azkv://kv/secrets/missing"}}) is None)


def test_command_backend_build():
    os.environ["AGENT_LAUNCH_CMD"] = "launch-agent.sh --id {agent_id} --caps {capabilities}"
    spec = al.AgentLaunchSpec(agent_id="swarm-abc-agent-1", capabilities="python,general",
                              env={"OPENAI_API_KEY": "sk-x"})
    argv = al.build_launch_command(spec)
    check("command template substitutes agent_id/capabilities",
          argv == ["launch-agent.sh", "--id", "swarm-abc-agent-1", "--caps", "python,general"])
    check("secret not on argv (passed via env)", all("sk-x" not in a for a in argv))
    os.environ.pop("AGENT_LAUNCH_CMD")


async def test_backend_none_noop():
    import asyncio  # noqa
    os.environ["AGENT_LAUNCH_BACKEND"] = "none"
    specs = [al.AgentLaunchSpec(agent_id="a1", capabilities="general", env={})]
    launched = await al.launch(specs, swarm_id="swarm-abc")
    check("backend=none launches nothing (external)", launched == [])
    os.environ.pop("AGENT_LAUNCH_BACKEND")


def test_k8s_manifests():
    import json
    spec = al.AgentLaunchSpec(
        agent_id="swarm-abc-agent-1", capabilities="python,general",
        env={"ORCHESTRATOR_URL": "ws://orch.svc:8000", "AGENT_ID": "swarm-abc-agent-1",
             "AGENT_CAPABILITIES": "python,general", "OPENAI_API_KEY": "sk-secret-xyz",
             "OPENAI_API_BASE": "https://hm/v1", "HEICODE_USER_ID": "u-1"},
    )
    # Secret manifest carries the key (applied via stdin, not argv).
    sec = al.build_secret_manifest(key_value="sk-secret-xyz", namespace="heicode-swarm",
                                   name="swarm-agent-key-swarm-abc", labels={"app": "heicode-swarm-agent"})
    check("secret kind/type", sec["kind"] == "Secret" and sec["type"] == "Opaque")
    check("secret stringData has the key", sec["stringData"]["OPENAI_API_KEY"] == "sk-secret-xyz")

    # Pod manifest references the key via secretKeyRef — NEVER inline.
    pod = al.build_pod_manifest(spec, namespace="heicode-swarm", swarm_id="swarm-abc",
                                image="img:1", secret_name="swarm-agent-key-swarm-abc",
                                service_account="swarm-agent-sa")
    blob = json.dumps(pod)
    check("pod kind/name", pod["kind"] == "Pod" and pod["metadata"]["name"] == "swarm-abc-agent-1")
    check("pod labels include swarm-id + user-id",
          pod["metadata"]["labels"].get("heicode-swarm-id") == "swarm-abc"
          and pod["metadata"]["labels"].get("heicode-user-id") == "u-1")
    env = {e["name"]: e for e in pod["spec"]["containers"][0]["env"]}
    check("OPENAI_API_KEY via secretKeyRef (not inline)",
          "value" not in env["OPENAI_API_KEY"] and env["OPENAI_API_KEY"]["valueFrom"]["secretKeyRef"]["name"] == "swarm-agent-key-swarm-abc")
    check("raw key value NOT in pod manifest (only in Secret)", "sk-secret-xyz" not in blob)
    check("non-secret env inline", env["ORCHESTRATOR_URL"]["value"] == "ws://orch.svc:8000")
    check("pod has resource limits + restartPolicy + serviceAccount",
          pod["spec"]["containers"][0]["resources"]["limits"]
          and pod["spec"]["restartPolicy"] == "OnFailure"
          and pod["spec"]["serviceAccountName"] == "swarm-agent-sa")
    # No secret_name -> key omitted entirely (never inlined).
    pod2 = al.build_pod_manifest(spec, namespace="heicode-swarm", swarm_id="swarm-abc", secret_name=None)
    env2 = {e["name"] for e in pod2["spec"]["containers"][0]["env"]}
    check("no secret -> OPENAI_API_KEY omitted from pod env", "OPENAI_API_KEY" not in env2)


def main():
    import asyncio
    test_launch_count()
    test_plan_specs()
    test_resolve_model_key()
    test_command_backend_build()
    test_k8s_manifests()
    asyncio.run(test_backend_none_noop())
    print()
    if failures:
        print(f"{len(failures)} agent-launcher check(s) FAILED: {failures}")
        sys.exit(1)
    print("all agent-launcher checks passed")


if __name__ == "__main__":
    main()