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():
    # Default hard pod-count window is [MIN, MAX] = [3, 16]. count = max(3, min(16, min(pool, cap-connected))).
    check("count = max(MIN, min(pool, limit-connected)) — pool 3, no connected -> 3",
          al.launch_count(pool_size=3, limit=10, connected_user_agents=0) == 3)
    # MIN floor takes PRIORITY over cap headroom: pool 5, cap 10, 8 connected -> headroom 2, floored to MIN 3.
    check("MIN floor (3) wins over cap headroom (headroom 2 -> 3)",
          al.launch_count(pool_size=5, limit=10, connected_user_agents=8) == 3)
    # Even when the user is already AT/over the cap, the MIN floor still applies (intentional [3,16]
    # hard-floor-first rule; WS register-time cap fail-closes any pods that exceed the cap).
    check("MIN floor still applies when already at cap (headroom 0 -> 3)",
          al.launch_count(pool_size=3, limit=10, connected_user_agents=10) == 3)
    check("count clamps to pool when cap is high (pool 3 -> 3)",
          al.launch_count(pool_size=3, limit=100, connected_user_agents=0) == 3)
    # MAX ceiling (16): a big pool with plenty of headroom is capped at 16.
    check("MAX ceiling (16) caps a large pool",
          al.launch_count(pool_size=50, limit=100, connected_user_agents=0) == 16)
    # A high pool but limited cap headroom: pool 50, cap 20, 12 connected -> headroom 8, within [3,16] -> 8.
    check("cap headroom (8) honored when within [MIN,MAX]",
          al.launch_count(pool_size=50, limit=20, connected_user_agents=12) == 8)
    # env-overridable window: MIN=1, MAX=4 -> headroom 2 honored (not floored), big pool capped at 4.
    os.environ["AGENT_LAUNCH_MIN_POOL"] = "1"
    os.environ["AGENT_LAUNCH_MAX_POOL"] = "4"
    check("env MIN=1 -> small headroom not floored",
          al.launch_count(pool_size=5, limit=10, connected_user_agents=8) == 2)
    check("env MAX=4 -> large pool capped at 4",
          al.launch_count(pool_size=50, limit=100, connected_user_agents=0) == 4)
    os.environ.pop("AGENT_LAUNCH_MIN_POOL")
    os.environ.pop("AGENT_LAUNCH_MAX_POOL")


def test_pool_window():
    # desired_pool_size() is clamped to [MIN, MAX] = [3, 16] by default.
    for k in ("AGENT_LAUNCH_POOL_SIZE", "AGENT_LAUNCH_MIN_POOL", "AGENT_LAUNCH_MAX_POOL"):
        os.environ.pop(k, None)
    check("desired_pool_size default = 3 (== MIN)", al.desired_pool_size() == 3)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "1"
    check("desired_pool_size floored to MIN 3 (requested 1)", al.desired_pool_size() == 3)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "99"
    check("desired_pool_size capped to MAX 16 (requested 99)", al.desired_pool_size() == 16)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "8"
    check("desired_pool_size passes through within window (8)", al.desired_pool_size() == 8)
    # env-overridable window.
    os.environ["AGENT_LAUNCH_MIN_POOL"] = "5"
    os.environ["AGENT_LAUNCH_MAX_POOL"] = "6"
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "1"
    check("env MIN=5 floors desired_pool_size", al.desired_pool_size() == 5)
    os.environ["AGENT_LAUNCH_POOL_SIZE"] = "20"
    check("env MAX=6 caps desired_pool_size", al.desired_pool_size() == 6)
    check("min_pool_size()/max_pool_size() read env", al.min_pool_size() == 5 and al.max_pool_size() == 6)
    for k in ("AGENT_LAUNCH_POOL_SIZE", "AGENT_LAUNCH_MIN_POOL", "AGENT_LAUNCH_MAX_POOL"):
        os.environ.pop(k, None)


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)
    # [3,16] pod-count window (#66): MIN floor (3) wins over cap headroom (10-8=2) -> launch 3.
    check("plan floors to MIN pool (8 connected, cap 10, headroom 2 -> floor 3)", len(specs) == 3)
    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))
    # #70: per-task timeout is transmitted into every launched agent's env (non-secret, inline).
    check("spec env has TASK_TIMEOUT_SECONDS (transmitted, #70)", s.env.get("TASK_TIMEOUT_SECONDS") == str(al.DEFAULT_TASK_TIMEOUT_SECONDS))
    # 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_task_timeout():
    # #70: per-task execution timeout transmitted to launched agents.
    os.environ.pop("TASK_TIMEOUT_SECONDS", None)
    check("default timeout = 300 (raised from 60)", al.resolve_task_timeout({}) == al.DEFAULT_TASK_TIMEOUT_SECONDS == 300)
    # env override.
    os.environ["TASK_TIMEOUT_SECONDS"] = "450"
    check("env TASK_TIMEOUT_SECONDS honored", al.resolve_task_timeout({}) == 450)
    # capped at run budget.duration_seconds (take the smaller).
    body = {"orchestration_plan": {"budget": {"duration_seconds": 120}}}
    check("timeout capped at budget.duration_seconds (smaller wins)", al.resolve_task_timeout(body) == 120)
    body_big = {"orchestration_plan": {"budget": {"duration_seconds": 3600}}}
    check("budget larger than base -> base wins", al.resolve_task_timeout(body_big) == 450)
    os.environ.pop("TASK_TIMEOUT_SECONDS")
    # max_duration_seconds alias also honored.
    check("max_duration_seconds alias honored",
          al.resolve_task_timeout({"orchestration_plan": {"budget": {"max_duration_seconds": 90}}}) == 90)
    # bad/zero env -> falls back to default; absent budget -> default.
    os.environ["TASK_TIMEOUT_SECONDS"] = "0"
    check("non-positive env -> default", al.resolve_task_timeout({}) == al.DEFAULT_TASK_TIMEOUT_SECONDS)
    os.environ["TASK_TIMEOUT_SECONDS"] = "notanint"
    check("non-int env -> default", al.resolve_task_timeout({}) == al.DEFAULT_TASK_TIMEOUT_SECONDS)
    os.environ.pop("TASK_TIMEOUT_SECONDS")


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_pool_window()
    test_plan_specs()
    test_resolve_task_timeout()
    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()