Agentswarm/scripts/test-max-agents-per-user.py

"""Per-user agent cap tests: at most MAX_AGENTS_PER_USER concurrent agents per user.

Covers:
  * unit — ConnectionManager per-user accounting (can_bind_user / bind_user / unbind / count) and
    the env-tunable limit (default 10);
  * integration — the REAL orchestrator over WebSocket: registrations beyond the cap are rejected
    (registration_rejected + socket close), a different user is unaffected, and freeing a slot
    (disconnect) lets a new agent in. Unbound agents (no user_id) are never capped.

Uses MAX_AGENTS_PER_USER=3 so the integration test stays small/fast.

Run from agent_swarm_v6 (install deps first):
    pip install -r orchestrator/requirements.txt -r agent/requirements.txt
    REDIS_FAKE=1 python scripts/test-max-agents-per-user.py
"""
import asyncio
import json
import logging
import os
import sys
import threading
from pathlib import Path

os.environ["REDIS_FAKE"] = "1"
os.environ["MAX_AGENTS_PER_USER"] = "3"

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import httpx
import uvicorn
import websockets
from orchestrator import main as orch

PORT = 8137
BASE = f"http://127.0.0.1:{PORT}"
WS = f"ws://127.0.0.1:{PORT}"

logging.getLogger("orchestrator.swarm_runtime").setLevel(logging.ERROR)

failures = []


def check(name, cond):
    print(("PASS" if cond else "FAIL"), "-", name)
    if not cond:
        failures.append(name)


def test_unit():
    cm = orch.ConnectionManager()
    limit = 3
    for i in range(limit):
        check(f"unit: agent a{i} allowed under cap", cm.can_bind_user(f"a{i}", "U", limit))
        cm.bind_user(f"a{i}", "U")
    check("unit: user count == 3", cm.user_agent_count("U") == 3)
    check("unit: 4th distinct agent rejected at cap", not cm.can_bind_user("a3", "U", limit))
    check("unit: reconnect by an already-counted agent allowed at cap", cm.can_bind_user("a0", "U", limit))
    check("unit: a different user is independent", cm.can_bind_user("b0", "V", limit))
    cm.unbind("a0")
    check("unit: count drops after unbind", cm.user_agent_count("U") == 2)
    check("unit: freed slot lets a new agent in", cm.can_bind_user("a3", "U", limit))
    cm.unbind("nonexistent")  # no-op, must not raise
    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}")
    msg = {"type": "register", "agent_id": agent_id, "capabilities": ["general"]}
    if user_id is not None:
        msg["user_id"] = user_id
    await ws.send(json.dumps(msg))
    resp = json.loads(await ws.recv())
    return ws, resp


async def wait_health(client):
    for _ in range(150):
        try:
            if (await client.get(f"{BASE}/health")).status_code == 200:
                return True
        except Exception:
            pass
        await asyncio.sleep(0.1)
    return False


async def test_integration():
    config = uvicorn.Config(orch.app, host="127.0.0.1", port=PORT, log_level="warning")
    server = uvicorn.Server(config)
    server.install_signal_handlers = lambda: None
    thread = threading.Thread(target=server.run, daemon=True)
    thread.start()
    open_ws = []
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            if not await wait_health(client):
                check("orchestrator started", False)
                return

            # 3 agents for one user all register OK
            for i in range(3):
                ws, resp = await _register(f"cap-{i}", user_id="user-cap")
                check(f"integration: agent cap-{i} registered", resp.get("type") == "registered")
                open_ws.append(ws)

            # 4th for the same user is rejected
            ws4, resp4 = await _register("cap-3", user_id="user-cap")
            check("integration: 4th agent rejected (over cap)", resp4.get("type") == "registration_rejected")
            check("integration: rejection states reason + limit",
                  resp4.get("reason") == "max_agents_per_user_exceeded" and resp4.get("limit") == 3)
            await ws4.close()

            # a different user is unaffected
            wsB, respB = await _register("other-0", user_id="user-other")
            check("integration: different user unaffected", respB.get("type") == "registered")
            open_ws.append(wsB)

            # freeing a slot (disconnect) lets a new agent for the capped user in
            await open_ws[0].close()
            open_ws.pop(0)
            await asyncio.sleep(0.4)  # let the server's finally-block run unbind
            ws5, resp5 = await _register("cap-4", user_id="user-cap")
            check("integration: freed slot admits a new agent", resp5.get("type") == "registered")
            open_ws.append(ws5)

            # unbound agents (no user_id) are never capped — many can register
            for i in range(5):
                ws, resp = await _register(f"unbound-{i}")  # no user_id
                check(f"integration: unbound agent unbound-{i} not capped", resp.get("type") == "registered")
                open_ws.append(ws)
    finally:
        for ws in open_ws:
            try:
                await ws.close()
            except Exception:
                pass
        server.should_exit = True
        thread.join(timeout=5)


async def main():
    test_unit()
    test_metadata_cap()
    await test_ws_cap_per_agent()
    await test_integration()
    print()
    if failures:
        print(f"{len(failures)} per-user-cap check(s) FAILED: {failures}")
        return 1
    print("all per-user agent-cap checks passed")
    return 0


if __name__ == "__main__":
    _code = asyncio.run(main())
    sys.stdout.flush()
    sys.stderr.flush()
    os._exit(_code or 0)