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)


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