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.

Git binding (agent_swarm#63 / HM #92): a swarm-run needs a repo to operate on. HM ships the git
binding as a ``resource_type=git`` entry in ``resource_grants`` — ``metadata.repo_url`` (clone URL,
non-secret) + ``secret_ref`` (``azkv://`` ref to the git credential, never plaintext). The launcher
resolves it here the same server-side way as the model key (``resolve_git_grant``) and injects
``GIT_REPO_URL`` (+ ``GIT_USERNAME``/``GIT_PASSWORD`` when a credential resolves) into the agent env;
the agent (`agent/main.py` + `agent/git_operations.py`) clones from those vars. Without this the
agent has no repo URL → cannot clone → the "需绑定 git 仓库才能使用" symptom in HM #92.

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 min_pool_size() -> int:
    """Hard floor on the launched-pod count, env-tunable (AGENT_LAUNCH_MIN_POOL, default 3)."""
    try:
        return max(0, int(os.getenv("AGENT_LAUNCH_MIN_POOL", "3") or 3))
    except ValueError:
        return 3


def max_pool_size() -> int:
    """Hard ceiling on the launched-pod count, env-tunable (AGENT_LAUNCH_MAX_POOL, default 16)."""
    try:
        return max(1, int(os.getenv("AGENT_LAUNCH_MAX_POOL", "16") or 16))
    except ValueError:
        return 16


def desired_pool_size() -> int:
    """How many agents to launch per run (before the per-user cap is applied), clamped to
    [AGENT_LAUNCH_MIN_POOL, AGENT_LAUNCH_MAX_POOL] = [3, 16] by default."""
    try:
        raw = max(0, int(os.getenv("AGENT_LAUNCH_POOL_SIZE", "3") or 3))
    except ValueError:
        raw = 3
    lo, hi = min_pool_size(), max_pool_size()
    if lo > hi:  # defensive: keep the window sane if mis-set
        lo = hi
    return max(lo, min(hi, raw))


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 _read_kv_secret_value(secret_ref: str) -> Optional[str]:
    """Read the **raw** secret value string from Azure Key Vault using the pod's workload identity.

    Shared by the model-key and git-credential resolvers (each applies its own extractor to the raw
    value). `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 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 _resolve_from_keyvault(secret_ref: str) -> Optional[str]:
    """Read the model key from Azure Key Vault (raw value → `_extract_model_key`). None on failure."""
    raw = _read_kv_secret_value(secret_ref)
    return _extract_model_key(raw) if raw is not None else None


# ── Git binding resolution (agent_swarm#63 / HM #92) ───────────────────────────
#: Default git username when the KV secret carries only a token/PAT (GitHub HTTPS convention: any
#: non-empty username + PAT-as-password works; ``x-access-token`` is the documented placeholder).
DEFAULT_GIT_USERNAME = "x-access-token"


def _extract_git_credentials(secret_value: str) -> Optional[Dict[str, str]]:
    """Pull git credentials out of a resolved KV secret value.

    Contract (mirrors the model-key JSON convention; pending HM #92 alignment): the git KV secret is
    JSON ``{"git_username": "...", "git_password": "..."}``. Accepted aliases:
    ``username``/``password``, and a token form ``git_token``/``token`` (username defaults to
    ``x-access-token``). A bare string value is treated as a token. Returns
    ``{"username","password"}`` or None (never fabricates — missing password → None).
    """
    value = (secret_value or "").strip()
    if not value:
        return None
    if value.startswith("{"):
        try:
            obj = json.loads(value)
        except ValueError:
            return None
        if not isinstance(obj, dict):
            return None
        username = (obj.get("git_username") or obj.get("username") or "").strip() or DEFAULT_GIT_USERNAME
        password = (obj.get("git_password") or obj.get("git_token")
                    or obj.get("password") or obj.get("token") or "").strip()
        return {"username": username, "password": password} if password else None
    return {"username": DEFAULT_GIT_USERNAME, "password": value}  # bare token string


def _resolve_git_secret_ref(secret_ref: str) -> Optional[Dict[str, str]]:
    """Resolve an azkv:// git ``secret_ref`` to ``{"username","password"}``.

    Same source order as the model key: dev/CI env map ``HEICODE_SECRET_<name>`` (hermetic) →
    production Azure Key Vault via workload identity. None when unavailable — never fabricated.
    """
    name = secret_ref.rstrip("/").rsplit("/", 1)[-1]
    raw = os.getenv(f"HEICODE_SECRET_{name}")
    if raw is None:
        raw = _read_kv_secret_value(secret_ref)
    return _extract_git_credentials(raw) if raw is not None else None


def _iter_resource_grants(body: Dict[str, Any]):
    """Yield every resource grant in the create body (top-level + per-agent), in declaration order."""
    yield from (body.get("resource_grants") or [])
    plan = body.get("orchestration_plan") or {}
    for agent in (plan.get("agents") or []):
        yield from (agent.get("resource_grants") or [])
    for agent in (body.get("agents") or []):
        yield from (agent.get("resource_grants") or [])


def _first_git_grant(body: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    """First grant with ``resource_type`` (or ``type``) == ``git``, or None."""
    for grant in _iter_resource_grants(body):
        if not isinstance(grant, dict):
            continue
        if (grant.get("resource_type") or grant.get("type")) == "git":
            return grant
    return None


def resolve_git_grant(body: Dict[str, Any]) -> Optional[Dict[str, str]]:
    """Resolve the git binding for launched agents — server-side, mirroring `resolve_model_key`.

    Returns the git env to inject (``GIT_REPO_URL`` + optional ``GIT_USERNAME``/``GIT_PASSWORD`` +
    optional ``GIT_BASE_BRANCH``) or None when there is no git grant / no repo URL. The credential
    (``secret_ref`` → azkv) is resolved here and never travels in the create body / events / logs.
    A grant with a repo URL but unresolved credential still yields ``GIT_REPO_URL`` (public-repo
    clone works; a private repo then fails to clone with a clear git error — never fabricated).
    """
    grant = _first_git_grant(body)
    if not grant:
        return None
    meta = grant.get("metadata") or {}
    repo_url = (meta.get("repo_url") or grant.get("repo_url") or "").strip()
    if not repo_url:
        return None
    env: Dict[str, str] = {"GIT_REPO_URL": repo_url}
    branch = (meta.get("base_branch") or meta.get("branch") or "").strip()
    if branch:
        env["GIT_BASE_BRANCH"] = branch
    # SWE-bench / benchmark: pin the exact commit the agent must work from. The agent shallow-fetches
    # and checks out this commit so reads/edits are against the correct baseline (not default HEAD).
    base_commit = (meta.get("base_commit") or grant.get("base_commit") or "").strip()
    if base_commit:
        env["GIT_BASE_COMMIT"] = base_commit
    secret_ref = (grant.get("secret_ref") or grant.get("ref") or "").strip()
    if secret_ref.startswith("azkv://"):
        creds = _resolve_git_secret_ref(secret_ref)
        if creds:
            env["GIT_USERNAME"] = creds["username"]
            env["GIT_PASSWORD"] = creds["password"]
        else:
            logger.warning("git grant: secret_ref present but unresolved; agent will clone keyless (private repo will fail)")
    return env


# ── Per-task execution timeout (agent_swarm#70) ────────────────────────────────
#: Default per-task timeout injected into launched agents. 60s (the old agent default) reliably
#: killed generation-class tasks before the model finished; raised to 300s. Overridable via
#: TASK_TIMEOUT_SECONDS on the orchestrator. Aligned (min) with the run's budget.duration_seconds
#: so a task can never outlive its run budget.
DEFAULT_TASK_TIMEOUT_SECONDS = 300


def _budget_duration_seconds(body: Dict[str, Any]) -> Optional[int]:
    """The run's wall-clock budget in seconds from the create body, or None when absent."""
    budget = ((body.get("orchestration_plan") or {}).get("budget") or {})
    raw = budget.get("duration_seconds") or budget.get("max_duration_seconds")
    try:
        val = int(raw)
    except (TypeError, ValueError):
        return None
    return val if val > 0 else None


def resolve_task_timeout(body: Dict[str, Any]) -> int:
    """Per-task execution timeout (seconds) to inject into each launched agent (agent_swarm#70).

    Base = ``TASK_TIMEOUT_SECONDS`` env (default ``DEFAULT_TASK_TIMEOUT_SECONDS`` = 300). When the
    run declares ``budget.duration_seconds`` we take the **smaller** of the two so a single task can
    never outlive the whole run budget. Always returns a positive int.
    """
    try:
        base = int(os.getenv("TASK_TIMEOUT_SECONDS", str(DEFAULT_TASK_TIMEOUT_SECONDS)))
    except ValueError:
        base = DEFAULT_TASK_TIMEOUT_SECONDS
    if base <= 0:
        base = DEFAULT_TASK_TIMEOUT_SECONDS
    budget_duration = _budget_duration_seconds(body)
    if budget_duration:
        return min(base, budget_duration)
    return base


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"))


def git_launch_env(body: Dict[str, Any]) -> Dict[str, str]:
    """Translate a bound git resource grant into the agent's workspace env so the agent clones the
    repo into WORKSPACE_DIR before working (without it /workspace is empty and any "read/repair the
    repo" task fails with "no source files to analyze"). Pure: reads resource_grants only, no I/O.

    Sets `GIT_REPO_URL` (+ `GIT_PROVIDER` / `GIT_DEFAULT_BRANCH` when present) from the grant, mirroring
    the HM git-binding → env contract (agent-capability-schema §3 / agent_template_env.go). PUBLIC repos
    clone with just `GIT_REPO_URL` (no credentials). Private-repo credential injection (`GIT_TOKEN` via
    the grant's azkv `secret_ref`, routed through the per-swarm k8s Secret) is intentionally NOT handled
    here yet — tracked as a follow-up. Returns {} when no git grant is bound.
    """
    grants: List[Dict[str, Any]] = []
    grants += body.get("resource_grants") or []
    plan = body.get("orchestration_plan") or {}
    grants += plan.get("resource_grants") or []
    for agent in (plan.get("agents") or []):
        grants += agent.get("resource_grants") or []
    for grant in grants:
        if not isinstance(grant, dict):
            continue
        meta = grant.get("metadata") or {}
        rtype = str(grant.get("resource_type") or grant.get("type") or "").lower()
        repo = str((meta or {}).get("repo_url") or "").strip()
        if not repo and rtype == "git":
            scope = str(grant.get("binding_scope") or "").strip()
            if scope.startswith(("http://", "https://", "git@", "ssh://")):
                repo = scope
        if not repo:
            continue
        env = {"GIT_REPO_URL": repo}
        provider = str((meta or {}).get("provider") or "").strip()
        if provider:
            env["GIT_PROVIDER"] = provider
        branch = str((meta or {}).get("default_branch") or "").strip()
        if branch:
            env["GIT_DEFAULT_BRANCH"] = branch
            env["GIT_BASE_BRANCH"] = branch  # agent/git_operations reads GIT_BASE_BRANCH for the base ref
        return env
    return {}


@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 agents to launch, with a HARD pod-count window [MIN, MAX] = [3, 16] by default.

        count = max(MIN, min(MAX, min(pool_size, limit - max(0, connected))))

    `limit` is the per-user cap (env MAX_AGENTS_PER_USER or the run's metadata.max_agents_per_user;
    see main.max_agents_per_user). `pool_size` is desired_pool_size() (already clamped to [MIN, MAX]).

    INTENTIONAL boundary: the MIN floor (3) takes priority over the per-user-cap headroom — when a
    user already has many agents connected, `limit - connected` can be < MIN, yet we still floor to
    MIN. This is the architect's [3,16] hard-floor-first rule. The WS register-time per-user cap
    (main.websocket_endpoint) still hard-rejects connections beyond the SAME cap, so any pods launched
    above `limit` simply fail to register (fail-closed) rather than over-provisioning the user."""
    lo, hi = min_pool_size(), max_pool_size()
    if lo > hi:  # defensive: keep the window sane if mis-set
        lo = hi
    headroom = min(pool_size, limit - max(0, connected_user_agents))
    return max(lo, min(hi, headroom))


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,
    git_env: Optional[Dict[str, str]] = None,
    count_override: Optional[int] = None,
    id_start: int = 0,
) -> 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). `git_env` (from
    git_launch_env) carries GIT_REPO_URL so each agent clones the bound repo into its workspace.

    Elastic top-up: pass ``count_override`` (the delta to add) + ``id_start`` (current agent count)
    to plan ADDITIONAL agents with non-colliding ids (``agent-{id_start+i+1}``). The caller is
    responsible for capping the delta against MAX_POOL and the per-user limit (the autoscale loop).
    """
    count = count_override if count_override is not None else 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)
    task_timeout = resolve_task_timeout(body)  # #70:透传执行超时，与 run budget.duration_seconds 取较小
    specs: List[AgentLaunchSpec] = []
    for i in range(count):
        idx = id_start + i
        cap_csv = caps[idx % len(caps)]
        env = {
            "ORCHESTRATOR_URL": orchestrator_url,
            "AGENT_ID": f"{run.swarm_id}-agent-{idx+1}",
            "AGENT_CAPABILITIES": cap_csv,
            "OPENAI_API_BASE": base,
            # Non-secret execution-timeout knob; inline like the other config env (never a Secret).
            "TASK_TIMEOUT_SECONDS": str(task_timeout),
        }
        if model_key:
            env["OPENAI_API_KEY"] = model_key
        if mid:
            env["OPENAI_MODEL"] = mid
        if user_id:
            env["HEICODE_USER_ID"] = user_id
        # Jina MCP key — passed through from orchestrator env so launched agents can call Jina
        # search/read tools. Secret-classed (SENSITIVE_ENV_KEYS) → delivered via per-swarm Secret.
        jina_key = os.getenv("JINA_API_KEY")
        if jina_key:
            env["JINA_API_KEY"] = jina_key
        if git_env:
            env.update(git_env)
        # LOCAL TEST ONLY (not for upstream): propagate the agent-side subtask-handoff switch so
        # launched agents decompose + hand off subtasks to peers (else each agent does its task alone).
        if os.getenv("ENABLE_SUBTASK_HANDOFF", "false").lower() in {"1", "true", "yes"}:
            env["ENABLE_SUBTASK_HANDOFF"] = "true"
        # 透传 plan/分解 token 上限到 agent pod：部分模型网关上限 32768，>之 400 → 分解调用失败、
        # fan-out 断。deployment 设 AGENT_PLAN_MAX_TOKENS 时透传；未设则不传，agent 用自身默认。
        plan_max_tokens = os.getenv("AGENT_PLAN_MAX_TOKENS")
        if plan_max_tokens:
            env["AGENT_PLAN_MAX_TOKENS"] = plan_max_tokens
        # 透传真实仓上下文预算(读多少文件/每文件多少字节)给 agent pod
        for _ctx in ("AGENT_CTX_MAX_FILES", "AGENT_CTX_MAX_BYTES"):
            _v = os.getenv(_ctx)
            if _v:
                env[_ctx] = _v
        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()

#: Env keys that carry secrets — delivered via the per-swarm k8s Secret (`secretKeyRef`), NEVER
#: inlined into the PodSpec (would leak into etcd / `kubectl get pod -o yaml`). `GIT_USERNAME` and
#: `GIT_REPO_URL` are non-secret and stay inline.
SENSITIVE_ENV_KEYS = ("OPENAI_API_KEY", "GIT_PASSWORD", "JINA_API_KEY")


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 pod_annotations() -> Dict[str, str]:
    """Annotations applied to every agent Pod.

    `karpenter.sh/do-not-disrupt` keeps the node autoscaler from consolidating/evicting an agent
    while its run is still live — mid-run eviction is what orphaned an agent on 2026-06-15 (its
    siblings were reaped as "Underutilized", leaving a lone idle agent). The run's own teardown
    (stop_launched) removes the pods when the run reaches a terminal state, so this only protects
    in-flight work. Set AGENT_POD_ALLOW_DISRUPTION=1 to opt out (e.g. cost-sensitive dev clusters).
    """
    if os.getenv("AGENT_POD_ALLOW_DISRUPTION", "").lower() in {"1", "true", "yes"}:
        return {}
    return {"karpenter.sh/do-not-disrupt": "true"}


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(*, namespace: str, name: str, labels: Dict[str, str],
                          secret_data: Dict[str, str]) -> Dict[str, Any]:
    """k8s Secret carrying the per-user sensitive env (model key + git password) — applied via stdin
    so values never hit argv. `secret_data` maps env key → value (only SENSITIVE_ENV_KEYS)."""
    return {
        "apiVersion": "v1", "kind": "Secret", "type": "Opaque",
        "metadata": {"name": name, "namespace": namespace, "labels": labels},
        "stringData": dict(secret_data),
    }


def build_pod_manifest(spec: AgentLaunchSpec, *, namespace: str, swarm_id: str,
                       image: Optional[str] = None, secret_name: Optional[str] = None,
                       secret_keys=None, service_account: Optional[str] = None,
                       resources: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
    """One agent Pod. Non-sensitive env inline; secret env (`OPENAI_API_KEY`/`GIT_PASSWORD`) via
    `secretKeyRef` (never inline). `secret_keys` = the keys actually present in the Secret."""
    keys_in_secret = set(secret_keys or ())
    env_list: List[Dict[str, Any]] = []
    for k, v in spec.env.items():
        if k in SENSITIVE_ENV_KEYS:
            if secret_name and k in keys_in_secret:  # reference the Secret; never inline into PodSpec
                env_list.append({"name": k, "valueFrom": {"secretKeyRef": {"name": secret_name, "key": k}}})
            # else: omit entirely — a secret value is never inlined
        else:
            env_list.append({"name": k, "value": v})
    pod_spec: Dict[str, Any] = {
        "restartPolicy": "OnFailure",
        "containers": [{
            "name": "agent",
            "image": image or pod_image(),
            "imagePullPolicy": "Always",  # always pull the latest :heicode-test tag (else node caches stale)
            "command": ["python", "-m", "agent.main"],
            "env": env_list,
            "resources": resources or pod_resources(),
        }],
    }
    if service_account:
        pod_spec["serviceAccountName"] = service_account
    metadata: Dict[str, Any] = {
        "name": spec.agent_id, "namespace": namespace, "labels": _k8s_labels(spec, swarm_id),
    }
    annotations = pod_annotations()
    if annotations:
        metadata["annotations"] = annotations
    return {
        "apiVersion": "v1", "kind": "Pod",
        "metadata": metadata,
        "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
    # Sensitive env (model key + git password) is per-user/per-run — same across the pool; collect
    # the first non-empty value of each into one per-swarm Secret.
    secret_data: Dict[str, str] = {}
    for skey in SENSITIVE_ENV_KEYS:
        val = next((s.env.get(skey) for s in specs if s.env.get(skey)), None)
        if val:
            secret_data[skey] = val
    if secret_data:
        secret_name = f"swarm-agent-key-{swarm_id}"
        try:
            await _kubectl_apply(build_secret_manifest(namespace=ns, name=secret_name, labels=labels, secret_data=secret_data))
        except Exception as exc:  # without the secret the pods can't get the values; still try without (clear error)
            logger.warning("k8s secret apply failed for %s: %s", swarm_id, exc)
            secret_name = None
            secret_data = {}
    secret_keys = set(secret_data)
    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, secret_keys=secret_keys,
                                                    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)
    # The in-memory `_k8s_swarms` set is lost on orchestrator restart, so don't gate teardown on
    # it: when the kubernetes backend is active, delete by label authoritatively. The label
    # selector + --ignore-not-found makes this idempotent and safe to call for any swarm_id
    # (including runs launched before a restart — which is exactly how agents got orphaned).
    if swarm_id in _k8s_swarms or launch_backend() == "kubernetes":
        _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