Agentswarm/orchestrator/agent_launcher.py

"""Swarm-side agent launcher (agent_swarm#16).

Team decision: the **Swarm runtime — not AM — launches the expert agent pool and enforces the
per-user limit**. On swarm-run create (decentralized seed flow), the orchestrator plans a
capability-diverse pool capped at `MAX_AGENTS_PER_USER`, composes each agent's launch env, and
spawns it via a pluggable backend. The launched agents connect back over WS (`/ws/{agent_id}`) and
self-select the seeded task.

Pluggable backend (`AGENT_LAUNCH_BACKEND`):
  - ``none``        (default): no-op — agents launched externally (preserves CI/e2e + lets a
                    deployment opt out). The orchestrator still seeds; P-guard reports
                    NO_AGENTS_CONNECTED if nothing connects.
  - ``subprocess``  : spawn ``python -m agent.main`` per agent (local/dev).
  - ``command``     : run a deployment-provided template ``AGENT_LAUNCH_CMD`` per agent (the env is
                    passed through; the template wraps the real spawn, e.g. a kubectl/pod-create).
  - ``kubernetes``  : **one Pod per agent** (production). Pods carry resource limits + labels
                    (`heicode-swarm-id`/`heicode-user-id`) for GC; the model key is delivered via a
                    per-swarm k8s **Secret** referenced by the pods (never inline PodSpec env, so it
                    is not exposed in etcd/`kubectl get -o yaml`/argv). Teardown deletes pods+secret
                    by label. Needs kubectl + a ServiceAccount with pod/secret RBAC in
                    AGENT_POD_NAMESPACE; ORCHESTRATOR_URL = in-cluster Service DNS.

Secret handling: the model key (per-user ``sk-``) is **resolved here from
``billing_context.secret_ref`` (azkv://)** and injected into the launched agent's env *server-side*
— it never travels in the create request body (runtime-contract §3.1) or in events/logs. Real
Azure Key Vault resolution is a deployment adapter (`SECRET_RESOLVER`); dev/CI use an env map.

Everything is **fail-soft**: a launch error is logged and swallowed so run creation never fails on
it. Pure planning (`plan_launch_specs`) is separated from side-effecting launch for testability.
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import shlex
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Default capability-diverse pool (mirrors the roles the swarm self-selects into). Override with
# AGENT_LAUNCH_CAPABILITIES (";"-separated per-agent capability CSVs).
DEFAULT_POOL_CAPABILITIES = [
    "python,code_generation,general",
    "testing,pytest,general",
    "technical-writing,documentation,general",
]


def launch_backend() -> str:
    return (os.getenv("AGENT_LAUNCH_BACKEND", "none") or "none").strip().lower()


def desired_pool_size() -> int:
    """How many agents to launch per run (before the per-user cap is applied)."""
    try:
        return max(0, int(os.getenv("AGENT_LAUNCH_POOL_SIZE", "3") or 3))
    except ValueError:
        return 3


def orchestrator_ws_url() -> str:
    """WS base the launched agent connects back to (deployment-set)."""
    return (os.getenv("ORCHESTRATOR_PUBLIC_URL")
            or os.getenv("AGENT_RUNTIME_WS_URL")
            or "ws://localhost:8000")


def pool_capabilities() -> List[str]:
    raw = os.getenv("AGENT_LAUNCH_CAPABILITIES")
    if raw:
        return [c.strip() for c in raw.split(";") if c.strip()]
    return list(DEFAULT_POOL_CAPABILITIES)


def resolve_model_key(body: Dict[str, Any]) -> Optional[str]:
    """Resolve the per-user model key (``sk-``) for the launched agents — server-side only.

    Precedence: explicit override (`AGENT_LAUNCH_MODEL_KEY`, dev) → `billing_context.secret_ref`
    (azkv:// via the deployment SecretResolver / dev env map `HEICODE_SECRET_<name>`) →
    orchestrator's own `OPENAI_API_KEY` (dev fallback). Returns None if unresolved (agents launch
    keyless and will error clearly — never fabricated).
    """
    override = os.getenv("AGENT_LAUNCH_MODEL_KEY")
    if override:
        return override
    secret_ref = ((body.get("billing_context") or {}).get("secret_ref") or "").strip()
    if secret_ref.startswith("azkv://"):
        resolved = _resolve_secret_ref(secret_ref)
        if resolved:
            return resolved
    return os.getenv("OPENAI_API_KEY") or None


#: Key Vault secret-value JSON field carrying the model key (HM #60 contract). HM writes the KV
#: secret as JSON ``{"openai_api_key": "sk-..."}`` — mirroring the callback-secret convention
#: (``{"callback_signing_secret": "..."}``) and leaving room for sibling fields later. A bare
#: ``sk-`` string value is still accepted (back-compat / hand-set dev secrets).
KV_MODEL_KEY_FIELD = "openai_api_key"


def _extract_model_key(secret_value: str) -> Optional[str]:
    """Pull the model key out of a resolved KV secret value.

    Accepts the contract JSON ``{"openai_api_key": "sk-..."}`` (HM #60) and, as a fallback, a bare
    ``sk-`` string. Never fabricates: unparseable/missing field → None.
    """
    value = (secret_value or "").strip()
    if not value:
        return None
    if value.startswith("{"):
        try:
            obj = json.loads(value)
        except ValueError:
            return None
        key = obj.get(KV_MODEL_KEY_FIELD) if isinstance(obj, dict) else None
        return str(key).strip() or None if key else None
    return value  # bare string secret value


def _resolve_secret_ref(secret_ref: str) -> Optional[str]:
    """Resolve an azkv:// ref to the model key.

    Order: (1) dev/CI env map ``HEICODE_SECRET_<name>`` — keeps tests/dev hermetic and offline;
    (2) production Azure Key Vault read via the pod's **workload identity** (``_resolve_from_keyvault``,
    agent_swarm#56). Either source yields the HM #60 JSON ``{"openai_api_key": "sk-..."}`` (bare string
    also accepted). Returns None when unavailable — we never fabricate a key.
    """
    name = secret_ref.rstrip("/").rsplit("/", 1)[-1]
    raw = os.getenv(f"HEICODE_SECRET_{name}")
    if raw is not None:
        return _extract_model_key(raw)
    return _resolve_from_keyvault(secret_ref)


def _azkv_enabled() -> bool:
    """Whether to attempt a real Key Vault read.

    Only when the orchestrator pod has **workload identity** injected (the AKS webhook sets
    ``AZURE_FEDERATED_TOKEN_FILE`` when the SA is annotated + the pod is labelled
    ``azure.workload.identity/use: "true"``), or a deployment explicitly opts in with
    ``SECRET_RESOLVER=azkv``. Keeps dev/CI/tests hermetic: without these the resolver never imports
    the azure SDK and never touches the network.
    """
    return bool(
        os.getenv("AZURE_FEDERATED_TOKEN_FILE")
        or os.getenv("SECRET_RESOLVER", "").strip().lower() == "azkv"
    )


def _parse_azkv_ref(secret_ref: str):
    """``azkv://<vault>/secrets/<name>[/<version>]`` → ``(vault_url, secret_name, version|None)``.

    ``<vault>`` may be a bare name (→ ``https://<name>.vault.azure.net``) or a full host. Also
    tolerates the short ``azkv://<vault>/<name>`` form. Returns None if it can't parse.
    """
    rest = secret_ref[len("azkv://"):].strip("/") if secret_ref.startswith("azkv://") else ""
    parts = [p for p in rest.split("/") if p]
    if len(parts) >= 3 and parts[1] == "secrets":
        host, name, version = parts[0], parts[2], (parts[3] if len(parts) > 3 else None)
    elif len(parts) == 2:
        host, name, version = parts[0], parts[1], None
    else:
        return None
    if not host or not name:
        return None
    vault_url = host if host.startswith("http") else (
        f"https://{host}" if "." in host else f"https://{host}.vault.azure.net")
    return vault_url, name, version


def _resolve_from_keyvault(secret_ref: str) -> Optional[str]:
    """Read the model key from Azure Key Vault using the pod's workload identity.

    `DefaultAzureCredential` picks up the federated token the AKS workload-identity webhook injects
    (see `_azkv_enabled`). Lazy-imports the azure SDK so dev/CI without it are unaffected. Returns
    None on ANY failure (not enabled / unparseable ref / SDK missing / no credential / network /
    secret absent) — never fabricates, never raises.
    """
    if not _azkv_enabled():
        return None
    parsed = _parse_azkv_ref(secret_ref)
    if not parsed:
        logger.warning("azkv resolver: unparseable secret_ref")
        return None
    vault_url, name, version = parsed
    try:
        from azure.identity import DefaultAzureCredential
        from azure.keyvault.secrets import SecretClient
    except Exception as exc:  # SDK not installed
        logger.warning("azkv resolver: azure SDK unavailable (%s); add azure-identity + azure-keyvault-secrets", exc)
        return None
    try:
        client = SecretClient(vault_url=vault_url, credential=DefaultAzureCredential())
        secret = client.get_secret(name, version) if version else client.get_secret(name)
        return _extract_model_key(secret.value)
    except Exception as exc:  # no credential / RBAC / network / missing secret
        logger.warning("azkv resolver: failed to read secret '%s' from %s: %s", name, vault_url, exc)
        return None


def model_api_base() -> str:
    return os.getenv("AGENT_OPENAI_API_BASE") or os.getenv("OPENAI_API_BASE") or "https://api.openai.com/v1"


def model_id(body: Dict[str, Any]) -> Optional[str]:
    plan = body.get("orchestration_plan") or {}
    return (((body.get("billing_context") or {}).get("default_model_id"))
            or plan.get("model_id") or body.get("model_id") or os.getenv("OPENAI_MODEL"))


@dataclass
class AgentLaunchSpec:
    agent_id: str
    capabilities: str               # CSV, as the agent reads AGENT_CAPABILITIES
    env: Dict[str, str] = field(default_factory=dict)


def launch_count(*, pool_size: int, limit: int, connected_user_agents: int) -> int:
    """How many to launch: pool size, but never pushing the user over the per-user cap."""
    return max(0, min(pool_size, limit - max(0, connected_user_agents)))


def plan_launch_specs(
    run,
    body: Dict[str, Any],
    *,
    connected_user_agents: int,
    limit: int,
    pool_size: int,
    model_key: Optional[str],
    orchestrator_url: str,
    user_id: Optional[str] = None,
) -> List[AgentLaunchSpec]:
    """Pure: plan the agent pool for a run (count capped by the per-user limit) + each agent's env.

    The model key is injected into the launch env (not the create body). Capabilities cycle through
    the configured pool so the pool is capability-diverse (the seed task has no required caps, so
    any agent can claim it; subtasks proposed later route by capability).
    """
    count = launch_count(pool_size=pool_size, limit=limit, connected_user_agents=connected_user_agents)
    caps = pool_capabilities() or ["general"]
    base = model_api_base()
    mid = model_id(body)
    specs: List[AgentLaunchSpec] = []
    for i in range(count):
        cap_csv = caps[i % len(caps)]
        env = {
            "ORCHESTRATOR_URL": orchestrator_url,
            "AGENT_ID": f"{run.swarm_id}-agent-{i+1}",
            "AGENT_CAPABILITIES": cap_csv,
            "OPENAI_API_BASE": base,
        }
        if model_key:
            env["OPENAI_API_KEY"] = model_key
        if mid:
            env["OPENAI_MODEL"] = mid
        if user_id:
            env["HEICODE_USER_ID"] = user_id
        specs.append(AgentLaunchSpec(agent_id=env["AGENT_ID"], capabilities=cap_csv, env=env))
    return specs


# Track subprocess-launched agents per run so they can be stopped (subprocess backend only).
_subprocess_agents: Dict[str, List[Any]] = {}


async def launch(specs: List[AgentLaunchSpec], *, swarm_id: str) -> List[str]:
    """Launch the planned agents via the configured backend. Fail-soft; returns launched agent ids."""
    backend = launch_backend()
    if backend == "none" or not specs:
        return []
    if backend == "kubernetes":
        # Pod-per-agent (production). Key goes into a per-swarm k8s Secret (referenced by the pods),
        # NOT inline PodSpec env — so it never lands in etcd/`kubectl get pod -o yaml`/argv.
        launched = await _launch_kubernetes_pool(specs, swarm_id)
        if launched:
            logger.info("launched %d agent pod(s) for swarm %s (backend=kubernetes)", len(launched), swarm_id)
        return launched
    launched: List[str] = []
    for spec in specs:
        try:
            if backend == "subprocess":
                await _launch_subprocess(spec, swarm_id)
            elif backend == "command":
                await _launch_command(spec)
            else:
                logger.warning("unknown AGENT_LAUNCH_BACKEND=%s; skipping launch", backend)
                continue
            launched.append(spec.agent_id)
        except Exception as exc:  # fail-soft: never break run creation on a launch error
            logger.warning("agent launch failed for %s (backend=%s): %s", spec.agent_id, backend, exc)
    if launched:
        logger.info("launched %d agent(s) for swarm %s (backend=%s)", len(launched), swarm_id, backend)
    return launched


async def _launch_subprocess(spec: AgentLaunchSpec, swarm_id: str) -> None:
    import sys
    child_env = {**os.environ, **spec.env}
    proc = await asyncio.create_subprocess_exec(
        sys.executable, "-m", "agent.main",
        env=child_env,
        stdout=asyncio.subprocess.DEVNULL,
        stderr=asyncio.subprocess.DEVNULL,
    )
    _subprocess_agents.setdefault(swarm_id, []).append(proc)


def build_launch_command(spec: AgentLaunchSpec) -> List[str]:
    """Build the argv for the `command` backend from the `AGENT_LAUNCH_CMD` template.

    The template is `shlex`-split; ``{agent_id}`` / ``{capabilities}`` are substituted. The agent
    env is passed via process env (see _launch_command), so the template wraps the real spawn
    (e.g. a kubectl/pod-create script) without secrets on the argv.
    """
    template = os.getenv("AGENT_LAUNCH_CMD", "")
    argv = shlex.split(template)
    return [a.format(agent_id=spec.agent_id, capabilities=spec.capabilities) for a in argv]


async def _launch_command(spec: AgentLaunchSpec) -> None:
    argv = build_launch_command(spec)
    if not argv:
        raise ValueError("AGENT_LAUNCH_CMD is empty but backend=command")
    child_env = {**os.environ, **spec.env}
    proc = await asyncio.create_subprocess_exec(*argv, env=child_env)
    await proc.wait()  # the template is a quick spawn-wrapper (e.g. kubectl apply), not the agent itself


# ── Kubernetes backend ────────────────────────────────────────────────────────
# Production runs agents as **one Pod per agent** in the swarm cluster. The model key is delivered
# via a per-swarm k8s Secret referenced by the pods (NOT inline PodSpec env → not exposed in
# etcd/`kubectl get -o yaml`/argv). Requires: kubectl in the orchestrator image + a ServiceAccount
# with RBAC to create/delete pods+secrets in AGENT_POD_NAMESPACE; ORCHESTRATOR_URL must be the
# in-cluster Service DNS (e.g. ws://swarm-orchestrator.<ns>.svc.cluster.local:8000). NetworkPolicy
# egress to HM /v1 + git. Hardened alternative: mount the azkv secret via a CSI SecretProviderClass
# (orchestrator never touches plaintext) — deployment/Infra option, see security-boundary §6.
_k8s_swarms: set = set()


def pod_image() -> str:
    return os.getenv("AGENT_POD_IMAGE", "heicode/swarm-agent:latest")


def pod_namespace() -> str:
    return os.getenv("AGENT_POD_NAMESPACE", "heicode-swarm")


def pod_resources() -> Dict[str, Any]:
    return {
        "requests": {"cpu": os.getenv("AGENT_POD_CPU_REQUEST", "250m"),
                     "memory": os.getenv("AGENT_POD_MEM_REQUEST", "256Mi")},
        "limits": {"cpu": os.getenv("AGENT_POD_CPU_LIMIT", "1"),
                   "memory": os.getenv("AGENT_POD_MEM_LIMIT", "1Gi")},
    }


def _k8s_labels(spec: AgentLaunchSpec, swarm_id: str) -> Dict[str, str]:
    labels = {"app": "heicode-swarm-agent", "heicode-swarm-id": swarm_id}
    uid = spec.env.get("HEICODE_USER_ID")
    if uid:
        labels["heicode-user-id"] = uid
    return labels


def build_secret_manifest(*, key_value: str, namespace: str, name: str, labels: Dict[str, str]) -> Dict[str, Any]:
    """k8s Secret carrying the per-user model key (applied via stdin → value never on argv)."""
    return {
        "apiVersion": "v1", "kind": "Secret", "type": "Opaque",
        "metadata": {"name": name, "namespace": namespace, "labels": labels},
        "stringData": {"OPENAI_API_KEY": key_value},
    }


def build_pod_manifest(spec: AgentLaunchSpec, *, namespace: str, swarm_id: str,
                       image: Optional[str] = None, secret_name: Optional[str] = None,
                       service_account: Optional[str] = None,
                       resources: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """One agent Pod. Non-sensitive env inline; `OPENAI_API_KEY` via `secretKeyRef` (never inline)."""
    env_list: List[Dict[str, Any]] = []
    for k, v in spec.env.items():
        if k == "OPENAI_API_KEY":
            if secret_name:  # reference the Secret; never inline the key into the PodSpec
                env_list.append({"name": k, "valueFrom": {"secretKeyRef": {"name": secret_name, "key": "OPENAI_API_KEY"}}})
        else:
            env_list.append({"name": k, "value": v})
    pod_spec: Dict[str, Any] = {
        "restartPolicy": "OnFailure",
        "containers": [{
            "name": "agent",
            "image": image or pod_image(),
            "command": ["python", "-m", "agent.main"],
            "env": env_list,
            "resources": resources or pod_resources(),
        }],
    }
    if service_account:
        pod_spec["serviceAccountName"] = service_account
    return {
        "apiVersion": "v1", "kind": "Pod",
        "metadata": {"name": spec.agent_id, "namespace": namespace, "labels": _k8s_labels(spec, swarm_id)},
        "spec": pod_spec,
    }


async def _kubectl_apply(manifest: Dict[str, Any]) -> None:
    proc = await asyncio.create_subprocess_exec(
        "kubectl", "apply", "-f", "-",
        stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE,
    )
    _, err = await proc.communicate(json.dumps(manifest).encode())
    if proc.returncode != 0:
        raise RuntimeError(f"kubectl apply failed: {(err or b'').decode()[:200]}")


async def _launch_kubernetes_pool(specs: List[AgentLaunchSpec], swarm_id: str) -> List[str]:
    ns, sa, res, img = pod_namespace(), pod_service_account(), pod_resources(), pod_image()
    labels = _k8s_labels(specs[0], swarm_id)
    secret_name = None
    key_value = next((s.env.get("OPENAI_API_KEY") for s in specs if s.env.get("OPENAI_API_KEY")), None)
    if key_value:
        secret_name = f"swarm-agent-key-{swarm_id}"
        try:
            await _kubectl_apply(build_secret_manifest(key_value=key_value, namespace=ns, name=secret_name, labels=labels))
        except Exception as exc:  # without the secret the pods can't get the key; still try keyless (clear error)
            logger.warning("k8s secret apply failed for %s: %s", swarm_id, exc)
            secret_name = None
    launched: List[str] = []
    for spec in specs:
        try:
            await _kubectl_apply(build_pod_manifest(spec, namespace=ns, swarm_id=swarm_id, image=img,
                                                    secret_name=secret_name, service_account=sa, resources=res))
            launched.append(spec.agent_id)
        except Exception as exc:
            logger.warning("k8s agent pod apply failed for %s: %s", spec.agent_id, exc)
    if launched:
        _k8s_swarms.add(swarm_id)
    return launched


def pod_service_account() -> Optional[str]:
    return os.getenv("AGENT_POD_SERVICE_ACCOUNT") or None


async def stop_launched(swarm_id: str) -> int:
    """Best-effort teardown of agents launched for a run. subprocess → terminate; kubernetes →
    delete pods+secret by label. `command` backend pods are deployment-managed."""
    stopped = 0
    for proc in _subprocess_agents.pop(swarm_id, []):
        try:
            if proc.returncode is None:
                proc.terminate()
                stopped += 1
        except Exception as exc:
            logger.warning("failed to stop launched agent proc for %s: %s", swarm_id, exc)
    if swarm_id in _k8s_swarms:
        _k8s_swarms.discard(swarm_id)
        try:
            proc = await asyncio.create_subprocess_exec(
                "kubectl", "delete", "pod,secret", "-l", f"heicode-swarm-id={swarm_id}",
                "-n", pod_namespace(), "--ignore-not-found=true",
                stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.DEVNULL,
            )
            await proc.wait()
        except Exception as exc:
            logger.warning("k8s teardown failed for %s: %s", swarm_id, exc)
    return stopped