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"}}
    git_env = {"GIT_REPO_URL": "https://git.example/r.git", "GIT_USERNAME": "x-access-token",
               "GIT_PASSWORD": "ghp_xyz"}
    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",
                                 git_env=git_env)
    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 env has GIT_REPO_URL (server-side injected)", s.env.get("GIT_REPO_URL") == "https://git.example/r.git")
    check("spec env has GIT_USERNAME + GIT_PASSWORD", s.env.get("GIT_USERNAME") == "x-access-token" and s.env.get("GIT_PASSWORD") == "ghp_xyz")
    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, no git -> GIT_* 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, git_env=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)
    check("no git grant -> GIT_REPO_URL omitted", "GIT_REPO_URL" not in specs2[0].env)


def test_resolve_git_grant():
    # ── credential extraction (pure) ──
    check("git creds: JSON git_username/git_password",
          al._extract_git_credentials('{"git_username":"u","git_password":"p"}') == {"username": "u", "password": "p"})
    check("git creds: token form -> default username",
          al._extract_git_credentials('{"git_token":"ghp_x"}') == {"username": al.DEFAULT_GIT_USERNAME, "password": "ghp_x"})
    check("git creds: bare string -> token",
          al._extract_git_credentials("ghp_bare") == {"username": al.DEFAULT_GIT_USERNAME, "password": "ghp_bare"})
    check("git creds: missing password -> None (never fabricated)", al._extract_git_credentials('{"git_username":"u"}') is None)
    check("git creds: empty -> None", al._extract_git_credentials("") is None)

    # ── grant lookup + resolution (dev env map, hermetic) ──
    os.environ.pop("HEICODE_SECRET_res_git_1", None)
    grant_inline = {"resource_type": "git", "secret_ref": "azkv://heicode-vault/secrets/res_git_1",
                    "metadata": {"repo_url": "https://git.example/acme/app.git", "base_branch": "develop"}}
    body = {"resource_grants": [grant_inline]}
    check("no secret in env -> repo url still resolved (public-repo path)",
          al.resolve_git_grant(body) == {"GIT_REPO_URL": "https://git.example/acme/app.git", "GIT_BASE_BRANCH": "develop"})
    os.environ["HEICODE_SECRET_res_git_1"] = '{"git_username":"bot","git_password":"ghp_kv"}'
    out = al.resolve_git_grant(body)
    check("secret_ref resolved via dev map -> creds injected",
          out.get("GIT_USERNAME") == "bot" and out.get("GIT_PASSWORD") == "ghp_kv"
          and out.get("GIT_REPO_URL") == "https://git.example/acme/app.git")
    os.environ.pop("HEICODE_SECRET_res_git_1")
    check("no git grant -> None", al.resolve_git_grant({"resource_grants": [{"resource_type": "newapi"}]}) is None)
    check("git grant without repo_url -> None",
          al.resolve_git_grant({"resource_grants": [{"resource_type": "git", "secret_ref": "azkv://v/secrets/x"}]}) is None)
    # per-agent grant (orchestration_plan.agents[].resource_grants) also discovered
    body_nested = {"orchestration_plan": {"agents": [{"resource_grants": [grant_inline]}]}}
    check("per-agent git grant discovered",
          (al.resolve_git_grant(body_nested) or {}).get("GIT_REPO_URL") == "https://git.example/acme/app.git")


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_azkv_resolver():
    # ── parse azkv:// refs (pure) ──
    check("azkv parse: bare vault name -> https URL",
          al._parse_azkv_ref("azkv://heicode-vault/secrets/swarm-model-key-u1")
          == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", None))
    check("azkv parse: full host preserved",
          al._parse_azkv_ref("azkv://heicode-vault.vault.azure.net/secrets/res_git_1")
          == ("https://heicode-vault.vault.azure.net", "res_git_1", None))
    check("azkv parse: version captured",
          al._parse_azkv_ref("azkv://heicode-vault/secrets/swarm-model-key-u1/abc")
          == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", "abc"))
    check("azkv parse: short form azkv://<vault>/<name>",
          al._parse_azkv_ref("azkv://heicode-vault/swarm-model-key-u1")
          == ("https://heicode-vault.vault.azure.net", "swarm-model-key-u1", None))
    check("azkv parse: bad ref -> None", al._parse_azkv_ref("azkv://heicode-vault") is None)

    # ── gating: hermetic unless workload identity injected / opt-in ──
    for k in ("AZURE_FEDERATED_TOKEN_FILE", "SECRET_RESOLVER"):
        os.environ.pop(k, None)
    check("azkv disabled without workload identity / opt-in", al._azkv_enabled() is False)
    check("disabled -> _resolve_from_keyvault returns None (no SDK/network touched)",
          al._resolve_from_keyvault("azkv://heicode-vault/secrets/swarm-model-key-u1") is None)
    os.environ["SECRET_RESOLVER"] = "azkv"
    check("azkv enabled via SECRET_RESOLVER=azkv", al._azkv_enabled() is True)
    os.environ.pop("SECRET_RESOLVER")
    os.environ["AZURE_FEDERATED_TOKEN_FILE"] = "/var/run/secrets/azure/tokens/azure-identity-token"
    check("azkv enabled via AZURE_FEDERATED_TOKEN_FILE", al._azkv_enabled() is True)
    os.environ.pop("AZURE_FEDERATED_TOKEN_FILE")


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",
             "GIT_REPO_URL": "https://git.example/acme/app.git", "GIT_USERNAME": "x-access-token",
             "GIT_PASSWORD": "ghp_secret_pw"},
    )
    # Secret manifest carries BOTH sensitive values (model key + git password), applied via stdin.
    sec = al.build_secret_manifest(namespace="heicode-swarm", name="swarm-agent-key-swarm-abc",
                                   labels={"app": "heicode-swarm-agent"},
                                   secret_data={"OPENAI_API_KEY": "sk-secret-xyz", "GIT_PASSWORD": "ghp_secret_pw"})
    check("secret kind/type", sec["kind"] == "Secret" and sec["type"] == "Opaque")
    check("secret stringData has model key + git password",
          sec["stringData"]["OPENAI_API_KEY"] == "sk-secret-xyz" and sec["stringData"]["GIT_PASSWORD"] == "ghp_secret_pw")

    # Pod manifest references secret env 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",
                                secret_keys={"OPENAI_API_KEY", "GIT_PASSWORD"},
                                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"]["key"] == "OPENAI_API_KEY")
    check("GIT_PASSWORD via secretKeyRef (not inline)",
          "value" not in env["GIT_PASSWORD"] and env["GIT_PASSWORD"]["valueFrom"]["secretKeyRef"]["key"] == "GIT_PASSWORD")
    check("raw secret values NOT in pod manifest (only in Secret)",
          "sk-secret-xyz" not in blob and "ghp_secret_pw" not in blob)
    check("non-secret env inline (orch url + git repo url + git username)",
          env["ORCHESTRATOR_URL"]["value"] == "ws://orch.svc:8000"
          and env["GIT_REPO_URL"]["value"] == "https://git.example/acme/app.git"
          and env["GIT_USERNAME"]["value"] == "x-access-token")
    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 -> secret env omitted entirely (never inlined); non-secret git env still inline.
    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 + GIT_PASSWORD omitted from pod env",
          "OPENAI_API_KEY" not in env2 and "GIT_PASSWORD" not in env2)
    check("no secret -> non-secret GIT_REPO_URL/GIT_USERNAME still present", "GIT_REPO_URL" in env2 and "GIT_USERNAME" in env2)


def main():
    import asyncio
    test_launch_count()
    test_plan_specs()
    test_resolve_model_key()
    test_resolve_git_grant()
    test_azkv_resolver()
    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()