Merge pull request #69 from xmindlab-heicode/fix/orphan-agent-lifecycle
fix(swarm): 孤儿 agent 终态回收 + 重试退避 + 任务错误上浮
This commit is contained in:
@@ -164,6 +164,46 @@ class GitOperations:
|
||||
logger.error(f"Error creating result branch: {e}")
|
||||
return False
|
||||
|
||||
async def add_task_worktree(self, path: str, task_id: str) -> Optional[str]:
|
||||
"""Create an isolated git worktree at `path` on a fresh per-task result branch.
|
||||
|
||||
Cut from this agent's clone (off origin/<base> or HEAD), the worktree has the full repo
|
||||
contents plus its own index/HEAD, so one agent can run several tasks concurrently without
|
||||
them clobbering a shared checkout. Returns the branch name, or None on failure.
|
||||
"""
|
||||
try:
|
||||
root = await self.repo_root() or self.workspace_dir
|
||||
safe_task_id = task_id.replace("/", "-")[:12]
|
||||
branch = f"agent/{self.agent_id}/{safe_task_id}-{int(time.time())}"
|
||||
|
||||
await self._run_git_command(["git", "fetch", "origin"], cwd=root)
|
||||
base_ref = f"origin/{self.base_branch}"
|
||||
if (await self._run_git_command(["git", "rev-parse", "--verify", base_ref], cwd=root)).returncode != 0:
|
||||
base_ref = "HEAD"
|
||||
logger.warning(f"Base branch origin/{self.base_branch} not found; worktree from HEAD")
|
||||
|
||||
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
||||
result = await self._run_git_command(
|
||||
["git", "worktree", "add", "-b", branch, path, base_ref], cwd=root,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
logger.error(f"git worktree add failed: {result.stderr}")
|
||||
return None
|
||||
logger.info(f"Created task worktree {path} on branch {branch}")
|
||||
return branch
|
||||
except Exception as e:
|
||||
logger.error(f"Error creating task worktree: {e}")
|
||||
return None
|
||||
|
||||
async def remove_task_worktree(self, path: str) -> None:
|
||||
"""Tear down a per-task worktree (best-effort). The branch ref is kept (already pushed)."""
|
||||
try:
|
||||
root = await self.repo_root() or self.workspace_dir
|
||||
await self._run_git_command(["git", "worktree", "remove", "--force", path], cwd=root)
|
||||
logger.info(f"Removed task worktree {path}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to remove task worktree {path}: {e}")
|
||||
|
||||
async def commit_changes(self, message: str) -> Optional[str]:
|
||||
try:
|
||||
cwd = await self.repo_root() or self.workspace_dir
|
||||
|
||||
+36
-14
@@ -72,6 +72,9 @@ class AgentRuntimeDisconnected(RuntimeError):
|
||||
class Agent:
|
||||
"""Agent that connects to orchestrator and executes tasks."""
|
||||
|
||||
# An agent may run several tasks at once; each executes in its OWN git worktree cut from this
|
||||
# agent's clone (see _execute_assignment / GitOperations.add_task_worktree), so concurrent tasks
|
||||
# never share a checkout/index. Swarm parallelism = many agents × this per-agent concurrency.
|
||||
MAX_CONCURRENT_TASKS = int(os.getenv("MAX_CONCURRENT_TASKS", "4"))
|
||||
TASK_TIMEOUT_SECONDS = int(os.getenv("TASK_TIMEOUT_SECONDS", "60"))
|
||||
PEER_REPLY_TIMEOUT_SECONDS = int(os.getenv("PEER_REPLY_TIMEOUT_SECONDS", "20"))
|
||||
@@ -111,14 +114,20 @@ class Agent:
|
||||
self._peer_executor: Optional[TaskExecutor] = None
|
||||
|
||||
self.workspace_git = GitOperations(str(self.workspace_dir), self.agent_id)
|
||||
# Serializes only the brief shared-repo git plumbing (worktree add/remove + fetch) so
|
||||
# concurrent tasks don't race on .git locks. Task EXECUTION stays parallel.
|
||||
self._git_admin_lock = asyncio.Lock()
|
||||
|
||||
def available_slots(self) -> int:
|
||||
# Return how many additional tasks this agent can currently accept.
|
||||
return max(0, self.MAX_CONCURRENT_TASKS - len(self.active_tasks))
|
||||
|
||||
def task_workspace(self, task_id: str) -> Path:
|
||||
# Compute the dedicated per-task working directory inside the agent workspace.
|
||||
return self.workspace_dir / ".agent_tasks" / task_id
|
||||
# Per-task git worktree path. Kept OUTSIDE the repo root (default /tmp/agent-worktrees) so
|
||||
# the main checkout never sees it as untracked, and namespaced by agent_id to avoid
|
||||
# collisions. `git worktree add` requires the leaf to not pre-exist, so we don't mkdir it.
|
||||
base = Path(os.getenv("AGENT_WORKTREE_BASE", "/tmp/agent-worktrees"))
|
||||
return base / self.agent_id / task_id.replace("/", "-")
|
||||
|
||||
async def safe_send(self, payload: dict):
|
||||
# Serialize and send a websocket message while holding a lock to prevent concurrent writes.
|
||||
@@ -383,17 +392,27 @@ class Agent:
|
||||
|
||||
start_time = time.time()
|
||||
task_workspace = self.task_workspace(task_id)
|
||||
task_workspace.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
git_enabled = False
|
||||
task_git = None
|
||||
try:
|
||||
executor = TaskExecutor(agent_id=self.agent_id, workspace_dir=str(task_workspace))
|
||||
git_enabled = await self.workspace_git.is_git_workspace()
|
||||
if git_enabled:
|
||||
branch_created = await self.workspace_git.create_result_branch(task_id)
|
||||
if not branch_created:
|
||||
logger.warning("Failed to create task result branch; continuing without git push")
|
||||
git_enabled = False
|
||||
# Each task runs in its OWN git worktree cut from this agent's clone, so an agent can
|
||||
# run several tasks concurrently without sharing a checkout/index. The worktree holds
|
||||
# the full repo contents on a fresh result branch; the executor reads/writes the real
|
||||
# source there and the per-task GitOperations (task_git) commits/pushes that branch.
|
||||
# This fixes the prior empty per-task subdir → "Empty workspace: no files detected".
|
||||
# Falls back to the repo root (no isolated branch) only if worktree creation fails.
|
||||
if await self.workspace_git.is_git_workspace():
|
||||
async with self._git_admin_lock: # serialize shared-repo git plumbing only
|
||||
result_branch = await self.workspace_git.add_task_worktree(str(task_workspace), task_id)
|
||||
if result_branch:
|
||||
git_enabled = True
|
||||
task_git = GitOperations(str(task_workspace), self.agent_id)
|
||||
task_git.result_branch = result_branch
|
||||
else:
|
||||
logger.warning("Failed to create task worktree; executing on repo root without git push")
|
||||
exec_dir = str(task_workspace) if git_enabled else str(self.workspace_dir)
|
||||
executor = TaskExecutor(agent_id=self.agent_id, workspace_dir=exec_dir)
|
||||
|
||||
result = await asyncio.wait_for(
|
||||
executor.execute_task(
|
||||
@@ -401,7 +420,7 @@ class Agent:
|
||||
description=description,
|
||||
context={
|
||||
**context,
|
||||
"workspace_dir": str(task_workspace),
|
||||
"workspace_dir": exec_dir,
|
||||
"repo_workspace_dir": str(self.workspace_dir),
|
||||
"git_repo_url": self.git_repo_url,
|
||||
"agent_id": self.agent_id,
|
||||
@@ -417,12 +436,12 @@ class Agent:
|
||||
if result.get("success"):
|
||||
self.last_summary = self._summarize_result(result) or self.last_summary
|
||||
if result.get("success") and not awaiting_handoff:
|
||||
if git_enabled:
|
||||
commit_sha = await self.workspace_git.commit_changes(
|
||||
if git_enabled and task_git:
|
||||
commit_sha = await task_git.commit_changes(
|
||||
message=f"Task {task_id}: {description[:50]}"
|
||||
)
|
||||
if commit_sha:
|
||||
branch_name = await self.workspace_git.push_results()
|
||||
branch_name = await task_git.push_results()
|
||||
result["git_branch"] = branch_name
|
||||
result["commit_sha"] = commit_sha
|
||||
else:
|
||||
@@ -465,6 +484,9 @@ class Agent:
|
||||
await self.send_task_result(task_id, False, {"error": str(e), "success": False})
|
||||
await self.send_status_update("idle", None, f"Task failed: {e}")
|
||||
finally:
|
||||
if git_enabled:
|
||||
async with self._git_admin_lock:
|
||||
await self.workspace_git.remove_task_worktree(str(task_workspace))
|
||||
self.active_tasks.pop(task_id, None)
|
||||
ACTIVE_TASKS.set(len(self.active_tasks))
|
||||
self.current_task_id = None
|
||||
|
||||
+39
-2
@@ -117,7 +117,7 @@ class TaskExecutor:
|
||||
|
||||
success = all(r["status"] in ["completed", "handed_off"] for r in results)
|
||||
awaiting_handoff = any(r["status"] == "handed_off" for r in results)
|
||||
return {
|
||||
payload = {
|
||||
"success": success,
|
||||
"task_id": task_id,
|
||||
"subtasks": results,
|
||||
@@ -125,6 +125,26 @@ class TaskExecutor:
|
||||
"agent_id": self.agent_id,
|
||||
"usage": self._usage_payload(time.time() - started_at),
|
||||
}
|
||||
if not success:
|
||||
# Surface the model's OWN failure explanation (what it saw / what was missing) as a
|
||||
# top-level `error`, so the orchestrator/Manager records WHY instead of a generic
|
||||
# "Task failed" (agent/main.py uses this as the failure reason). Lead with the model's
|
||||
# summary/error — NOT the task prompt — so the reason reads as the diagnosis, not the
|
||||
# ask. summary is usually the fuller "saw X, missing Y" narrative; append a distinct
|
||||
# error for the concise cause.
|
||||
failed = [r for r in results if r.get("status") not in ("completed", "handed_off")]
|
||||
explanations = []
|
||||
for r in failed:
|
||||
summary = str(r.get("summary") or "").strip()
|
||||
error = str(r.get("error") or "").strip()
|
||||
text = summary or error or "subtask failed without detail"
|
||||
if summary and error and error not in summary:
|
||||
text = f"{summary} ({error})"
|
||||
explanations.append(text)
|
||||
detail = "; ".join(explanations) or "subtask(s) failed without detail"
|
||||
payload["error"] = detail
|
||||
logger.error(f"Task {task_id} failed: {detail}")
|
||||
return payload
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing task {task_id}: {e}")
|
||||
return {
|
||||
@@ -166,6 +186,7 @@ Return ONLY the JSON array, no other text."""
|
||||
}]
|
||||
|
||||
async def _execute_subtask(self, subtask: dict, task_id: str, context: dict, peer_collaboration_callback: Optional[Callable]) -> dict:
|
||||
content = ""
|
||||
try:
|
||||
description = subtask["description"]
|
||||
workspace_files = self._summarize_workspace()
|
||||
@@ -270,10 +291,26 @@ Return ONLY the JSON, no other text."""
|
||||
if apply_result["errors"]:
|
||||
result["status"] = "failed"
|
||||
result["error"] = "; ".join(apply_result["errors"])
|
||||
logger.warning(
|
||||
f"Subtask for task {task_id} failed applying file changes: {result['error']}"
|
||||
)
|
||||
else:
|
||||
# LLM returned 2xx but declared the subtask not completed: log its reason and a
|
||||
# bounded snippet of the model output so the failure is diagnosable from agent logs.
|
||||
logger.warning(
|
||||
"Subtask for task %s reported status=%r (error=%s summary=%s); llm_response[:500]=%s",
|
||||
task_id, result.get("status"), result.get("error"),
|
||||
result.get("summary"), (content or "").strip()[:500],
|
||||
)
|
||||
result["subtask"] = subtask
|
||||
return result
|
||||
except Exception as e:
|
||||
logger.error(f"Error executing subtask: {e}")
|
||||
# Parse failures / API errors land here; include a bounded model-output snippet
|
||||
# (model-generated text, no secrets) to explain why parsing/execution failed.
|
||||
logger.error(
|
||||
f"Error executing subtask for task {task_id}: {e}; "
|
||||
f"llm_response[:500]={(content or '').strip()[:500]}"
|
||||
)
|
||||
return {
|
||||
"subtask": subtask,
|
||||
"status": "failed",
|
||||
|
||||
@@ -61,8 +61,10 @@ create 响应 `data`:`deployment_id`、`runtime_deployment_id`、`manager_depl
|
||||
### 3.3 蜂群专家 agent 拉起环境契约(agent launch env,回应 agent_swarm#16)
|
||||
|
||||
> **状态(团队决议)**:**Swarm 运行时负责拉起专家 agent 并执行每用户限额**(不再由 AM 拉起)。下表 env 字段 + 拉起方/限额口径**已定**。实现:`orchestrator/agent_launcher.py` + `main.launch_swarm_agents`,测试 `scripts/test-agent-launcher.py`。
|
||||
>
|
||||
> **架构师裁定(2026-06-15,取代此前逃生门口径)**:**所有蜂群一律走去中心化主路——Swarm 永远播种单一目标任务(`build_seed_task_specs`),并按用户拉起 agent 池。** `orchestration_plan.agents`(HM 现会发 `[{role:general}]`)**不控制任务创建,也不控制 agent 拉起**:既不会让运行时跳过播种、改按 breakdown 建任务,也不会让运行时认为「caller 自行拉起 agent」而不拉池。此前「Manager 显式 agent breakdown 由 caller 拉起 / honored as-is」的逃生门(`_manager_provided_agents`)已**彻底退役删除**,本节其余口径(拉起后端、限额、模型 key 注入 §3.3.1)不变。这与 CLAUDE.md「去中心化蜂群是唯一行为(cutover 已完成)」一致。
|
||||
|
||||
**拓扑**:专家 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 新职责。
|
||||
**拓扑**:专家 agent 仍是**独立进程**,主动出站连编排器 WS(`/ws/{agent_id}`,见 §5/security-boundary §6)。**拉起方 = Swarm 运行时**:编排器在 create **无条件播种**单一目标任务后,由 `agent_launcher` 按 provisioning 策略拉起一个能力多样的 agent 池。去中心化下**无「按 run 分解算出 agent 数/角色」这一步**,且**无「按 caller 提供的 agent breakdown 拉起」这一步**(种子任务无 required caps、任意 agent 可认领,子任务按能力自路由),故 agent 池**按用户**拉起、为固定能力集。§1 的「派发」= 把任务指派给**已连入**的 agent;**拉起 agent 进程**是本节定义的 Swarm 新职责。
|
||||
|
||||
**拉起后端**(`AGENT_LAUNCH_BACKEND`,fail-soft——拉起失败不影响 create):
|
||||
- `none`(默认):不自动拉起,agent 由外部供给(保留 CI/e2e 与「外部拉起」部署);
|
||||
|
||||
+113
-12
@@ -61,14 +61,35 @@ 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)."""
|
||||
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_POOL_SIZE", "3") or 3))
|
||||
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")
|
||||
@@ -322,6 +343,47 @@ def model_id(body: Dict[str, Any]) -> Optional[str]:
|
||||
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
|
||||
@@ -330,8 +392,23 @@ class AgentLaunchSpec:
|
||||
|
||||
|
||||
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)))
|
||||
"""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(
|
||||
@@ -348,10 +425,10 @@ def plan_launch_specs(
|
||||
) -> List[AgentLaunchSpec]:
|
||||
"""Pure: plan the agent pool for a run (count capped by the per-user limit) + each agent's env.
|
||||
|
||||
The model key and git binding are resolved server-side (see `resolve_model_key` /
|
||||
`resolve_git_grant`) and injected into the launch env, never 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).
|
||||
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.
|
||||
"""
|
||||
count = launch_count(pool_size=pool_size, limit=limit, connected_user_agents=connected_user_agents)
|
||||
caps = pool_capabilities() or ["general"]
|
||||
@@ -373,7 +450,7 @@ def plan_launch_specs(
|
||||
if user_id:
|
||||
env["HEICODE_USER_ID"] = user_id
|
||||
if git_env:
|
||||
env.update(git_env) # GIT_REPO_URL (+ GIT_USERNAME/GIT_PASSWORD/GIT_BASE_BRANCH)
|
||||
env.update(git_env)
|
||||
specs.append(AgentLaunchSpec(agent_id=env["AGENT_ID"], capabilities=cap_csv, env=env))
|
||||
return specs
|
||||
|
||||
@@ -478,6 +555,20 @@ def pod_resources() -> Dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
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")
|
||||
@@ -524,9 +615,15 @@ def build_pod_manifest(spec: AgentLaunchSpec, *, namespace: str, swarm_id: str,
|
||||
}
|
||||
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": {"name": spec.agent_id, "namespace": namespace, "labels": _k8s_labels(spec, swarm_id)},
|
||||
"metadata": metadata,
|
||||
"spec": pod_spec,
|
||||
}
|
||||
|
||||
@@ -590,7 +687,11 @@ async def stop_launched(swarm_id: str) -> int:
|
||||
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:
|
||||
# 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(
|
||||
|
||||
+99
-32
@@ -145,14 +145,73 @@ manager = ConnectionManager()
|
||||
AGENT_SLOTS: Dict[str, int] = {}
|
||||
|
||||
|
||||
def max_agents_per_user() -> int:
|
||||
"""Max concurrent agents one user may have connected (per-user swarm cap). Env-tunable; default 10."""
|
||||
# Absolute ceiling on concurrent agents per user, enforced regardless of what HM/plan/env requests.
|
||||
# Product rule: a tenant may run at most 16 agents at once (architect ruling 2026-06-15).
|
||||
MAX_AGENTS_PER_USER_HARD_CAP = 16
|
||||
|
||||
|
||||
def _clamp_user_cap(value: int) -> int:
|
||||
"""Clamp a requested per-user agent cap into [1, MAX_AGENTS_PER_USER_HARD_CAP]."""
|
||||
return max(1, min(MAX_AGENTS_PER_USER_HARD_CAP, value))
|
||||
|
||||
|
||||
def _default_task_retries() -> int:
|
||||
"""Default retry budget for runtime tasks (env TASK_MAX_RETRIES, default 3). Paired with the
|
||||
task-queue retry backoff so retries span minutes rather than seconds."""
|
||||
try:
|
||||
return max(1, int(os.getenv("MAX_AGENTS_PER_USER", "10") or 10))
|
||||
return max(1, int(os.getenv("TASK_MAX_RETRIES", "3") or 3))
|
||||
except ValueError:
|
||||
return 3
|
||||
|
||||
|
||||
def _env_max_agents_per_user() -> int:
|
||||
try:
|
||||
return _clamp_user_cap(int(os.getenv("MAX_AGENTS_PER_USER", "10") or 10))
|
||||
except ValueError:
|
||||
return 10
|
||||
|
||||
|
||||
def max_agents_per_user(body: Optional[Dict[str, Any]] = None) -> int:
|
||||
"""Max concurrent agents one user may have connected (per-user swarm cap / ceiling).
|
||||
|
||||
HM sends `metadata.max_agents_per_user = min(plan SwarmMaxAgents, user override)` on the swarm
|
||||
create body, and only when > 0. When `body` carries a positive integer there, it WINS (this is
|
||||
the per-user CEILING for that run); otherwise we fall back to env `MAX_AGENTS_PER_USER` (default
|
||||
10). `body` is OPTIONAL: callers without a create body (e.g. the WS register-time cap on an agent
|
||||
whose owning run can't be resolved) still get the env value — DO NOT make body required.
|
||||
"""
|
||||
if isinstance(body, dict):
|
||||
cap = ((body.get("metadata") or {}).get("max_agents_per_user"))
|
||||
if isinstance(cap, int) and not isinstance(cap, bool) and cap > 0:
|
||||
# HM's value wins, but never above the hard product ceiling (16).
|
||||
return _clamp_user_cap(cap)
|
||||
return _env_max_agents_per_user()
|
||||
|
||||
|
||||
async def _per_user_cap_for_agent(agent_id: str) -> int:
|
||||
"""Per-user cap to enforce at WS registration for `agent_id`.
|
||||
|
||||
Keeps the register-time cap consistent with what launch_swarm_agents planned: launcher-minted
|
||||
agent ids are `{swarm_id}-agent-{N}`, so we strip the `-agent-N` suffix, resolve the owning run,
|
||||
and reuse ITS metadata.max_agents_per_user ceiling. If the run can't be resolved (externally
|
||||
supplied agent, malformed id, or a run not yet/no longer persisted) we fall back to the env cap.
|
||||
Fail-soft: any lookup error falls back to env (never blocks registration on a lookup hiccup).
|
||||
"""
|
||||
try:
|
||||
marker = "-agent-"
|
||||
idx = agent_id.rfind(marker)
|
||||
if idx <= 0:
|
||||
return _env_max_agents_per_user()
|
||||
swarm_id = agent_id[:idx]
|
||||
run = await swarm_runtime.get_run_by_identifier(swarm_id)
|
||||
if run is None:
|
||||
return _env_max_agents_per_user()
|
||||
return max_agents_per_user({"metadata": run.metadata or {}})
|
||||
except Exception as exc: # never block registration on a cap-lookup error
|
||||
logger.debug("per-user cap lookup failed for agent %s: %s; using env cap", agent_id, exc)
|
||||
return _env_max_agents_per_user()
|
||||
|
||||
|
||||
def record_agent_slots(agent_id: str, message: dict):
|
||||
"""Update the cached free-slot count for an agent from a protocol message."""
|
||||
slots = message.get("available_slots")
|
||||
@@ -692,6 +751,17 @@ async def refresh_swarm_run_status(run):
|
||||
except Exception as exc:
|
||||
logger.warning("benchmark capture failed for run %s: %s", run.swarm_id, exc)
|
||||
|
||||
# Reap the run's agent pods/secret now that it has reached a terminal state. Previously teardown
|
||||
# ran only on a Manager-requested stop (stop_swarm_run), so naturally completed/failed runs left
|
||||
# their agents Running — those idle, registered agents then lingered in the shared pool and could
|
||||
# self-select another run's tasks (incident 2026-06-15: a leftover agent ran another swarm's seed
|
||||
# against the wrong workspace). Best-effort; never fails the terminal transition. Does NOT touch
|
||||
# task self-selection — only the worker pod lifecycle.
|
||||
try:
|
||||
await agent_launcher.stop_launched(run.swarm_id)
|
||||
except Exception as exc:
|
||||
logger.warning("agent teardown on terminal run %s failed: %s", run.swarm_id, exc)
|
||||
|
||||
|
||||
# Lifespan context manager
|
||||
@asynccontextmanager
|
||||
@@ -807,13 +877,6 @@ def request_context_headers(request: Request) -> Dict[str, Optional[str]]:
|
||||
}
|
||||
|
||||
|
||||
def _manager_provided_agents(body: Dict[str, Any]) -> bool:
|
||||
"""Whether the Manager supplied an explicit agent breakdown (Manager-first; never overridden)."""
|
||||
normalized = swarm_runtime.normalize_create_request(body)
|
||||
plan = normalized.get("orchestration_plan") or {}
|
||||
return bool(plan.get("agents") or normalized.get("agents"))
|
||||
|
||||
|
||||
def build_seed_task_specs(run, body: Dict[str, Any], base_specs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Decentralized rework P2: seed the run with ONE objective-carrying task instead of a
|
||||
Master-decomposed plan.
|
||||
@@ -1177,12 +1240,12 @@ async def create_tasks_for_run(run, body: Dict[str, Any]) -> int:
|
||||
"""Create the runtime task graph and emit task.created events."""
|
||||
created_count = 0
|
||||
task_specs = swarm_runtime.build_task_descriptions(body)
|
||||
# Swarm task creation (the ONLY path): when the Manager supplied an explicit agent breakdown we
|
||||
# honor it (Manager-first contract); otherwise we SEED a single objective-carrying task and the
|
||||
# Swarm task creation (the ONLY path): ALWAYS seed a single objective-carrying task and let the
|
||||
# agents grow the graph bottom-up via proposals (handle_task_proposal). No up-front Master
|
||||
# decomposition — this repo is the swarm runtime (see rework plan §0).
|
||||
if not _manager_provided_agents(body):
|
||||
task_specs = build_seed_task_specs(run, body, task_specs)
|
||||
# decomposition and no escape hatch for a caller-supplied agent breakdown — this repo is the
|
||||
# decentralized swarm runtime and seeding is unconditional. `orchestration_plan.agents` does NOT
|
||||
# control task creation (see rework plan §0 / runtime-contract §3.3, architect ruling 2026-06-15).
|
||||
task_specs = build_seed_task_specs(run, body, task_specs)
|
||||
task_id_map = {
|
||||
task_spec["task_id"]: f"{run.swarm_id}-{task_spec['task_id']}"
|
||||
for task_spec in task_specs
|
||||
@@ -1230,7 +1293,7 @@ async def create_tasks_for_run(run, body: Dict[str, Any]) -> int:
|
||||
root_task_id=runtime_root_task_id,
|
||||
source=task_spec.get("source", "runtime_bridge"),
|
||||
context=task_context,
|
||||
max_retries=body.get("max_retries", 3),
|
||||
max_retries=body.get("max_retries") or _default_task_retries(),
|
||||
)
|
||||
await swarm_runtime.attach_task(run, task.task_id)
|
||||
await swarm_runtime.emit_event(
|
||||
@@ -1663,8 +1726,10 @@ def _launch_note(backend: str, planned: int, launched: int, model_key_resolved:
|
||||
"agents externally, or set AGENT_LAUNCH_BACKEND=kubernetes (prod) / subprocess (dev). "
|
||||
"With no agent connected the seeded task is never claimed (see health.blockers).")
|
||||
if planned == 0:
|
||||
return ("0 agents planned — the per-user cap (MAX_AGENTS_PER_USER) is already met by connected "
|
||||
"agents, or AGENT_LAUNCH_POOL_SIZE is 0.")
|
||||
return ("0 agents planned — the launched-pod window is [AGENT_LAUNCH_MIN_POOL, "
|
||||
"AGENT_LAUNCH_MAX_POOL] (default [3,16]); a 0 here means the window was mis-set "
|
||||
"(MIN=0 with no headroom). Note the per-user cap = the run's "
|
||||
"metadata.max_agents_per_user ceiling (HM-sent) or env MAX_AGENTS_PER_USER.")
|
||||
if launched == 0:
|
||||
note = (f"backend={backend}: planned {planned} but launched 0 — the launch backend failed "
|
||||
f"(see orchestrator logs; for kubernetes verify kubectl/RBAC + Pod Workload Identity, "
|
||||
@@ -1685,21 +1750,19 @@ async def launch_swarm_agents(run, body: Dict[str, Any]) -> None:
|
||||
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.
|
||||
Every swarm goes through this unconditionally — the runtime always seeds + launches the pool.
|
||||
`orchestration_plan.agents` does NOT control launch; there is no caller-provisioned escape hatch
|
||||
(architect ruling 2026-06-15, runtime-contract §3.3). The launched-pod count self-regulates via
|
||||
plan_launch_specs / launch_count: max(MIN, min(MAX, min(pool_size, cap − already-connected))),
|
||||
with a HARD window [AGENT_LAUNCH_MIN_POOL, AGENT_LAUNCH_MAX_POOL] = [3, 16] by default (the MIN
|
||||
floor takes priority over cap headroom). `cap` = the run's metadata.max_agents_per_user ceiling
|
||||
(HM-sent, min(plan SwarmMaxAgents, user override)) when present, else env MAX_AGENTS_PER_USER.
|
||||
|
||||
Records the launch outcome on run.metadata["agent_launch"] (backend / planned / launched /
|
||||
model_key_resolved / note) — surfaced in /result + /diagnostics so a run with 0 expert agents
|
||||
explains itself instead of hanging silently (agent_swarm#56). The note carries no secret — only
|
||||
a bool for whether the model key resolved.
|
||||
"""
|
||||
if _manager_provided_agents(body):
|
||||
run.metadata["agent_launch"] = {
|
||||
"backend": "manager_provided", "planned": 0, "launched": 0, "launched_ids": [],
|
||||
"note": "Manager provided explicit agents; the runtime did not auto-launch a pool.",
|
||||
}
|
||||
await swarm_runtime.save_run(run)
|
||||
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
|
||||
backend = agent_launcher.launch_backend()
|
||||
@@ -1716,7 +1779,7 @@ async def launch_swarm_agents(run, body: Dict[str, Any]) -> None:
|
||||
specs = agent_launcher.plan_launch_specs(
|
||||
run, body,
|
||||
connected_user_agents=connected,
|
||||
limit=max_agents_per_user(),
|
||||
limit=max_agents_per_user(body),
|
||||
pool_size=agent_launcher.desired_pool_size(),
|
||||
model_key=model_key,
|
||||
orchestrator_url=agent_launcher.orchestrator_ws_url(),
|
||||
@@ -2686,11 +2749,15 @@ async def websocket_endpoint(websocket: WebSocket, agent_id: str):
|
||||
capabilities = data.get("capabilities", [])
|
||||
user_id = data.get("user_id")
|
||||
|
||||
# Per-user swarm cap: a user may have at most MAX_AGENTS_PER_USER (default 10) agents
|
||||
# connected at once. Agents that omit user_id are unbound and not capped. A reconnect by
|
||||
# an already-counted agent_id is allowed (can_bind_user handles it).
|
||||
# Per-user swarm cap: a user may have at most `limit` agents connected at once. The
|
||||
# limit MUST match what launch_swarm_agents planned with, or we'd "plan 12 but reject at
|
||||
# 10". The launcher uses the owning run's metadata.max_agents_per_user ceiling (HM-sent);
|
||||
# so here we resolve THAT run by agent_id and reuse its cap, falling back to env when the
|
||||
# run can't be found (externally-supplied agents, or pre-cap legacy ids).
|
||||
# Agents that omit user_id are unbound and not capped. A reconnect by an already-counted
|
||||
# agent_id is allowed (can_bind_user handles it).
|
||||
if user_id:
|
||||
limit = max_agents_per_user()
|
||||
limit = await _per_user_cap_for_agent(agent_id)
|
||||
if not manager.can_bind_user(agent_id, user_id, limit):
|
||||
logger.warning(
|
||||
"Rejecting agent %s: user %s already at max agents (%d connected)",
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Task queue management with dependency-aware failure recovery."""
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
import logging
|
||||
@@ -12,6 +13,24 @@ from .agent_registry import agent_registry, AgentStatus
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _retry_backoff_seconds(retry_count: int) -> float:
|
||||
"""Delay before a failed task becomes dispatchable again (exponential, capped).
|
||||
|
||||
Defaults: 5s → 30s → 180s … capped at 300s. Spreads retries over minutes instead of burning
|
||||
the whole budget in seconds, so a swarm's own agents (which may be cold-starting / waiting on
|
||||
node scale-up) have time to register before the seed exhausts its retries (incident 2026-06-15).
|
||||
Tunable via TASK_RETRY_BACKOFF_{BASE,FACTOR,CAP}.
|
||||
"""
|
||||
try:
|
||||
base = float(os.getenv("TASK_RETRY_BACKOFF_BASE", "5") or 5)
|
||||
factor = float(os.getenv("TASK_RETRY_BACKOFF_FACTOR", "6") or 6)
|
||||
cap = float(os.getenv("TASK_RETRY_BACKOFF_CAP", "300") or 300)
|
||||
except ValueError:
|
||||
base, factor, cap = 5.0, 6.0, 300.0
|
||||
exponent = max(0, retry_count - 1)
|
||||
return min(cap, base * (factor ** exponent))
|
||||
|
||||
|
||||
class TaskStatus(str, Enum):
|
||||
"""Task status enumeration."""
|
||||
PENDING = "pending"
|
||||
@@ -44,6 +63,7 @@ class Task(BaseModel):
|
||||
child_task_ids: List[str] = Field(default_factory=list)
|
||||
retry_count: int = 0
|
||||
max_retries: int = 3
|
||||
next_retry_at: Optional[float] = None
|
||||
context: Dict = Field(default_factory=dict)
|
||||
|
||||
|
||||
@@ -172,6 +192,11 @@ class TaskQueue:
|
||||
if task.status != TaskStatus.PENDING:
|
||||
return False
|
||||
|
||||
# Respect retry backoff: a task re-queued after a failure is not dispatchable until its
|
||||
# next_retry_at has passed (see fail_task / _retry_backoff_seconds).
|
||||
if task.next_retry_at and time.time() < task.next_retry_at:
|
||||
return False
|
||||
|
||||
for dependency_id in task.depends_on:
|
||||
dependency = await self.get_task(dependency_id)
|
||||
if not dependency or dependency.status != TaskStatus.COMPLETED:
|
||||
@@ -322,11 +347,17 @@ class TaskQueue:
|
||||
task.assigned_agent_id = None
|
||||
task.started_at = None
|
||||
|
||||
# Backoff: gate re-dispatch until next_retry_at (is_task_ready enforces it) so retries
|
||||
# spread over minutes rather than all firing within seconds.
|
||||
delay = _retry_backoff_seconds(task.retry_count)
|
||||
task.next_retry_at = time.time() + delay
|
||||
|
||||
# Re-add to pending queue
|
||||
await redis_client.lpush(self.PENDING_QUEUE_KEY, task_id)
|
||||
|
||||
logger.warning(
|
||||
f"Task {task_id} failed (retry {task.retry_count}/{task.max_retries}): {reason}"
|
||||
f"Task {task_id} failed (retry {task.retry_count}/{task.max_retries}): {reason} "
|
||||
f"(next retry in {delay:.0f}s)"
|
||||
)
|
||||
else:
|
||||
task.status = TaskStatus.FAILED
|
||||
@@ -384,6 +415,7 @@ class TaskQueue:
|
||||
task.assigned_agent_id = None
|
||||
task.started_at = None
|
||||
task.blocked_reason = None
|
||||
task.next_retry_at = None # a capacity release is not a failure — no backoff
|
||||
await self._save_task(task)
|
||||
await self.requeue_task(task_id)
|
||||
|
||||
@@ -417,6 +449,7 @@ class TaskQueue:
|
||||
task.started_at = None
|
||||
task.completed_at = None
|
||||
task.blocked_reason = None
|
||||
task.next_retry_at = None # a review reopen is not a failure — no backoff
|
||||
await self._save_task(task)
|
||||
await self.requeue_task(task_id)
|
||||
logger.info(f"Re-opened task {task_id} for another review cycle")
|
||||
|
||||
@@ -36,10 +36,56 @@ class FakeRun:
|
||||
|
||||
|
||||
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)
|
||||
# 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():
|
||||
@@ -223,7 +269,9 @@ def test_k8s_manifests():
|
||||
def main():
|
||||
import asyncio
|
||||
test_launch_count()
|
||||
test_pool_window()
|
||||
test_plan_specs()
|
||||
test_git_launch_env()
|
||||
test_resolve_model_key()
|
||||
test_resolve_git_grant()
|
||||
test_azkv_resolver()
|
||||
|
||||
@@ -64,6 +64,50 @@ def test_unit():
|
||||
check("unit: env limit default respected (=3 here)", orch.max_agents_per_user() == 3)
|
||||
|
||||
|
||||
def test_metadata_cap():
|
||||
# HM sends metadata.max_agents_per_user = min(plan SwarmMaxAgents, user override), only when > 0.
|
||||
# When present as a positive int it WINS over env; otherwise we fall back to env (=3 here).
|
||||
check("metadata cap (>0 int) wins over env",
|
||||
orch.max_agents_per_user({"metadata": {"max_agents_per_user": 7}}) == 7)
|
||||
check("body=None -> env cap (callers without a body unaffected)",
|
||||
orch.max_agents_per_user() == 3)
|
||||
check("body without the field -> env cap",
|
||||
orch.max_agents_per_user({"metadata": {}}) == 3)
|
||||
check("metadata cap = 0 ignored -> env cap (HM only sends when > 0)",
|
||||
orch.max_agents_per_user({"metadata": {"max_agents_per_user": 0}}) == 3)
|
||||
check("metadata cap negative ignored -> env cap",
|
||||
orch.max_agents_per_user({"metadata": {"max_agents_per_user": -5}}) == 3)
|
||||
check("metadata cap bool ignored -> env cap (True is not a valid count)",
|
||||
orch.max_agents_per_user({"metadata": {"max_agents_per_user": True}}) == 3)
|
||||
check("metadata cap non-int ignored -> env cap",
|
||||
orch.max_agents_per_user({"metadata": {"max_agents_per_user": "9"}}) == 3)
|
||||
|
||||
|
||||
async def test_ws_cap_per_agent():
|
||||
"""The register-time per-user cap resolves the owning run's metadata cap by agent_id, so it
|
||||
stays consistent with what launch_swarm_agents planned (no 'plan 12 but reject at 10')."""
|
||||
await orch.redis_client.connect() # REDIS_FAKE in-memory store (this test runs before the server boots)
|
||||
# Seed a run whose metadata carries a HM-sent cap of 5; launcher mints `{swarm_id}-agent-N`.
|
||||
run = await orch.swarm_runtime.get_or_create_run(
|
||||
{
|
||||
"mode": "swarm",
|
||||
"orchestration_plan": {"objective": "x"},
|
||||
"callback": {"url": "https://hm.example/cb"},
|
||||
"metadata": {"manager_deployment_id": "mdep-1", "max_agents_per_user": 5},
|
||||
},
|
||||
idempotency_key=None,
|
||||
correlation_id="corr-ws-cap",
|
||||
)
|
||||
run_obj = run[0]
|
||||
swarm_id = run_obj.swarm_id
|
||||
cap = await orch._per_user_cap_for_agent(f"{swarm_id}-agent-2")
|
||||
check("WS cap for launcher agent uses run metadata cap (=5, not env 3)", cap == 5)
|
||||
cap_ext = await orch._per_user_cap_for_agent("externally-supplied-agent")
|
||||
check("WS cap for non-launcher id falls back to env (=3)", cap_ext == 3)
|
||||
cap_missing = await orch._per_user_cap_for_agent("swarm-doesnotexist-agent-1")
|
||||
check("WS cap for unknown run falls back to env (=3)", cap_missing == 3)
|
||||
|
||||
|
||||
async def _register(agent_id, user_id=None):
|
||||
"""Open a WS, send register (optionally with user_id), return (ws, response_dict)."""
|
||||
ws = await websockets.connect(f"{WS}/ws/{agent_id}")
|
||||
@@ -142,6 +186,8 @@ async def test_integration():
|
||||
|
||||
async def main():
|
||||
test_unit()
|
||||
test_metadata_cap()
|
||||
await test_ws_cap_per_agent()
|
||||
await test_integration()
|
||||
print()
|
||||
if failures:
|
||||
|
||||
@@ -75,9 +75,13 @@ async def test_seeder():
|
||||
check("seed preserves base context", specs[0]["context"].get("orchestration_plan") == {"x": 1})
|
||||
check("seed marks is_seed + objective", specs[0]["context"].get("is_seed") is True
|
||||
and specs[0]["context"].get("objective") == "Build a calculator")
|
||||
# Manager-provided agents are honored (Manager-first; seeder does not override).
|
||||
check("Manager-provided agent breakdown bypasses the seeder",
|
||||
orch._manager_provided_agents({"orchestration_plan": {"agents": [{"role": "impl"}]}}) is True)
|
||||
# Seeding is UNCONDITIONAL (architect ruling 2026-06-15): a caller-supplied
|
||||
# orchestration_plan.agents no longer bypasses the seeder — every swarm still gets exactly one
|
||||
# seed task and the pool is auto-launched (the old _manager_provided_agents escape hatch is gone).
|
||||
specs_with_agents = orch.build_seed_task_specs(
|
||||
run, {"orchestration_plan": {"agents": [{"role": "impl"}]}}, base_specs)
|
||||
check("orchestration_plan.agents does NOT bypass the seeder (one seed task regardless)",
|
||||
len(specs_with_agents) == 1 and specs_with_agents[0]["source"] == "seed")
|
||||
|
||||
|
||||
async def test_agent_peer_routing():
|
||||
|
||||
Reference in New Issue
Block a user