Swarm 负责拉起 agent + 执行限额（team 决议，反转 AM 拉起）（Refs #16）
团队决议：由 Swarm 运行时（非 AM）拉起专家 agent 池并执行每用户限额。 代码： - orchestrator/agent_launcher.py（新）：plan_launch_specs（纯，按 min(池大小, MAX_AGENTS_PER_USER −已连) 限额 + 组装每 agent env）、resolve_model_key（override→azkv secret_ref 解析(部署 SecretResolver/dev HEICODE_SECRET_<name>)→OPENAI_API_KEY 兜底，解析不到不伪造）、可插拔后端 launch()（none 默认/subprocess/command 模板，fail-soft）、stop_launched。 - orchestrator/main.py：create 播种后调 launch_swarm_agents（仅去中心化、非 Manager 显式 agent； 从 create x-user-id 取 user_id；key 服务端解析，不入 create 体）；stop_swarm_run 调 stop_launched。 文档：runtime-contract §3.3 由「AM 拉起（提案待确认）」改为「Swarm 拉起 + 限额（已定）」， 更新 env 来源列（key=Swarm 从 secret_ref 解析、AGENT_ID/CAPABILITIES=Swarm launcher、 HEICODE_USER_ID=从 create 透传）+ 后端/限额/解析约束；security-boundary §6 增 Swarm 拉起 + 服务端解析 key（不上 argv/日志）说明。 测试：scripts/test-agent-launcher.py（限额封顶、env 组装、key 解析优先级、command 模板、 none no-op）接入 CI。e2e/contract 回归通过（默认 backend=none，行为不变）。 影响范围：仅 agent_swarm（orchestrator + docs + 测试 + CI）。默认 backend=none 不自动拉起、 向后兼容；密钥仅服务端 env 注入、不入 create 体/回调/日志/argv（满足 §3.1 + security-boundary）。 不改 Manager↔Swarm 契约鉴权/计费账本/审批链。 Refs #16 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Songhaoz666
co-authored by
Claude Opus 4.8
Claude Opus 4.8
parent
fb367c4373
commit
9516b5c696
@@ -59,6 +59,9 @@ jobs:
|
||||
env: { REDIS_FAKE: "1" }
|
||||
run: python scripts/test-max-agents-per-user.py
|
||||
|
||||
- name: Swarm agent launcher (plan/limit/key-resolve) (#16)
|
||||
run: python scripts/test-agent-launcher.py
|
||||
|
||||
- name: Manager event-contract test
|
||||
env: { REDIS_FAKE: "1" }
|
||||
run: python scripts/test-contract-events.py
|
||||
|
||||
@@ -58,30 +58,34 @@ create 响应 `data`：`deployment_id`、`runtime_deployment_id`、`manager_depl
|
||||
|
||||
### 3.3 蜂群专家 agent 拉起环境契约（agent launch env，回应 agent_swarm#16）
|
||||
|
||||
> **冻结状态**：下表的 **env 字段清单本身已稳定**（`agent/main.py` 实读，无异议）；但**「拉起触发链」（谁调 AM、何时拉起、agent 池由谁提供）仍为提案，待 @azgy(AM) + @zsbgnw12(HM) 确认**（见 #16），**未冻结**。在触发链敲定前，本节按「env 契约已定、触发流程待定」对待。
|
||||
> **状态（团队决议）**：**Swarm 运行时负责拉起专家 agent 并执行每用户限额**（不再由 AM 拉起）。下表 env 字段 + 拉起方/限额口径**已定**。实现：`orchestrator/agent_launcher.py` + `main.launch_swarm_agents`，测试 `scripts/test-agent-launcher.py`。
|
||||
|
||||
**关键拓扑事实**：编排器（swarm runtime）**不拉起 agent、无 AM 客户端、不持有模型 key**。专家 agent 是**外部进程主动出站**连编排器 WS（`/ws/{agent_id}`，见 §5/security-boundary §6）。注意 §1「分解→派发→执行」里的「派发」= 把任务指派给**已连入**的 agent，**不是**编排器拉起 agent pod——§1 与本节不矛盾，但二者都未定义「谁拉起 pod」（即 #16 的缺口）。
|
||||
**拓扑**：专家 agent 仍是**独立进程**，主动出站连编排器 WS（`/ws/{agent_id}`，见 §5/security-boundary §6）。**拉起方 = Swarm 运行时**：编排器在 create 播种后由 `agent_launcher` 按 provisioning 策略拉起一个能力多样的 agent 池。去中心化下**无「按 run 分解算出 agent 数/角色」这一步**（种子任务无 required caps、任意 agent 可认领，子任务按能力自路由），故 agent 池**按用户**拉起、为固定能力集。§1 的「派发」= 把任务指派给**已连入**的 agent；**拉起 agent 进程**是本节定义的 Swarm 新职责。
|
||||
|
||||
**拉起触发链（提案，待确认，见 #16）**：去中心化模型下编排器只播种单一种子任务、agent 自选+自主分解，**无「按 run 分解算出 agent 数/角色」这一步**；故 agent 是**按用户的常驻池**（能力多样），**由 HM 触发 AM 拉起**（HM 持 per-user `sk-` + 已有 AM 客户端；若 AM 无批量接口则 HM 循环调模板 `POST /agents`）。模型 key 全程 **HM→AM**，**不经编排器 create 请求体**（满足 §3.1）。`ORCHESTRATOR_URL` 为部署期常量，由 HM 注入。
|
||||
**拉起后端**（`AGENT_LAUNCH_BACKEND`，fail-soft——拉起失败不影响 create）：
|
||||
- `none`（默认）：不自动拉起，agent 由外部供给（保留 CI/e2e 与「外部拉起」部署）；
|
||||
- `subprocess`：本地起 `python -m agent.main`（dev）；
|
||||
- `command`：执行部署注入的模板 `AGENT_LAUNCH_CMD`（env 经进程环境传入，模板包装真实拉起如 kubectl/pod-create；**密钥不上 argv**）。
|
||||
|
||||
AM 拉起每个蜂群专家 agent 时，**须注入以下进程环境变量**（本仓 `agent/main.py` 实读，为接线清单——此部分无异议）：
|
||||
**限额（Swarm 执行）**：池大小 `AGENT_LAUNCH_POOL_SIZE`（默认 3）；**实际拉起数 = min(池大小, `MAX_AGENTS_PER_USER` − 该用户已连接数)**——Swarm 在**拉起时**限额，并在 **agent 注册时**兜底硬拒（security-boundary §6 / PR#32）。
|
||||
|
||||
Swarm 拉起每个专家 agent 时注入以下进程环境变量（`agent/main.py` 实读）：
|
||||
|
||||
| env | 含义 | 来源 |
|
||||
|---|---|---|
|
||||
| `OPENAI_API_KEY` | 模型调用凭据 = **HM 为该用户现签的 per-user `sk-`**（NewAPI token），扣发起用户 `user.Quota` | HM 现签 |
|
||||
| `OPENAI_API_BASE` | **HM 模型网关 `/v1`**（OpenAI 兼容） | HM |
|
||||
| `OPENAI_MODEL` | 所选模型 id | HM/计划 |
|
||||
| `ORCHESTRATOR_URL` | 编排器 WS 基址（如 `ws://<swarm-runtime>`）；agent 据此回连、自选任务 | 部署 |
|
||||
| `AGENT_ID` | agent 实例 id（唯一） | AM |
|
||||
| `AGENT_CAPABILITIES` | 能力集合（逗号分隔） | 计划/角色 |
|
||||
| `HEICODE_USER_ID` | 发起用户；用于**每用户并发 Agent 上限** `MAX_AGENTS_PER_USER`（注册时强制，见 security-boundary §6） | HM |
|
||||
| `OPENAI_API_KEY` | 模型调用凭据 = 该用户 per-user `sk-`（NewAPI token），扣发起用户 `user.Quota` | **Swarm 从 `billing_context.secret_ref` 解析**（HM 现签写入 KV） |
|
||||
| `OPENAI_API_BASE` | HM 模型网关 `/v1`（OpenAI 兼容） | 部署（`AGENT_OPENAI_API_BASE`/`OPENAI_API_BASE`） |
|
||||
| `OPENAI_MODEL` | 所选模型 id | `billing_context.default_model_id`/计划 |
|
||||
| `ORCHESTRATOR_URL` | 编排器 WS 基址；agent 据此回连、自选任务 | 部署（`ORCHESTRATOR_PUBLIC_URL`/`AGENT_RUNTIME_WS_URL`） |
|
||||
| `AGENT_ID` | agent 实例 id（唯一） | **Swarm launcher** |
|
||||
| `AGENT_CAPABILITIES` | 能力集合（逗号分隔） | **Swarm launcher**（`AGENT_LAUNCH_CAPABILITIES` 池策略） |
|
||||
| `HEICODE_USER_ID` | 发起用户；用于每用户并发上限 `MAX_AGENTS_PER_USER`（注册时强制） | **Swarm 从 create `x-user-id` 透传** |
|
||||
| `WORKSPACE_DIR` / `GIT_REPO_URL` | 可选：工作区 / 代码仓 | 资源授权 |
|
||||
|
||||
**约束**：
|
||||
- key 注入是 **server→server env**，**不入** create 请求 / 回调 / 日志 / 事件（与模板 Agent 一致）；编排器 create 的 `billing_context` 仅为归因元数据（`newapi_user_ref`/`quota_ref`），**不承载 key**。见 [usage-billing-schema.md §2](./usage-billing-schema.md)。
|
||||
- token 生命周期 = **agent 部署生命周期**（stop/delete 时由 HM 吊销），**非 run 结束**——专家 agent 长驻、跨 run 自选任务。
|
||||
- 所有专家 agent 注入**同一把**该用户的 `sk-` + `OPENAI_API_BASE=HM/v1`，保证计费归一到发起用户；`task_executor` 随模型请求带 `X-Agent-*` 归因头供 HM/NewAPI 关联。
|
||||
- **运行时事件前置**：只有 (i) HM 真把 create 派发到编排器（`SWARM_RUNTIME_ENABLED=true`，非 manager-local 适配器）且 (ii) 专家 agent 已按上表拉起连入，编排器才会回推 `task.*`/`swarm.*` 等运行时事件（否则 events feed 仅有 HM 控制面的 `deployment.status_changed`）。
|
||||
- 模型 key 由 **Swarm 服务端从 `billing_context.secret_ref`（`azkv://`）解析**（`resolve_model_key`：override → azkv（部署 SecretResolver / dev `HEICODE_SECRET_<name>`）→ 编排器 `OPENAI_API_KEY` 兜底）后注入被拉起 agent 的 env。key **不入** create 请求体 / 回调 / 日志 / argv（满足 §3.1）；解析不到则 agent keyless 启动并明确报错（不伪造）。`billing_context` 的 `newapi_user_ref`/`quota_ref` 仍仅为归因元数据。见 [usage-billing-schema.md §2](./usage-billing-schema.md)。
|
||||
- 同一用户的专家 agent 注入**同一把** `sk-` + `OPENAI_API_BASE=HM/v1`，计费归一到发起用户；`task_executor` 随模型请求带 `X-Agent-*` 归因头供 HM/NewAPI 关联。token 由 HM 按 agent 生命周期吊销（stop/delete）。
|
||||
- **运行时事件前置**：只有 (i) HM 真把 create 派发到编排器（`SWARM_RUNTIME_ENABLED=true`，非 manager-local 适配器）且 (ii) Swarm 已按上表拉起 agent 并连入，编排器才回推 `task.*`/`swarm.*` 运行时事件（否则 events feed 仅有 HM 控制面 `deployment.status_changed`，见 #39）。
|
||||
|
||||
## 4. 状态机
|
||||
|
||||
|
||||
@@ -59,6 +59,7 @@
|
||||
- **Swarm 模型**：Agent 主动出站连编排器 WebSocket（`/ws/{agent_id}`），**不**对公网暴露每 Agent 子域名。AM 单 Agent 模型里的「客户端↔agent 直连 + `AGENT_ACCESS_TOKEN` 本地校验」**不适用于** swarm（无直连回路）。
|
||||
- 服务间鉴权：HM→Swarm 用 `AGENT_RUNTIME_SERVICE_TOKEN`（Bearer）；回调 HMAC 签名。
|
||||
- **每用户并发 Agent 配额**：一个 `user_id` 同时连接的 Agent 数上限为 `MAX_AGENTS_PER_USER`（env，默认 10）。注册（WS `register` 消息携带 `user_id`）超额即被拒绝（回 `registration_rejected` 并关闭，code 1008），断开后释放名额。归因主轴仍为 `user.id`/`channelId`。未带 `user_id` 的 Agent 为 unbound，不计入该配额。实现：`ConnectionManager.can_bind_user/bind_user/unbind` + 注册处强制；测试 `scripts/test-max-agents-per-user.py`。
|
||||
- **Swarm 拉起 agent + 服务端解析 key（team 决议，runtime-contract §3.3）**：由 **Swarm 运行时**（`orchestrator/agent_launcher.py`）拉起专家 agent 池（拉起数 `min(池大小, MAX_AGENTS_PER_USER − 已连)`，与上面的注册兜底一致）。模型 key 由 **Swarm 从 `billing_context.secret_ref`（`azkv://`）服务端解析**后注入被拉起 agent 的 env——**不入** create 请求体 / 回调 / 日志 / argv（`command` 后端的密钥经进程 env 传入，不上命令行）。azkv 真实解析为部署侧 SecretResolver；dev/CI 用 `HEICODE_SECRET_<name>`。解析不到即 keyless 启动并明确报错（不伪造）。
|
||||
- 🟡 待接入：多租户运行时隔离（命名空间/网络/配额）由 Agent 平台（AKS Workload Identity）承载，非本仓编排器；归因主轴为 `user.id`/`channelId`（见 `usage-billing-schema.md`），不引入 tenant 概念。
|
||||
|
||||
## 7. 外部 API 与传输
|
||||
|
||||
@@ -0,0 +1,234 @@
|
||||
"""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).
|
||||
|
||||
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 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
|
||||
|
||||
|
||||
def _resolve_secret_ref(secret_ref: str) -> Optional[str]:
|
||||
"""Resolve an azkv:// ref to its secret value.
|
||||
|
||||
Production: an Azure Key Vault adapter (out of this repo; wire via deployment). Dev/CI: read
|
||||
env `HEICODE_SECRET_<name>` where <name> is the trailing path segment of the ref. Returns None
|
||||
when unavailable — we never fabricate a key.
|
||||
"""
|
||||
name = secret_ref.rstrip("/").rsplit("/", 1)[-1]
|
||||
return os.getenv(f"HEICODE_SECRET_{name}")
|
||||
|
||||
|
||||
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 []
|
||||
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
|
||||
|
||||
|
||||
async def stop_launched(swarm_id: str) -> int:
|
||||
"""Best-effort stop of subprocess-launched agents for a run (other backends: deployment-managed)."""
|
||||
procs = _subprocess_agents.pop(swarm_id, [])
|
||||
stopped = 0
|
||||
for proc in procs:
|
||||
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)
|
||||
return stopped
|
||||
@@ -37,6 +37,7 @@ from . import autonomous_tasks as autonomous_mod
|
||||
from . import task_competition as competition_mod
|
||||
from . import cross_review as cross_review_mod
|
||||
from . import guard as guard_mod
|
||||
from . import agent_launcher
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
@@ -1599,6 +1600,36 @@ async def finalize_parent_after_child(run, child_task, agent_id: str, success: b
|
||||
)
|
||||
|
||||
|
||||
async def launch_swarm_agents(run, body: Dict[str, Any]) -> None:
|
||||
"""Swarm-owned agent launch (agent_swarm#16): launch the per-user expert pool for this run,
|
||||
capped at MAX_AGENTS_PER_USER, with the model key resolved server-side from
|
||||
billing_context.secret_ref. No-op unless AGENT_LAUNCH_BACKEND is set (default 'none'). Fail-soft.
|
||||
|
||||
Manager-provided explicit agent breakdowns are honored as-is (those agents are provisioned by
|
||||
the caller), so we only auto-launch the pool for the decentralized seed flow.
|
||||
"""
|
||||
if _manager_provided_agents(body):
|
||||
return
|
||||
user_id = ((body.get("metadata") or {}).get("runtime_headers") or {}).get("x_user_id")
|
||||
connected = manager.user_agent_count(user_id) if user_id else 0
|
||||
try:
|
||||
specs = agent_launcher.plan_launch_specs(
|
||||
run, body,
|
||||
connected_user_agents=connected,
|
||||
limit=max_agents_per_user(),
|
||||
pool_size=agent_launcher.desired_pool_size(),
|
||||
model_key=agent_launcher.resolve_model_key(body),
|
||||
orchestrator_url=agent_launcher.orchestrator_ws_url(),
|
||||
user_id=user_id,
|
||||
)
|
||||
launched = await agent_launcher.launch(specs, swarm_id=run.swarm_id)
|
||||
if launched:
|
||||
run.metadata["launched_agents"] = launched
|
||||
await swarm_runtime.save_run(run)
|
||||
except Exception as exc: # never fail run creation on launch
|
||||
logger.warning("launch_swarm_agents failed for run %s: %s", run.swarm_id, exc)
|
||||
|
||||
|
||||
async def create_swarm_run_from_request(
|
||||
request: Request,
|
||||
body: Dict[str, Any],
|
||||
@@ -1628,6 +1659,7 @@ async def create_swarm_run_from_request(
|
||||
|
||||
if created and run.status == "running":
|
||||
await create_tasks_for_run(run, body)
|
||||
await launch_swarm_agents(run, body)
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
@@ -1663,6 +1695,12 @@ async def stop_swarm_run(deployment_or_swarm_id: str, reason: str):
|
||||
"summary": "Swarm stopped by Manager",
|
||||
"reason": reason,
|
||||
})
|
||||
# Tear down any Swarm-launched agents for this run (best-effort; subprocess backend only —
|
||||
# command/k8s backends are torn down by the deployment). agent_swarm#16.
|
||||
try:
|
||||
await agent_launcher.stop_launched(run.swarm_id)
|
||||
except Exception as exc:
|
||||
logger.warning("stop_launched failed for run %s: %s", run.swarm_id, exc)
|
||||
return run
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
"""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 main():
|
||||
import asyncio
|
||||
test_launch_count()
|
||||
test_plan_specs()
|
||||
test_resolve_model_key()
|
||||
test_command_backend_build()
|
||||
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()
|
||||
Reference in New Issue
Block a user