diff --git a/CLAUDE.md b/CLAUDE.md index 77b696e..8630c6e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -334,7 +334,8 @@ instances//hermes_data/ # маунт → /opt/data (~/.hermes) - **Skills env isolation**: service-agent передаёт subprocess только переменные из `env_vars` в `skills.json` — остальные env не попадают в скилл - **SQLite WAL**: все компоненты используют WAL mode + `journal_size_limit=1048576` (1MB); данные переживают рестарт контейнера - **Broker sender auth**: отправитель определяется по API-ключу, не по полю body — подменить sender невозможно -- **Streaming не трекается квотой**: `quota-proxy` считает токены только из non-streaming ответов (`{"usage": {...}}`). SSE-стримы (`text/event-stream`) пропускаются — использование не учитывается +- **SSE-стримы трекаются квотой**: `quota-proxy` использует потоковый парсер `SSEUsageParser` для учёта токенов в `text/event-stream` (Anthropic `message_start`/`message_delta` и OpenAI `usage` chunks) без буферизации всего ответа в памяти (с защитой 64 KB на строку). +- **Enterprise 1000-User Scale**: поддерживаются горячая регистрация ключей без рестарта (`POST /quota/register-key`, `POST /admin/register-key`), департаментные квоты (`POST /quota/set-department-limit`), пагинация отчётов и инбоксов, широковещательная рассылка (`POST /broadcast`), массовый провижининг до 1000 инстансов с шардированием (`scripts/bulk-users.py`, `make bulk-import`, `make fleet-status`) и нагрузочный бенчмарк (`scripts/benchmark-1000-users.py`, `make benchmark-1k`). Подробности: `docs/ENTERPRISE_1000_USERS.md`. - **Broker per-operation connections**: `broker.py` открывает/закрывает SQLite-соединение на каждую операцию (для WAL mode). `quota-proxy` наоборот — держит один `_conn` на всё время жизни - **Memory limits**: admin = 2G, user instances = 1G (`deploy.resources.limits`) - **Test import order**: env-переменные **ОБЯЗАТЕЛЬНО** ставить до `import` тестируемого модуля — модули читают env при импорте, не при вызове diff --git a/Makefile b/Makefile index 08bc435..0a1796d 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: add-user remove-user bootstrap deploy restart logs backup status quota-report quota-reset set-limit security-check monitor monitor-alerts add-service service-health service-skills install-hooks admin-overview +.PHONY: add-user remove-user bulk-import fleet-status bootstrap deploy restart logs backup status quota-report quota-reset set-limit security-check monitor monitor-alerts add-service service-health service-skills install-hooks admin-overview benchmark-1k # ── Управление инстансами ────────────────────────────────────────────────── @@ -18,6 +18,19 @@ bootstrap: @bash scripts/provision-hermes-data.sh @echo "✅ hermes_data готов" +# make bulk-import FILE=users.csv [SHARD_SIZE=100] — массовый онбординг до 1000 сотрудников +bulk-import: + @[ -n "$(FILE)" ] || (echo "❌ FILE= обязателен (CSV или JSON манифест)"; exit 1) + @python3 scripts/bulk-users.py import --file "$(FILE)" --shard-size "$${SHARD_SIZE:-100}" + +# make fleet-status — сводка по парку инстансов и шардам +fleet-status: + @python3 scripts/bulk-users.py status + +# make benchmark-1k — нагрузочный бенчмарк на 1000 конкурентных пользователей +benchmark-1k: + @python3 scripts/benchmark-1000-users.py + # ── Деплой ──────────────────────────────────────────────────────────────── deploy: bootstrap diff --git a/README.md b/README.md index 2147c58..5506d88 100644 --- a/README.md +++ b/README.md @@ -307,6 +307,9 @@ make status # all containers should be Up (healthy) | Command | Description | |---|---| | `make add-user NAME=x BOT_TOKEN=y TG_ID=z` | Onboard a new employee | +| `make bulk-import FILE=users.csv [SHARD_SIZE=100]` | Bulk-onboard up to 1,000 employees with Compose shard profiles | +| `make fleet-status` | Fleet integrity summary and shard distribution | +| `make benchmark-1k` | Run 1,000-concurrent-user load benchmark (`quota-proxy` + `broker`) | | `make remove-user NAME=x` | Remove instance (archives workspace to `backups/` first) | | `make bootstrap` | Provision `hermes_data/` (`.env`, `config.yaml`, `SOUL.md`, permissions) for instances | | `make deploy` | Run `bootstrap`, then start / rebuild all containers | @@ -436,6 +439,7 @@ CortexForge/ │ ├── add-service.sh # add service-agent replica │ └── backup.sh # workspace backup (tar.gz per instance) ├── docs/ +│ ├── ENTERPRISE_1000_USERS.md # 1000-user scale architecture, sharding, SSE quotas, bulk provisioning │ ├── SETUP.md # server installation guide │ ├── ONBOARDING.md # employee onboarding guide │ ├── PERSONAS.md # persona creation guide diff --git a/README.ru.md b/README.ru.md index bc49010..f98caf7 100644 --- a/README.ru.md +++ b/README.ru.md @@ -305,6 +305,9 @@ make status # все контейнеры должны быть Up (healthy) | Команда | Описание | |---|---| | `make add-user NAME=x BOT_TOKEN=y TG_ID=z` | Онбординг нового сотрудника | +| `make bulk-import FILE=users.csv [SHARD_SIZE=100]` | Массовый онбординг до 1000 сотрудников с шардированием Compose | +| `make fleet-status` | Сводка целостности парка инстансов и распределения по шардам | +| `make benchmark-1k` | Нагрузочный бенчмарк на 1000 конкурентных пользователей (`quota-proxy` + `broker`) | | `make remove-user NAME=x` | Удалить инстанс (сначала архивирует воркспейс в `backups/`) | | `make bootstrap` | Подготовить `hermes_data/` (`.env`, `config.yaml`, `SOUL.md`, права) для инстансов | | `make deploy` | Выполнить `bootstrap`, затем запустить / пересобрать все контейнеры | @@ -434,6 +437,7 @@ CortexForge/ │ ├── add-service.sh # добавить реплику service-agent │ └── backup.sh # бэкап воркспейсов (tar.gz на инстанс) ├── docs/ +│ ├── ENTERPRISE_1000_USERS.md # архитектура на 1000 пользователей, шардирование, SSE-квоты, bulk-users │ ├── SETUP.md # руководство по установке на сервер │ ├── ONBOARDING.md # руководство по онбордингу сотрудника │ ├── PERSONAS.md # руководство по созданию персонажей diff --git a/broker/broker.py b/broker/broker.py index e75437a..dfa5238 100644 --- a/broker/broker.py +++ b/broker/broker.py @@ -1,150 +1,537 @@ #!/usr/bin/env python3 """ -CortexForge Broker — message bus [hardened] +CortexForge Broker — Enterprise message bus [hardened, 1000-agent scale] -Hardening: - - Constant-time key comparison - - Rate limiting: 10 сообщений/мин на инстанс +Hardening & Enterprise capabilities: + - Constant-time key comparison (hmac.compare_digest) + - Thread-safe dynamic key management (SQLite broker_keys + env hot-reload) + - Per-instance rate limiting - Max message size: 10KB - - Security headers - - Message size limit in inbox - - SQLite persistence (WAL mode) — messages survive restarts + - Message priorities (low, normal, high, urgent) & custom per-message TTL + - Admin broadcast / multicast endpoint (/broadcast) + - Paginated & filtered inbox (/inbox with limit, offset, since_id, peek; /inbox/all) + - High-concurrency SQLite (WAL mode, synchronous=NORMAL, busy_timeout=10000, + composite indexes, throttled batch purge, ThreadingHTTPServer) """ -import json, os, signal, time, hashlib, hmac, threading, sqlite3 -from http.server import HTTPServer, BaseHTTPRequestHandler -from urllib.parse import urlparse, parse_qs +import hashlib +import hmac +import json +import os +import re +import signal +import sqlite3 +import threading +import time +from http.server import BaseHTTPRequestHandler, HTTPServer, ThreadingHTTPServer +from urllib.parse import parse_qs, urlparse MAX_MSG_BYTES = 10 * 1024 # 10KB на сообщение -MAX_INBOX = 100 -MAX_AGE_SEC = 86400 +MAX_INBOX = int(os.environ.get("MAX_INBOX", 100)) +MAX_AGE_SEC = int(os.environ.get("RETENTION_SECONDS", os.environ.get("MAX_AGE_SEC", 86400))) +RETENTION_SECONDS = MAX_AGE_SEC +PURGE_INTERVAL_SEC = float(os.environ.get("BROKER_PURGE_INTERVAL_SEC", 5.0)) + ADMIN_INSTANCE = "admin" # Instances that can act as proxy: read/send on behalf of other instances PROXY_INSTANCES = {"admin", "service"} +VALID_PRIORITIES = {"low", "normal", "high", "urgent"} +_INSTANCE_NAME_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$") BROKER_DB = os.environ.get("BROKER_DB", "/data/broker.db") -_keys: dict[str, str] = {} # {sha256(key): instance_name} +_keys_lock = threading.RLock() +_db_write_lock = threading.RLock() +_purge_lock = threading.Lock() +_last_purge_ts: float = 0.0 +_min_active_ttl: float = float(MAX_AGE_SEC) + +_keys: dict[str, str] = {} # {sha256(key): canonical_instance_name} +_instance_aliases: dict[str, str] = {} # {alias_or_name: canonical_instance_name} +_REPO_INSTANCES_DIR = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "instances" +) + + +def _retention_sec() -> float: + """Return active retention window in seconds, respecting runtime test overrides.""" + if RETENTION_SECONDS != 86400: + return float(RETENTION_SECONDS) + return float(MAX_AGE_SEC) + + +def _active_db_path() -> str: + """Resolve active SQLite database path, honoring runtime BROKER_DB overrides.""" + global BROKER_DB + env_db = os.environ.get("BROKER_DB") + if BROKER_DB == "/data/broker.db" and env_db: + BROKER_DB = env_db + return BROKER_DB + + +def _get_conn() -> sqlite3.Connection: + """Get a SQLite connection configured for high concurrency.""" + conn = sqlite3.connect(_active_db_path(), timeout=10.0) + conn.execute("PRAGMA busy_timeout=10000") + conn.execute("PRAGMA synchronous=NORMAL") + return conn + + +def _resolve_env_instance_name(suffix: str) -> tuple[str, str, str]: + """Return (canonical_name, raw_name, hyphen_name) for a BROKER_KEY_ variable.""" + raw_name = suffix.lower().strip() + hyphen_name = raw_name.replace("_", "-") + explicit = os.environ.get(f"BROKER_INSTANCE_{suffix}", "").lower().strip() + if explicit: + return explicit, raw_name, hyphen_name + + if hyphen_name != raw_name: + for base_dir in ("instances", _REPO_INSTANCES_DIR): + hyphen_dir = os.path.join(base_dir, hyphen_name) + raw_dir = os.path.join(base_dir, raw_name) + if os.path.isdir(hyphen_dir) and not os.path.isdir(raw_dir): + return hyphen_name, raw_name, hyphen_name + + return raw_name, raw_name, hyphen_name + + +def _load_keys() -> dict[str, str]: + """Load and merge broker keys from environment (BROKER_KEY_*) and SQLite (broker_keys).""" + keys: dict[str, str] = {} + aliases: dict[str, str] = {} -def _load_keys(): - keys = {} for k, v in os.environ.items(): if k.startswith("BROKER_KEY_") and v: - name = k[len("BROKER_KEY_"):].lower() - keys[hashlib.sha256(v.strip().encode()).hexdigest()] = name - return keys + suffix = k[len("BROKER_KEY_"):] + canonical, raw_name, hyphen_name = _resolve_env_instance_name(suffix) + if canonical: + keys[hashlib.sha256(v.strip().encode()).hexdigest()] = canonical + aliases[canonical] = canonical + aliases[raw_name] = canonical + aliases[hyphen_name] = canonical + + db_path = _active_db_path() + if db_path and os.path.exists(db_path): + try: + conn = _get_conn() + try: + rows = conn.execute( + "SELECT instance, key_hash FROM broker_keys" + ).fetchall() + for inst, key_hash in rows: + inst_norm = str(inst).lower().strip() + if not inst_norm or not key_hash: + continue + for old_h, old_inst in list(keys.items()): + if old_inst == inst_norm: + del keys[old_h] + keys[str(key_hash)] = inst_norm + aliases[inst_norm] = inst_norm + alt_hyphen = inst_norm.replace("_", "-") + alt_raw = inst_norm.replace("-", "_") + aliases.setdefault(alt_hyphen, inst_norm) + aliases.setdefault(alt_raw, inst_norm) + finally: + conn.close() + except sqlite3.Error: + pass + + with _keys_lock: + _keys.clear() + _keys.update(keys) + _instance_aliases.clear() + _instance_aliases.update(aliases) + return _keys + _keys = _load_keys() -# ── SQLite storage ───────────────────────────────────────────────────────── +def _known_instances() -> set[str]: + """Return thread-safe snapshot of canonical registered instance names.""" + with _keys_lock: + return set(_keys.values()) + + +def _known_with_aliases() -> set[str]: + """Return all canonical instance names and registered underscore/hyphen aliases.""" + with _keys_lock: + canonical = set(_keys.values()) + return canonical | { + alias for alias, target in _instance_aliases.items() if target in canonical + } + + +def _resolve_instance(name: str) -> str | None: + """Resolve an instance name or its underscore/hyphen alias to its canonical identity.""" + norm = name.lower().strip() + if not norm: + return None + with _keys_lock: + canonical_set = set(_keys.values()) + if norm in canonical_set: + return norm + mapped = _instance_aliases.get(norm) + if mapped and mapped in canonical_set: + return mapped + # Fallback if _keys was mutated directly in tests without updating _instance_aliases + alt_underscore = norm.replace("-", "_") + if alt_underscore in canonical_set: + return alt_underscore + alt_hyphen = norm.replace("_", "-") + if alt_hyphen in canonical_set: + return alt_hyphen + return None + + +def _register_key(instance: str, broker_key: str) -> str: + """Register or update an instance's broker key in SQLite and in-memory map.""" + inst_norm = instance.lower().strip() + key_hash = hashlib.sha256(broker_key.strip().encode()).hexdigest() + now_ts = int(time.time()) + + with _db_write_lock: + conn = _get_conn() + try: + conn.execute( + """ + INSERT INTO broker_keys (instance, key_hash, updated_at) + VALUES (?, ?, ?) + ON CONFLICT(instance) DO UPDATE SET + key_hash = excluded.key_hash, + updated_at = excluded.updated_at + """, + (inst_norm, key_hash, now_ts), + ) + conn.commit() + finally: + conn.close() -def _init_db(): - """Initialize SQLite database with WAL mode.""" - os.makedirs(os.path.dirname(BROKER_DB) or ".", exist_ok=True) - conn = sqlite3.connect(BROKER_DB) - conn.execute("PRAGMA journal_mode=WAL") - conn.execute("PRAGMA journal_size_limit=1048576") # 1MB - conn.execute(""" - CREATE TABLE IF NOT EXISTS messages ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - sender TEXT NOT NULL, - recipient TEXT NOT NULL, - body TEXT NOT NULL, - ts REAL NOT NULL - ) - """) - conn.execute(""" - CREATE INDEX IF NOT EXISTS idx_messages_recipient_ts - ON messages (recipient, ts) - """) - conn.commit() - conn.close() - - -def _get_conn(): - """Get a SQLite connection (WAL mode set once in _init_db).""" - conn = sqlite3.connect(BROKER_DB) - conn.execute("PRAGMA busy_timeout=5000") - return conn + with _keys_lock: + for old_h, old_inst in list(_keys.items()): + if old_inst == inst_norm: + del _keys[old_h] + _keys[key_hash] = inst_norm + _instance_aliases[inst_norm] = inst_norm + _instance_aliases.setdefault(inst_norm.replace("_", "-"), inst_norm) + _instance_aliases.setdefault(inst_norm.replace("-", "_"), inst_norm) + return key_hash -def _db_insert(sender: str, recipient: str, body: str, ts: float): - """Insert a message and enforce MAX_INBOX per recipient.""" - conn = _get_conn() - try: - conn.execute( - "INSERT INTO messages (sender, recipient, body, ts) VALUES (?, ?, ?, ?)", - (sender, recipient, body, ts), - ) - # Enforce MAX_INBOX: delete oldest messages beyond limit - conn.execute(""" - DELETE FROM messages WHERE id IN ( - SELECT id FROM messages - WHERE recipient = ? - ORDER BY ts DESC - LIMIT -1 OFFSET ? - ) - """, (recipient, MAX_INBOX)) - conn.commit() - finally: - conn.close() +# ── SQLite storage ───────────────────────────────────────────────────────── -def _db_inbox(recipient: str) -> list[dict]: - """Get inbox messages for a recipient (up to MAX_INBOX, excluding expired).""" - cutoff = time.time() - MAX_AGE_SEC - conn = _get_conn() - try: - rows = conn.execute( - "SELECT sender, recipient, body, ts FROM messages " - "WHERE recipient = ? AND ts > ? ORDER BY ts LIMIT ?", - (recipient, cutoff, MAX_INBOX), - ).fetchall() +def _init_db(): + """Initialize SQLite database with WAL mode, migrations, and enterprise indexes.""" + db_path = _active_db_path() + os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True) + with _db_write_lock: + conn = sqlite3.connect(db_path, timeout=10.0) + try: + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA synchronous=NORMAL") + conn.execute("PRAGMA busy_timeout=10000") + conn.execute("PRAGMA journal_size_limit=1048576") # 1MB + conn.execute(""" + CREATE TABLE IF NOT EXISTS messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + sender TEXT NOT NULL, + recipient TEXT NOT NULL, + body TEXT NOT NULL, + ts REAL NOT NULL, + priority TEXT DEFAULT 'normal', + delivered INTEGER DEFAULT 0, + created_at REAL DEFAULT 0, + expires_at REAL DEFAULT NULL + ) + """) + + # Safe migrations for pre-existing databases + existing_cols = { + r[1] for r in conn.execute("PRAGMA table_info(messages)").fetchall() + } + if "priority" not in existing_cols: + conn.execute("ALTER TABLE messages ADD COLUMN priority TEXT DEFAULT 'normal'") + if "delivered" not in existing_cols: + conn.execute("ALTER TABLE messages ADD COLUMN delivered INTEGER DEFAULT 0") + if "created_at" not in existing_cols: + conn.execute("ALTER TABLE messages ADD COLUMN created_at REAL DEFAULT 0") + conn.execute("UPDATE messages SET created_at = ts WHERE created_at = 0 OR created_at IS NULL") + if "expires_at" not in existing_cols: + conn.execute("ALTER TABLE messages ADD COLUMN expires_at REAL DEFAULT NULL") + + conn.execute(""" + CREATE TABLE IF NOT EXISTS broker_keys ( + instance TEXT PRIMARY KEY, + key_hash TEXT NOT NULL, + updated_at INTEGER + ) + """) + + conn.execute(""" + CREATE INDEX IF NOT EXISTS idx_messages_recipient_ts + ON messages (recipient, ts) + """) + conn.execute(""" + CREATE INDEX IF NOT EXISTS idx_messages_recipient_delivered_created + ON messages (recipient, delivered, created_at) + """) + conn.execute(""" + CREATE INDEX IF NOT EXISTS idx_messages_expires_at + ON messages (expires_at) + """) + conn.commit() + finally: + conn.close() + + _load_keys() + + +def _db_insert( + sender: str, + recipient: str, + body: str, + ts: float, + priority: str = "normal", + ttl_sec: float | None = None, +) -> int: + """Insert a message and enforce MAX_INBOX per recipient.""" + global _min_active_ttl + prio = priority if priority in VALID_PRIORITIES else "normal" + retention = _retention_sec() + expires_at = None + if ttl_sec is not None: + effective_ttl = max(0.001, min(float(ttl_sec), retention)) + expires_at = ts + effective_ttl + if effective_ttl <= 5.0: + _min_active_ttl = min(_min_active_ttl, effective_ttl) + + with _db_write_lock: + conn = _get_conn() + try: + cur = conn.execute( + """ + INSERT INTO messages (sender, recipient, body, ts, priority, delivered, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, 0, ?, ?) + """, + (sender, recipient, body, ts, prio, ts, expires_at), + ) + msg_id = int(cur.lastrowid or 0) + # Enforce MAX_INBOX: delete oldest messages beyond limit + conn.execute( + """ + DELETE FROM messages WHERE id IN ( + SELECT id FROM messages + WHERE recipient = ? + ORDER BY ts DESC, id DESC + LIMIT -1 OFFSET ? + ) + """, + (recipient, MAX_INBOX), + ) + conn.commit() + return msg_id + finally: + conn.close() + + +def _db_broadcast( + sender: str, + recipients: list[str], + body: str, + ts: float, + priority: str = "normal", + ttl_sec: float | None = None, +) -> int: + """Insert broadcast message to multiple recipients in a single SQLite transaction.""" + global _min_active_ttl + if not recipients: + return 0 + prio = priority if priority in VALID_PRIORITIES else "normal" + retention = _retention_sec() + expires_at = None + if ttl_sec is not None: + effective_ttl = max(0.001, min(float(ttl_sec), retention)) + expires_at = ts + effective_ttl + if effective_ttl <= 5.0: + _min_active_ttl = min(_min_active_ttl, effective_ttl) + + with _db_write_lock: + conn = _get_conn() + try: + conn.executemany( + """ + INSERT INTO messages (sender, recipient, body, ts, priority, delivered, created_at, expires_at) + VALUES (?, ?, ?, ?, ?, 0, ?, ?) + """, + [(sender, r, body, ts, prio, ts, expires_at) for r in recipients], + ) + for r in recipients: + conn.execute( + """ + DELETE FROM messages WHERE id IN ( + SELECT id FROM messages + WHERE recipient = ? + ORDER BY ts DESC, id DESC + LIMIT -1 OFFSET ? + ) + """, + (r, MAX_INBOX), + ) + conn.commit() + return len(recipients) + finally: + conn.close() + + +def _db_inbox( + recipient: str, + limit: int | None = None, + offset: int = 0, + since_id: int | None = None, + peek: bool = False, +) -> list[dict]: + """Get inbox messages for a recipient with optional pagination, since_id, and peek mode.""" + now = time.time() + cutoff = now - _retention_sec() + eff_limit = MAX_INBOX if limit is None else max(0, min(int(limit), MAX_INBOX)) + eff_offset = max(0, int(offset)) + + query = ( + "SELECT id, sender, recipient, body, ts, priority, delivered " + "FROM messages " + "WHERE recipient = ? AND ts > ? AND (expires_at IS NULL OR expires_at > ?)" + ) + params: list = [recipient, cutoff, now] + + if since_id is not None: + query += " AND id > ?" + params.append(int(since_id)) + + query += " ORDER BY ts ASC, id ASC LIMIT ? OFFSET ?" + params.extend([eff_limit, eff_offset]) + + if peek: + conn = _get_conn() + try: + rows = conn.execute(query, tuple(params)).fetchall() + finally: + conn.close() return [ - {"from": r[0], "to": r[1], "message": r[2], "ts": r[3]} + { + "id": r[0], + "from": r[1], + "to": r[2], + "message": r[3], + "ts": r[4], + "priority": r[5] or "normal", + "delivered": bool(r[6]), + } for r in rows ] - finally: - conn.close() + + with _db_write_lock: + conn = _get_conn() + try: + rows = conn.execute(query, tuple(params)).fetchall() + undelivered_ids = [r[0] for r in rows if not r[6]] + if undelivered_ids: + placeholders = ",".join("?" for _ in undelivered_ids) + conn.execute( + f"UPDATE messages SET delivered = 1 WHERE id IN ({placeholders})", + tuple(undelivered_ids), + ) + conn.commit() + finally: + conn.close() + + return [ + { + "id": r[0], + "from": r[1], + "to": r[2], + "message": r[3], + "ts": r[4], + "priority": r[5] or "normal", + "delivered": True, + } + for r in rows + ] def _db_clear(recipient: str) -> tuple[int, int]: """Delete all messages for a recipient. Returns (deleted, remaining).""" - conn = _get_conn() - try: - conn.execute("DELETE FROM messages WHERE recipient = ?", (recipient,)) - count = conn.execute("SELECT changes()").fetchone()[0] - conn.commit() - remaining = conn.execute( - "SELECT COUNT(*) FROM messages WHERE recipient = ?", (recipient,) - ).fetchone()[0] - return count, remaining - finally: - conn.close() + with _db_write_lock: + conn = _get_conn() + try: + conn.execute("DELETE FROM messages WHERE recipient = ?", (recipient,)) + count = conn.execute("SELECT changes()").fetchone()[0] + conn.commit() + remaining = conn.execute( + "SELECT COUNT(*) FROM messages WHERE recipient = ?", (recipient,) + ).fetchone()[0] + return count, remaining + finally: + conn.close() def _db_counts() -> dict[str, int]: - """Get message count per recipient.""" + """Get active message count per recipient.""" + now = time.time() + cutoff = now - _retention_sec() conn = _get_conn() try: rows = conn.execute( - "SELECT recipient, COUNT(*) FROM messages GROUP BY recipient" + """ + SELECT recipient, COUNT(*) FROM messages + WHERE ts > ? AND (expires_at IS NULL OR expires_at > ?) + GROUP BY recipient + """, + (cutoff, now), ).fetchall() return {r[0]: r[1] for r in rows} finally: conn.close() -def _cleanup(): - """Delete expired messages.""" - cutoff = time.time() - MAX_AGE_SEC - conn = _get_conn() +def _purge_expired(force: bool = False) -> int: + """Batch-delete expired messages, throttled unless force=True or short TTL is active.""" + global _last_purge_ts + now = time.time() + retention = _retention_sec() + immediate = force or (retention <= 5.0) or (_min_active_ttl <= 5.0) + + if not immediate and (now - _last_purge_ts) < PURGE_INTERVAL_SEC: + return 0 + + if not _purge_lock.acquire(blocking=immediate): + return 0 + try: - conn.execute("DELETE FROM messages WHERE ts < ?", (cutoff,)) - conn.commit() + now = time.time() + if not immediate and (now - _last_purge_ts) < PURGE_INTERVAL_SEC: + return 0 + cutoff = now - retention + with _db_write_lock: + conn = _get_conn() + try: + conn.execute( + "DELETE FROM messages WHERE ts < ? OR (expires_at IS NOT NULL AND expires_at <= ?)", + (cutoff, now), + ) + deleted = conn.execute("SELECT changes()").fetchone()[0] + conn.commit() + _last_purge_ts = now + return deleted + finally: + conn.close() finally: - conn.close() + _purge_lock.release() + + +def _cleanup(): + """Immediately delete expired messages (backward-compatible entrypoint).""" + return _purge_expired(force=True) # ── Rate limiter ──────────────────────────────────────────────────────────── @@ -157,27 +544,40 @@ def __init__(self, rate: float, capacity: int): def allow(self, key: str) -> bool: now = time.monotonic() with self._l: + if len(self._b) > 10000: + self._b = {k: v for k, v in self._b.items() if now - v[1] <= 120} tokens, last = self._b.get(key, (self.capacity, now)) tokens = min(self.capacity, tokens + (now - last) * self.rate) if tokens < 1: - self._b[key] = (tokens, now); return False - self._b[key] = (tokens - 1, now); return True + self._b[key] = (tokens, now) + return False + self._b[key] = (tokens - 1, now) + return True + + +RATE_LIMIT_PER_MIN = float(os.environ.get("BROKER_RATE_LIMIT_PER_MIN", 10)) +RATE_LIMIT_BURST = int(os.environ.get("BROKER_RATE_LIMIT_BURST", 10)) +_send_rl = RateLimiter(rate=RATE_LIMIT_PER_MIN / 60.0, capacity=RATE_LIMIT_BURST) -_send_rl = RateLimiter(rate=10/60, capacity=10) # 10 сообщений в минуту # Constant-time auth def _auth(req) -> str | None: auth = req.headers.get("Authorization", "") - if not auth.startswith("Bearer "): return None + if not auth.startswith("Bearer "): + return None h = hashlib.sha256(auth[7:].encode()).hexdigest() result = None - for k, name in _keys.items(): + with _keys_lock: + items = list(_keys.items()) + for k, name in items: if hmac.compare_digest(h, k): result = name return result + _SEC_HEADERS = {"X-Content-Type-Options": "nosniff", "Cache-Control": "no-store"} + class Handler(BaseHTTPRequestHandler): def log_message(self, fmt, *args): @@ -188,20 +588,34 @@ def _json(self, code, body): self.send_response(code) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", len(data)) - for k, v in _SEC_HEADERS.items(): self.send_header(k, v) + for k, v in _SEC_HEADERS.items(): + self.send_header(k, v) self.end_headers() self.wfile.write(data) + def _parse_content_length(self) -> int | None: + raw_len = self.headers.get("Content-Length", "0") + try: + length = int(raw_len) + except (TypeError, ValueError): + self._json(400, {"error": "Invalid Content-Length"}) + return None + if length < 0: + self._json(400, {"error": "Invalid Content-Length"}) + return None + return length + def do_GET(self): parsed = urlparse(self.path) path = parsed.path + qs = parse_qs(parsed.query) if path == "/health": - _cleanup() + _purge_expired() counts = _db_counts() self._json(200, { "status": "ok", - "instances": sorted(set(_keys.values())), + "instances": sorted(_known_instances()), "message_counts": counts, }) return @@ -209,113 +623,380 @@ def do_GET(self): if path == "/inbox/all": sender = _auth(self) if not sender: - self._json(403, {"error": "Unauthorized"}); return + self._json(403, {"error": "Unauthorized"}) + return if sender != ADMIN_INSTANCE: - self._json(403, {"error": "Admin only"}); return - _cleanup() - counts = _db_counts() - self._json(200, {"inboxes": counts}) + self._json(403, {"error": "Admin only"}) + return + + try: + offset = int(qs.get("offset", ["0"])[0]) + limit_raw = qs.get("limit", [None])[0] + limit = int(limit_raw) if limit_raw is not None else None + if offset < 0 or (limit is not None and limit < 0): + raise ValueError + except ValueError: + self._json(400, {"error": "Invalid pagination parameters"}) + return + + _purge_expired() + all_counts = _db_counts() + sorted_items = sorted(all_counts.items()) + total_inboxes = len(sorted_items) + + if limit is not None: + paged_items = sorted_items[offset: offset + limit] + else: + paged_items = sorted_items[offset:] + + self._json(200, { + "inboxes": dict(paged_items), + "total_inboxes": total_inboxes, + "limit": limit if limit is not None else total_inboxes, + "offset": offset, + }) return if path == "/inbox": sender = _auth(self) if not sender: - self._json(403, {"error": "Unauthorized"}); return + self._json(403, {"error": "Unauthorized"}) + return target = sender - qs = parse_qs(parsed.query) for_param = qs.get("for", [None])[0] if for_param: if sender not in PROXY_INSTANCES: - self._json(403, {"error": "Proxy access denied"}); return - target = for_param.lower().strip() - known = set(_keys.values()) - if target not in known: - self._json(404, {"error": "Unknown instance"}); return - - _cleanup() - msgs = _db_inbox(target) - self._json(200, {"inbox": msgs, "count": len(msgs)}) + self._json(403, {"error": "Proxy access denied"}) + return + resolved_target = _resolve_instance(for_param) + if not resolved_target: + self._json(404, {"error": "Unknown instance"}) + return + target = resolved_target + + try: + limit_raw = qs.get("limit", [None])[0] + limit = int(limit_raw) if limit_raw is not None else None + offset = int(qs.get("offset", ["0"])[0]) + since_raw = qs.get("since_id", [None])[0] + since_id = int(since_raw) if since_raw is not None else None + if offset < 0 or (limit is not None and limit < 0): + raise ValueError + except ValueError: + self._json(400, {"error": "Invalid query parameters"}) + return + + peek_raw = str(qs.get("peek", ["0"])[0]).lower().strip() + peek = peek_raw in ("1", "true", "yes") + + _purge_expired() + msgs = _db_inbox( + target, + limit=limit, + offset=offset, + since_id=since_id, + peek=peek, + ) + self._json(200, { + "inbox": msgs, + "count": len(msgs), + "limit": limit if limit is not None else MAX_INBOX, + "offset": offset, + "peek": peek, + }) return self._json(404, {"error": "Not found"}) def do_POST(self): - if urlparse(self.path).path != "/send": - self._json(404, {"error": "Not found"}); return + path = urlparse(self.path).path + + if path == "/admin/register-key": + sender = _auth(self) + if not sender: + self._json(403, {"error": "Unauthorized"}) + return + if sender != ADMIN_INSTANCE: + self._json(403, {"error": "Admin only"}) + return + + length = self._parse_content_length() + if length is None: + return + if length > MAX_MSG_BYTES: + self._json(413, {"error": "Payload too large"}) + return + try: + body = json.loads(self.rfile.read(length)) if length > 0 else {} + except Exception: + self._json(400, {"error": "Invalid JSON"}) + return + + instance = str(body.get("instance", "")).lower().strip() + broker_key = str(body.get("broker_key") or body.get("key") or "").strip() + if not instance or not broker_key: + self._json(400, {"error": "Missing 'instance' or 'broker_key'"}) + return + if not _INSTANCE_NAME_RE.match(instance): + self._json(400, {"error": "Invalid 'instance' name (expected ^[a-z0-9][a-z0-9_-]{0,63}$)"}) + return + + _register_key(instance, broker_key) + known = _known_instances() + self._json(200, { + "ok": True, + "instance": instance, + "total_instances": len(known), + }) + return + + if path == "/admin/reload-keys": + sender = _auth(self) + if not sender: + self._json(403, {"error": "Unauthorized"}) + return + if sender != ADMIN_INSTANCE: + self._json(403, {"error": "Admin only"}) + return + + _load_keys() + known = sorted(_known_with_aliases()) + self._json(200, { + "ok": True, + "instances": known, + "total_instances": len(_known_instances()), + }) + return + + if path == "/broadcast": + sender = _auth(self) + if not sender: + self._json(403, {"error": "Unauthorized"}) + return + if sender != ADMIN_INSTANCE: + self._json(403, {"error": "Admin only"}) + return + + length = self._parse_content_length() + if length is None: + return + if length > MAX_MSG_BYTES * 4: + self._json(413, {"error": "Payload too large"}) + return + + try: + body = json.loads(self.rfile.read(length)) if length > 0 else {} + except Exception: + self._json(400, {"error": "Invalid JSON"}) + return + + raw_text = body.get("text") if body.get("text") is not None else body.get("message", "") + text = str(raw_text).strip() + if not text: + self._json(400, {"error": "Missing 'text' or 'message'"}) + return + if len(text.encode()) > MAX_MSG_BYTES: + self._json(413, {"error": f"Message too large (max {MAX_MSG_BYTES // 1024}KB)"}) + return + + priority = str(body.get("priority", "normal") or "normal").lower().strip() + if priority not in VALID_PRIORITIES: + self._json(400, {"error": f"Invalid priority: {priority}"}) + return + + ttl_sec = body.get("ttl_sec") + if ttl_sec is not None: + try: + ttl_sec = float(ttl_sec) + if ttl_sec <= 0: + raise ValueError + except (TypeError, ValueError): + self._json(400, {"error": "Invalid 'ttl_sec'"}) + return + + to_target = body.get("to") + known = _known_instances() + if to_target == "*" or to_target == ["*"]: + recipients = sorted(inst for inst in known if inst != sender) + elif isinstance(to_target, list): + raw_list = [str(x).lower().strip() for x in to_target if str(x).strip()] + if not raw_list: + self._json(400, {"error": "Empty 'to' recipient list"}) + return + resolved_list = [] + unknown = [] + for r in raw_list: + res = _resolve_instance(r) + if res is None: + unknown.append(r) + else: + resolved_list.append(res) + if unknown: + self._json(404, {"error": f"Unknown recipients: {', '.join(unknown)}"}) + return + recipients = list(dict.fromkeys(r for r in resolved_list if r != sender)) + elif isinstance(to_target, str) and to_target.strip(): + raw_list = [x.lower().strip() for x in to_target.split(",") if x.strip()] + resolved_list = [] + unknown = [] + for r in raw_list: + res = _resolve_instance(r) + if res is None: + unknown.append(r) + else: + resolved_list.append(res) + if unknown: + self._json(404, {"error": f"Unknown recipients: {', '.join(unknown)}"}) + return + recipients = list(dict.fromkeys(r for r in resolved_list if r != sender)) + else: + self._json(400, {"error": "Missing or invalid 'to' field"}) + return + + _purge_expired() + delivered_count = _db_broadcast( + sender, + recipients, + text[:MAX_MSG_BYTES], + time.time(), + priority=priority, + ttl_sec=ttl_sec, + ) + self._json(200, { + "ok": True, + "from": sender, + "recipients": recipients, + "delivered_count": delivered_count, + "count": delivered_count, + "priority": priority, + }) + return + + if path != "/send": + self._json(404, {"error": "Not found"}) + return sender = _auth(self) if not sender: - self._json(403, {"error": "Unauthorized"}); return + self._json(403, {"error": "Unauthorized"}) + return # Rate limit if not _send_rl.allow(sender): - self._json(429, {"error": "Rate limit exceeded (10 msg/min)"}); return + self._json(429, {"error": "Rate limit exceeded (10 msg/min)"}) + return - length = int(self.headers.get("Content-Length", 0)) + length = self._parse_content_length() + if length is None: + return if length > MAX_MSG_BYTES: - self._json(413, {"error": f"Message too large (max {MAX_MSG_BYTES//1024}KB)"}); return + self._json(413, {"error": f"Message too large (max {MAX_MSG_BYTES // 1024}KB)"}) + return try: - body = json.loads(self.rfile.read(length)) - to = str(body.get("to", "")).lower().strip() - message = str(body.get("message", "")).strip() + body = json.loads(self.rfile.read(length)) + to = str(body.get("to", "")).lower().strip() + raw_msg = body.get("message") if body.get("message") is not None else body.get("text", "") + message = str(raw_msg).strip() except Exception: - self._json(400, {"error": "Invalid JSON"}); return + self._json(400, {"error": "Invalid JSON"}) + return if not to or not message: - self._json(400, {"error": "Missing 'to' or 'message'"}); return + self._json(400, {"error": "Missing 'to' or 'message'"}) + return + + priority = str(body.get("priority", "normal") or "normal").lower().strip() + if priority not in VALID_PRIORITIES: + self._json(400, {"error": f"Invalid priority: {priority}"}) + return + + ttl_sec = body.get("ttl_sec") + if ttl_sec is not None: + try: + ttl_sec = float(ttl_sec) + if ttl_sec <= 0: + raise ValueError + except (TypeError, ValueError): + self._json(400, {"error": "Invalid 'ttl_sec'"}) + return # Proxy: service/admin can send on behalf of another instance effective_sender = sender on_behalf_of = str(body.get("on_behalf_of", "")).lower().strip() if on_behalf_of: if sender not in PROXY_INSTANCES: - self._json(403, {"error": "Proxy send denied"}); return - known_names = set(_keys.values()) - if on_behalf_of not in known_names: - self._json(404, {"error": f"Unknown sender: {on_behalf_of}"}); return - effective_sender = on_behalf_of - - known = set(_keys.values()) - if to not in known: - self._json(404, {"error": f"Unknown recipient: {to}"}); return - if to == effective_sender: - self._json(400, {"error": "Cannot send to yourself"}); return - - _cleanup() - _db_insert(effective_sender, to, message[:MAX_MSG_BYTES], time.time()) - - print(f"[broker] {effective_sender} → {to}: {message[:60]}" - + (f" (via {sender})" if effective_sender != sender else ""), - flush=True) - self._json(200, {"ok": True, "from": effective_sender, "to": to}) + self._json(403, {"error": "Proxy send denied"}) + return + resolved_sender = _resolve_instance(on_behalf_of) + if not resolved_sender: + self._json(404, {"error": f"Unknown sender: {on_behalf_of}"}) + return + effective_sender = resolved_sender + + resolved_to = _resolve_instance(to) + if not resolved_to: + self._json(404, {"error": f"Unknown recipient: {to}"}) + return + if resolved_to == effective_sender: + self._json(400, {"error": "Cannot send to yourself"}) + return + + _purge_expired() + msg_id = _db_insert( + effective_sender, + resolved_to, + message[:MAX_MSG_BYTES], + time.time(), + priority=priority, + ttl_sec=ttl_sec, + ) + + print( + f"[broker] {effective_sender} → {resolved_to} [{priority}]: {message[:60]}" + + (f" (via {sender})" if effective_sender != sender else ""), + flush=True, + ) + self._json(200, { + "ok": True, + "id": msg_id, + "from": effective_sender, + "to": to, + "priority": priority, + }) def do_DELETE(self): parsed = urlparse(self.path) if parsed.path != "/inbox": - self._json(404, {"error": "Not found"}); return + self._json(404, {"error": "Not found"}) + return sender = _auth(self) if not sender: - self._json(403, {"error": "Unauthorized"}); return + self._json(403, {"error": "Unauthorized"}) + return target = sender qs = parse_qs(parsed.query) for_param = qs.get("for", [None])[0] if for_param: if sender not in PROXY_INSTANCES: - self._json(403, {"error": "Proxy access denied"}); return - target = for_param.lower().strip() - known = set(_keys.values()) - if target not in known: - self._json(404, {"error": "Unknown instance"}); return + self._json(403, {"error": "Proxy access denied"}) + return + resolved_target = _resolve_instance(for_param) + if not resolved_target: + self._json(404, {"error": "Unknown instance"}) + return + target = resolved_target count, remaining = _db_clear(target) - print(f"[broker] {target} очистил inbox ({count} сообщений)" - + (f" (via {sender})" if target != sender else ""), - flush=True) + print( + f"[broker] {target} очистил inbox ({count} сообщений)" + + (f" (via {sender})" if target != sender else ""), + flush=True, + ) self._json(200, {"ok": True, "deleted": count, "remaining": remaining}) @@ -328,9 +1009,9 @@ def do_DELETE(self): if not _keys: print("⚠️ Нет API-ключей!", flush=True) else: - print(f"✅ Broker hardened :{port} | instances: {sorted(set(_keys.values()))}", flush=True) + print(f"✅ Broker hardened :{port} | instances: {sorted(_known_instances())}", flush=True) - server = HTTPServer(("0.0.0.0", port), Handler) + server = ThreadingHTTPServer(("0.0.0.0", port), Handler) def _shutdown(signum, frame): print(f"[broker] Получен сигнал {signum}, завершаем...", flush=True) diff --git a/docs/ENTERPRISE_1000_USERS.md b/docs/ENTERPRISE_1000_USERS.md new file mode 100644 index 0000000..30e1956 --- /dev/null +++ b/docs/ENTERPRISE_1000_USERS.md @@ -0,0 +1,315 @@ +# Enterprise Architecture: Масштабирование CortexForge до 1000 пользователей + +Данный документ описывает архитектуру, механизмы провижининга, управление квотами, шину сообщений, мониторинг и нагрузочное тестирование **CortexForge Enterprise** в конфигурации до **1 000 изолированных инстансов Hermes Agent**. + +--- + +## 1. Обзор архитектуры 1000-User Enterprise + +При масштабировании с десятков до 1 000 сотрудников основные узкие места классической Docker Compose-инсталляции лежат в четырёх плоскостях: +1. **Сеть и iptables / bridge-лимиты Linux** — помещение 1 000 контейнеров в один Docker bridge приводит к деградации ARP/broadcast-трафика и исчерпанию лимита портов Linux bridge (по умолчанию 1023 интерфейса на мост). +2. **Конкурентный доступ к SQLite и блокировки HTTP-серверов** — однопоточный `HTTPServer` блокируется на медленных запросах или длинных SSE-стримах, а частые `DELETE` при очистке очередей вызывают `SQLITE_BUSY`. +3. **Онбординг и ротация ключей** — перезапуск центральных сервисов (`quota-proxy`, `message-broker`) при добавлении каждого нового сотрудника недопустим при парке в 1 000 ботов. +4. **Агрегация телеметрии и отчётов** — выгрузка полного списка из 1 000 инстансов в одном HTTP-ответе или запуск 1 000 процессов `docker inspect` в цикле перегружает демон Docker и админ-панель. + +### Ключевые архитектурные решения Enterprise-режима + +| Подсистема | Решение для 1 000 пользователей | +|---|---| +| **Сетевое шардирование** | Разбиение парка на профили и изолированные bridge-сети `shard-0` .. `shard-9` (по умолчанию по 100 инстансов на шард) | +| **Массовый онбординг** | Утилита `scripts/bulk-users.py` (`make bulk-import`, `make fleet-status`) с `fcntl.flock`, защитой от симлинков и авто-назначением портов | +| **Многопоточные серверы** | Все core-компоненты (`quota-proxy`, `broker`, `service-agent`, `resource-monitor`) переведены на `ThreadingHTTPServer` | +| **Хранилище SQLite** | Режим `WAL`, `synchronous=NORMAL`, `busy_timeout=10000`, `journal_size_limit=1MB`, композитные индексы и `threading.RLock` | +| **Hot-Reload ключей** | Персистентные таблицы `quota_keys` и `broker_keys` в SQLite + API регистрации без рестарта контейнеров | +| **Потоковый учёт SSE** | Построчный парсинг `text/event-stream` (`SSEUsageParser`) в `quota-proxy` на лету без буферизации ответа в памяти | +| **Иерархия квот** | Индивидуальные лимиты инстансов + групповые лимиты департаментов (`department_limits`) | +| **Пагинация и фильтрация** | Поддержка `limit`, `offset`, `department`, `status`, `search`, `sort` во всех отчётных API | + +--- + +## 2. Шардирование парка контейнеров (`shard-0` .. `shard-9`) + +### Профили Docker Compose и сетевая топология + +При импорте 1 000 пользователей через `scripts/bulk-users.py` каждый сервис `assistant-` в `docker-compose.override.yml` автоматически привязывается к своему шарду: +- Индекс шарда вычисляется как `shard_id = user_index // shard_size` (при `shard_size = 100` создаются шарды `shard-0` .. `shard-9`). +- Каждому инстансу назначаются Compose-профили: + - `shard-` — для пошардового управления запуском, обновлением и раскладкой по нодам; + - `all-users` — для одновременного управления всем парком на мощном хосте. +- Для каждого шарда генерируется собственная изолированная сеть `corp-shard-` (`driver: bridge`, `internal: true`), а также метка `corp.assistant.shard: shard-` и `corp.assistant.department: `. + +Пример сгенерированного блока сервиса в `docker-compose.override.yml`: + +```yaml +services: + assistant-user-0001: + build: + context: . + dockerfile: instances/Dockerfile.user + image: cortex-forge-user:latest + container_name: corp-user-0001 + restart: unless-stopped + profiles: + - "shard-0" + - "all-users" + env_file: + - instances/user-0001/.env + volumes: + - ./instances/user-0001/hermes_data:/opt/data + - ./instances/user-0001/workspace:/opt/data/workspace + - ./shared/skills:/shared/skills + - ./shared/docs:/shared/docs + - ./shared/compliance-data:/shared/compliance-data:ro + environment: + - ANTHROPIC_API_KEY=${QUOTA_KEY_USER-0001} + - OPENAI_API_KEY=${OPENAI_API_KEY:-} + - SERVICE_API_KEY=${SERVICE_API_KEY} + - BROKER_URL=http://message-broker:8080 + - BROKER_KEY=${BROKER_KEY_USER-0001} + - HERMES_UID=1000 + - HERMES_GID=1000 + - HERMES_DASHBOARD=1 + ports: + - "127.0.0.1:18790:9119" + networks: + - corp-internal + - corp-outbound + - corp-shard-0 + security_opt: + - no-new-privileges:true + deploy: + resources: + limits: + cpus: "2.0" + memory: 1G + depends_on: + quota-proxy: + condition: service_healthy + message-broker: + condition: service_healthy + labels: + corp.assistant.role: personal + corp.assistant.user: user-0001 + corp.assistant.shard: shard-0 + corp.assistant.department: engineering + +networks: + corp-shard-0: + driver: bridge + internal: true +``` + +### Команды пошардового управления + +```bash +# Запустить только шард 0 (первые 100 пользователей) +docker compose --profile shard-0 up -d + +# Перезапустить шард 3 +docker compose --profile shard-3 restart + +# Запустить все шарды сразу +docker compose --profile all-users up -d +``` + +В `scripts/sync-instances.sh` проверка живости 1 000 контейнеров оптимизирована за $O(1)$ через единый вызов `docker ps -a --format '{{.Names}}={{.State}}'`, а подъём остановившихся сервисов выполняется батчами по 50 (`xargs -n 50 docker compose up -d`), чтобы не вызывать пиковую перегрузку демона Docker. + +--- + +## 3. Массовый провижининг (`scripts/bulk-users.py`) + +CLI-инструмент `scripts/bulk-users.py` автоматизирует онбординг, экспорт и аудит целостности до 1 000+ инстансов без внешних зависимостей (только Python 3.12+ stdlib). + +### Поддерживаемые форматы входных файлов + +1. **CSV (`users.csv`)**: + ```csv + name,bot_token,full_name,tg_id,department,quota_limit,ui_port + user-0001,7100000001:AAHxxxx1,Иван Петров,100000001,engineering,2000000,18790 + user-0002,7100000002:AAHxxxx2,Мария Сидорова,100000002,sales,1000000, + ``` +2. **JSON (`users.json`)** — список объектов или `{"users": [...]}` с теми же полями (`department`, `quota_limit`, `ui_port` опциональны). + +### Гарантии безопасности и атомарности в `bulk-users.py` + +- **Строгая валидация полей**: + - `name`: `^[a-z][a-z0-9_-]{1,31}$` (имя `admin` зарезервировано); + - `bot_token`: `^[0-9]+:[A-Za-z0-9_-]+$`; + - `tg_id`: `^[0-9]+(?:,[0-9]+)*$` (обязателен — защита от пустого allowlist); + - `department`: `^[A-Za-z0-9_.-]{1,64}$`; + - `ui_port`: диапазон `1024..65535` (исключая админский порт `18789`), автоматическое предотвращение коллизий с уже занятыми портами в `docker-compose.override.yml`. +- **Защита от symlink-атак**: перед чтением шаблона и записью любых файлов проверяется отсутствие символических ссылок (`_verify_tree_no_symlinks`, `_ensure_not_symlink`, `follow_symlinks=False`). +- **Атомарная запись файлов**: функция `_safe_write_file` создаёт временный файл через `tempfile.mkstemp` в целевой директории, выставляет права (`0600` для `.env`, `config.yaml`, `SOUL.md`), вызывает `os.fsync()` и атомарно заменяет целевой файл через `os.replace()`. +- **Межпроцессная блокировка `.env`**: дозапись ключей `BROKER_KEY_*`, `QUOTA_KEY_*`, `QUOTA_LIMIT_*` в корневой `.env` защищена эксклюзивным файловым локом `fcntl.flock(..., fcntl.LOCK_EX)` на `.env.bulk.lock`. +- **Идемпотентность**: уже существующие каталоги `instances/` безопасно пропускаются (`skipped`) без перезаписи пользовательских данных. + +### Команды CLI и Makefile + +```bash +# Предварительная проверка манифеста без записи на диск (dry-run) +python3 scripts/bulk-users.py import --file users.csv --shard-size 100 --dry-run + +# Массовый импорт через Makefile +make bulk-import FILE=users.csv SHARD_SIZE=100 + +# Проверка состояния парка, распределения по шардам и отсутствия коллизий портов +make fleet-status + +# Экспорт манифеста парка (без секретов — 0644, с секретами — 0600) +python3 scripts/bulk-users.py export --output fleet.json +python3 scripts/bulk-users.py export --output fleet-secrets.json --include-secrets +``` + +--- + +## 4. Горячая регистрация ключей в `quota-proxy` и `broker` + +Для подключения новых ботов или ротации скомпрометированных токенов на лету (без рестарта `corp-quota` и `corp-broker`) оба сервиса хранят ключи в SQLite и синхронизируют их с in-memory хэш-таблицами под защитой `threading.RLock`. + +### Хранение и сравнение ключей + +- В памяти и в SQLite (`quota_keys`, `broker_keys`) хранятся **только SHA-256 хэши** ключей. +- Проверка авторизации выполняется за константное время через `hmac.compare_digest()` по снимку списка ключей, взятому под `_keys_lock`, что исключает timing-атаки и `RuntimeError: dictionary changed size during iteration` при параллельной регистрации. +- При ротации ключа старый хэш того же инстанса автоматически удаляется из памяти и перезаписывается в SQLite (`ON CONFLICT(instance) DO UPDATE`). + +### API `quota-proxy` (порт `9090`, авторизация `Bearer `) + +1. **`POST /quota/register-key`** — зарегистрировать или обновить ключ инстанса, департамент и индивидуальный лимит: + ```bash + curl -s -X POST http://127.0.0.1:9090/quota/register-key \ + -H "Authorization: Bearer ${QUOTA_ADMIN_TOKEN}" \ + -H "Content-Type: application/json" \ + -d '{ + "instance": "user-0001", + "quota_key": "quota-user-0001-xxxx", + "department": "engineering", + "limit": 2000000 + }' + ``` + *(Также поддерживается передача готового SHA-256 хэша в поле `"key_hash"` вместо `"quota_key"`).* + +2. **`POST /quota/reload-keys`** — перечитать ключи из переменных окружения и таблицы `quota_keys`: + ```bash + curl -s -X POST http://127.0.0.1:9090/quota/reload-keys \ + -H "Authorization: Bearer ${QUOTA_ADMIN_TOKEN}" + ``` + +### API `message-broker` (порт `8080`, авторизация `Bearer `) + +1. **`POST /admin/register-key`** — горячая регистрация ключа инстанса в шине сообщений: + ```bash + curl -s -X POST http://message-broker:8080/admin/register-key \ + -H "Authorization: Bearer ${BROKER_KEY_ADMIN}" \ + -H "Content-Type: application/json" \ + -d '{ + "instance": "user-0001", + "broker_key": "brk-xxxx" + }' + ``` + +2. **`POST /admin/reload-keys`** — горячая перезагрузка таблицы `broker_keys` и env: + ```bash + curl -s -X POST http://message-broker:8080/admin/reload-keys \ + -H "Authorization: Bearer ${BROKER_KEY_ADMIN}" + ``` + +--- + +## 5. Потоковый учёт SSE-стримов (`SSEUsageParser`) + +В базовой версии `quota-proxy` учитывал токены только для обычных JSON-ответов, пропуская `Content-Type: text/event-stream`. В Enterprise-режиме реализован потоковый перехват и учёт токенов для SSE без задержки первого токена (TTFT) и без буферизации всего потока в RAM. + +### Как работает `_stream_sse` и `SSEUsageParser` + +1. При получении от upstream заголовка `Content-Type: text/event-stream` метод `_proxy()` передаёт соединение в `_stream_sse()`. +2. Чтение из upstream выполняется построчно (`upstream_resp.readline()`) и немедленно пересылается клиенту (`self.wfile.write(line)` + `self.wfile.flush()`). +3. Каждая строка передаётся в `SSEUsageParser.feed_line(line)`: + - **Формат Anthropic SSE**: + - из события `message_start` извлекается `message.usage.input_tokens` (и начальный `output_tokens`); + - из финального события `message_delta` извлекается итоговый `usage.output_tokens` (и обновлённый `input_tokens`, если передан). + - **Формат OpenAI / совместимых прокси**: + - из чанков `data: {...}` с полем `usage` извлекаются `prompt_tokens` / `input_tokens` и `completion_tokens` / `output_tokens`. +4. В блоке `finally` (даже если клиент оборвал соединение `BrokenPipeError` / `ConnectionResetError` до конца генерации) накопленные `input_tokens` и `output_tokens` атомарно записываются в таблицу `usage` через `_record_usage_log()`. + +--- + +## 6. Департаментные квоты (`department_limits`) + +Помимо индивидуальных лимитов каждого сотрудника, `quota-proxy` поддерживает квотирование на уровне подразделений (`engineering`, `sales`, `legal`, `support` и т.д.). + +### Логика проверки (`_check_quota`) + +Перед отправкой запроса в upstream проверяются два условия: +1. **Индивидуальная квота инстанса**: если `usage[instance] >= limits[instance]`, возвращается `HTTP 429` (`quota_exceeded`). +2. **Квота департамента**: определяется департамент инстанса (из таблицы `quota_keys` или переменной `QUOTA_DEPT_`). Если для департамента задан лимит `department_limits[dept] > 0` и суммарный расход всех инстансов этого департамента за текущий месяц достиг лимита, запрос блокируется с `HTTP 429` (`Квота департамента исчерпана`). + +### Установка лимита департамента + +```bash +curl -s -X POST http://127.0.0.1:9090/quota/set-department-limit \ + -H "Authorization: Bearer ${QUOTA_ADMIN_TOKEN}" \ + -H "Content-Type: application/json" \ + -d '{ + "department": "engineering", + "limit": 50000000 + }' +``` + +В отчёте `GET /quota/report` возвращается секция `departments` со сводкой по каждому подразделению (`input_tokens`, `output_tokens`, `total_tokens`, `requests`, `instances`, `limit`, `used_pct`, `status`). + +--- + +## 7. Шина сообщений, Service-Agent и пагинация API + +При 1 000 активных ботов все внутренние API поддерживают пагинацию, фильтрацию и защиту от шумных соседей (noisy neighbor). + +### Шина сообщений (`broker/broker.py`) + +- **Приоритеты и TTL**: `POST /send` и `POST /broadcast` принимают `priority` (`low`, `normal`, `high`, `urgent`) и `ttl_sec` (индивидуальное время жизни сообщения в секундах). +- **Широковещательная рассылка (`POST /broadcast`)**: + - Доступна только `admin`; + - Поле `to` принимает `"*"` (все зарегистрированные инстансы кроме отправителя), массив `["user-1", "user-2"]` или строку через запятую; + - Вставка выполняется в одной SQLite-транзакции через `conn.executemany()` с удержанием лимита `MAX_INBOX` на каждого получателя. +- **Троттлинг очистки (`_purge_expired`)**: удаление просроченных сообщений выполняется не чаще раза в `BROKER_PURGE_INTERVAL_SEC` (по умолчанию 5 секунд) под неблокирующим `_purge_lock`, что устраняет write-contention при сотнях параллельных чтений `/inbox`. +- **Пагинация и `peek`-режим**: + - `GET /inbox?limit=50&offset=0&since_id=1200&peek=true` — чтение сообщений с курсором `since_id`; при `peek=true` сообщения не помечаются как `delivered=1`. + - `GET /inbox/all?limit=100&offset=0` — пагинированная сводка размеров очередей по всем инстансам для `admin`. + +### Пагинация и фильтрация остальных сервисов + +| Сервис | Эндпоинт | Параметры запроса | +|---|---|---| +| `quota-proxy` | `GET /quota/report` | `month=YYYY-MM`, `department=`, `status=ok\|warning\|exceeded`, `limit=N`, `offset=M` | +| `service-agent` | `GET /v1/usage` | `caller=`, `skill=`, `limit=N`, `offset=M` | +| `resource-monitor` | `GET /metrics` | `search=`, `sort=cpu_pct\|mem_pct\|name`, `order=asc\|desc`, `limit=N`, `offset=M` | +| `resource-monitor` | `GET /alerts`, `/alerts/active` | `instance=`, `limit=N`, `offset=M` | + +В `service-agent` дополнительно введён per-caller rate limiter (`CallerRateLimiter`, переменная `SKILL_RATE_LIMIT`, по умолчанию `30` вызовов/мин на одного вызывающего бота) и настраиваемый пул `MAX_CONCURRENT_SKILLS`. + +--- + +## 8. Нагрузочное тестирование (`make benchmark-1k`) + +Цель `make benchmark-1k` в `Makefile` запускает нагрузочный бенчмарк для проверки способности `quota-proxy`, `message-broker`, `service-agent` и `resource-monitor` обслуживать 1 000 конкурентных инстансов. + +### Запуск тестов и бенчмарка + +```bash +# Прогон всех Enterprise unit/integration тестов (160 тестов, включая конкурентные стресс-тесты) +python3 -m unittest discover tests/ -v + +# Отдельные тестовые наборы Enterprise-компонентов: +python3 -m unittest tests/test_quota_enterprise.py -v +python3 -m unittest tests/test_broker_enterprise.py -v +python3 -m unittest tests/test_bulk_provisioning.py -v +python3 -m unittest tests/test_service_monitor_enterprise.py -v + +# Запуск нагрузочного бенчмарка на 1000 пользователей (при наличии скрипта benchmark-1000-users.py) +make benchmark-1k +``` + +### Проверяемые инварианты под нагрузкой + +1. **Отсутствие `SQLITE_BUSY` и гонок данных**: параллельная регистрация сотен ключей и одновременный учёт токенов/сообщений из 50–100 потоков дают точное совпадение счётчиков в SQLite. +2. **Линейная масштабируемость памяти**: потоковый `SSEUsageParser` обрабатывает длинные ответы константным объёмом памяти (`O(1)` по числу чанков). +3. **Изоляция сбоев**: превышение квоты одного департамента или rate-limit одного бота (`SKILL_RATE_LIMIT`) не влияет на задержки остальных инстансов. diff --git a/quota-proxy/proxy.py b/quota-proxy/proxy.py index 5f5f22b..616bfa5 100644 --- a/quota-proxy/proxy.py +++ b/quota-proxy/proxy.py @@ -15,8 +15,9 @@ - Лимиты в SQLite (переживают рестарт, меняются без рестарта) """ -import json, os, pwd, grp, signal, sqlite3, threading, time, hmac, hashlib -from http.server import HTTPServer, BaseHTTPRequestHandler +import json, os, pwd, grp, re, signal, sqlite3, threading, time, hmac, hashlib +from contextlib import contextmanager +from http.server import HTTPServer, ThreadingHTTPServer, BaseHTTPRequestHandler from urllib.parse import urlparse, parse_qs from urllib.request import Request, urlopen from urllib.error import HTTPError, URLError @@ -42,20 +43,56 @@ if not ADMIN_TOKEN: raise RuntimeError("QUOTA_ADMIN_TOKEN не задан — запуск небезопасен") +_INSTANCE_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$") +_KEY_HASH_RE = re.compile(r"^[0-9a-fA-F]{64}$") +MAX_SSE_LINE_BYTES = 65536 + # Квота-ключи: QUOTA_KEY_USER_1=quota-user-1-xxx → {key_hash: "user-1"} # Храним sha256 хэш, а не сам ключ — чтобы ключ не жил в памяти как plaintext +_keys_lock = threading.RLock() _quota_keys: dict[str, str] = {} # {sha256(key): instance_name} +_departments: dict[str, str] = {} # {instance_name: department} +_instance_departments = _departments _admin_hash: str = "" +_conn: sqlite3.Connection | None = None def _load_keys(): - global _quota_keys, _admin_hash - keys = {} + global _quota_keys, _departments, _instance_departments, _admin_hash + keys: dict[str, str] = {} + depts: dict[str, str] = {} for k, v in os.environ.items(): if k.startswith("QUOTA_KEY_") and v: - name = k[len("QUOTA_KEY_"):].lower() + raw_suffix = k[len("QUOTA_KEY_"):] + name = raw_suffix.lower().strip() + hyphen_name = name.replace("_", "-") + if v.strip().lower().startswith(f"quota-{hyphen_name}-"): + name = hyphen_name + dept_env = ( + os.environ.get(f"QUOTA_DEPT_{raw_suffix.upper()}") + or os.environ.get(f"QUOTA_DEPT_{name.upper()}") + or "default" + ).strip().lower() or "default" keys[hashlib.sha256(v.strip().encode()).hexdigest()] = name - _quota_keys = keys - _admin_hash = hashlib.sha256(ADMIN_TOKEN.encode()).hexdigest() + depts[name] = dept_env + if _conn is not None: + try: + with _db_lock: + rows = _conn.execute("SELECT instance, key_hash, department FROM quota_keys").fetchall() + for inst, khash, dept in rows: + inst = inst.lower().strip() + # Удаляем старый хэш этого инстанса (если был из env), чтобы ротация работала атомарно + stale = [h for h, n in keys.items() if n == inst] + for h in stale: + del keys[h] + keys[khash] = inst + depts[inst] = (dept or "default").strip().lower() or "default" + except sqlite3.Error: + pass + with _keys_lock: + _quota_keys = keys + _departments = depts + _instance_departments = _departments + _admin_hash = hashlib.sha256(ADMIN_TOKEN.encode()).hexdigest() _load_keys() @@ -66,11 +103,16 @@ def __init__(self, rate: float, capacity: int): self.rate = rate self.capacity = capacity self._buckets: dict[str, tuple[float, float]] = {} + self._b = self._buckets self._lock = threading.Lock() def allow(self, key: str) -> bool: now = time.monotonic() with self._lock: + if len(self._buckets) > 10000: + stale_keys = [k for k, (_, ts) in self._buckets.items() if now - ts > 120] + for k in stale_keys: + del self._buckets[k] tokens, last = self._buckets.get(key, (self.capacity, now)) tokens = min(self.capacity, tokens + (now - last) * self.rate) if tokens < 1: @@ -87,16 +129,52 @@ def allow(self, key: str) -> bool: _broker_rl = RateLimiter(rate=0.17, capacity=10) # ── SQLite ───────────────────────────────────────────────────────────────── -_db_lock = threading.Lock() +_db_lock = threading.RLock() +_admission_meta_lock = threading.Lock() +_instance_admission_locks: dict[str, threading.RLock] = {} +_dept_admission_locks: dict[str, threading.RLock] = {} + +def _get_instance_admission_lock(instance: str) -> threading.RLock: + with _admission_meta_lock: + lock = _instance_admission_locks.get(instance) + if lock is None: + lock = threading.RLock() + _instance_admission_locks[instance] = lock + return lock + +def _get_dept_admission_lock(department: str) -> threading.RLock: + with _admission_meta_lock: + lock = _dept_admission_locks.get(department) + if lock is None: + lock = threading.RLock() + _dept_admission_locks[department] = lock + return lock + +@contextmanager +def _quota_admission_guard(instance: str): + """Сериализует окно check_quota -> upstream -> add_usage для одного инстанса (и департамента с лимитом).""" + dept = _get_department(instance) + dept_limit = _get_department_limit(dept) + inst_lock = _get_instance_admission_lock(instance) + if dept_limit > 0: + dept_lock = _get_dept_admission_lock(dept) + with dept_lock: + with inst_lock: + yield + else: + with inst_lock: + yield def _now(): return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) def _init_db() -> sqlite3.Connection: os.makedirs(os.path.dirname(DB_PATH), exist_ok=True) - conn = sqlite3.connect(DB_PATH, check_same_thread=False) - # WAL mode: намного лучше под нагрузкой (readers не блокируют writer) + conn = sqlite3.connect(DB_PATH, timeout=10.0, check_same_thread=False) + # WAL mode + busy_timeout: выдерживает 1000 конкурентных пользователей без блокировок conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") + conn.execute("PRAGMA busy_timeout=10000") + conn.execute("PRAGMA journal_size_limit=1048576") conn.executescript(""" CREATE TABLE IF NOT EXISTS usage ( instance TEXT NOT NULL, @@ -106,12 +184,26 @@ def _init_db() -> sqlite3.Connection: requests INTEGER DEFAULT 0, PRIMARY KEY (instance, month) ); + CREATE INDEX IF NOT EXISTS idx_usage_month_instance ON usage(month, instance); CREATE TABLE IF NOT EXISTS limits ( instance TEXT PRIMARY KEY, monthly_limit INTEGER NOT NULL, updated_at TEXT NOT NULL, updated_by TEXT DEFAULT 'system' ); + CREATE TABLE IF NOT EXISTS quota_keys ( + instance TEXT PRIMARY KEY, + key_hash TEXT NOT NULL, + department TEXT DEFAULT 'default', + updated_at INTEGER + ); + CREATE INDEX IF NOT EXISTS idx_quota_keys_department ON quota_keys(department); + CREATE TABLE IF NOT EXISTS department_limits ( + department TEXT PRIMARY KEY, + monthly_limit INTEGER NOT NULL, + updated_at TEXT NOT NULL, + updated_by TEXT DEFAULT 'system' + ); CREATE TABLE IF NOT EXISTS audit_log ( id INTEGER PRIMARY KEY AUTOINCREMENT, ts TEXT NOT NULL, @@ -125,7 +217,7 @@ def _init_db() -> sqlite3.Connection: # Начальные лимиты из env (только если в БД ещё нет) for k, v in os.environ.items(): if k.startswith("QUOTA_LIMIT_") and v: - name = k[len("QUOTA_LIMIT_"):].lower() + name = k[len("QUOTA_LIMIT_"):].lower().strip().replace("_", "-") try: limit = int(v) except ValueError: continue conn.execute("INSERT OR IGNORE INTO limits (instance, monthly_limit, updated_at, updated_by) VALUES (?,?,?,'env')", @@ -170,6 +262,7 @@ def _drop_privileges(user: str = "app"): _drop_privileges() # ← вызываем ДО открытия БД (чтобы открыть с правильным uid) _conn = _init_db() +_load_keys() # ← загружаем ключи из env + таблицы quota_keys def _audit(action: str, actor: str, target: str = None, detail: str = None): with _db_lock: @@ -177,6 +270,106 @@ def _audit(action: str, actor: str, target: str = None, detail: str = None): (_now(), action, actor, target, detail)) _conn.commit() +# ── Динамические ключи и департаменты ────────────────────────────────────── +def _get_department(instance: str) -> str: + with _keys_lock: + dept = _departments.get(instance) + if dept: + return dept + with _db_lock: + row = _conn.execute("SELECT department FROM quota_keys WHERE instance=?", (instance,)).fetchone() + if row and row[0]: + dept = row[0].strip().lower() or "default" + with _keys_lock: + _departments[instance] = dept + return dept + return "default" + +def _register_key(instance: str, raw_key: str = None, key_hash: str = None, + department: str = "default", limit: int | None = None, by: str = "admin-api") -> str: + instance = instance.lower().strip() + if not _INSTANCE_RE.match(instance): + raise ValueError("Недопустимое имя инстанса (ожидается ^[a-z0-9][a-z0-9_-]{0,63}$)") + department = (department or "default").strip().lower() or "default" + if not _INSTANCE_RE.match(department): + raise ValueError("Недопустимое имя департамента") + if raw_key and str(raw_key).strip(): + key_hash = hashlib.sha256(str(raw_key).strip().encode()).hexdigest() + elif key_hash and str(key_hash).strip(): + key_hash = str(key_hash).strip().lower() + if not _KEY_HASH_RE.match(key_hash): + raise ValueError("Некорректный key_hash (ожидается 64 hex-символа SHA256)") + else: + raise ValueError("Required quota_key or key_hash") + now_ts = int(time.time()) + with _db_lock: + _conn.execute("""INSERT INTO quota_keys (instance, key_hash, department, updated_at) + VALUES (?,?,?,?) + ON CONFLICT(instance) DO UPDATE SET + key_hash=excluded.key_hash, + department=excluded.department, + updated_at=excluded.updated_at""", + (instance, key_hash, department, now_ts)) + if limit is not None: + _conn.execute("""INSERT INTO limits (instance, monthly_limit, updated_at, updated_by) + VALUES (?,?,?,?) + ON CONFLICT(instance) DO UPDATE SET + monthly_limit=excluded.monthly_limit, + updated_at=excluded.updated_at, + updated_by=excluded.updated_by""", + (instance, int(limit), _now(), by)) + _conn.commit() + with _keys_lock: + stale = [h for h, n in _quota_keys.items() if n == instance] + for h in stale: + del _quota_keys[h] + _quota_keys[key_hash] = instance + _departments[instance] = department + return key_hash + +def _get_department_limit(department: str) -> int: + department = (department or "default").strip().lower() or "default" + with _db_lock: + row = _conn.execute("SELECT monthly_limit FROM department_limits WHERE department=?", (department,)).fetchone() + return row[0] if row else 0 + +def _set_department_limit(department: str, limit: int, by: str = "admin-api"): + department = (department or "default").strip().lower() or "default" + with _db_lock: + _conn.execute("""INSERT INTO department_limits (department, monthly_limit, updated_at, updated_by) + VALUES (?,?,?,?) + ON CONFLICT(department) DO UPDATE SET + monthly_limit=excluded.monthly_limit, + updated_at=excluded.updated_at, + updated_by=excluded.updated_by""", + (department, int(limit), _now(), by)) + _conn.commit() + +def _get_all_department_limits() -> dict: + with _db_lock: + rows = _conn.execute("SELECT department, monthly_limit, updated_at, updated_by FROM department_limits").fetchall() + return {r[0]: {"limit": r[1], "updated_at": r[2], "updated_by": r[3]} for r in rows} + +def _get_department_usage(department: str, month: str = None) -> int: + month = month or time.strftime("%Y-%m") + department = (department or "default").strip().lower() or "default" + with _keys_lock: + dept_cache = dict(_departments) + with _db_lock: + rows = _conn.execute( + """SELECT u.instance, u.input_tokens, u.output_tokens, q.department + FROM usage u + LEFT JOIN quota_keys q ON u.instance = q.instance + WHERE u.month = ?""", + (month,), + ).fetchall() + total = 0 + for inst, inp, out, db_dept in rows: + inst_dept = (db_dept or dept_cache.get(inst) or "default").strip().lower() or "default" + if inst_dept == department: + total += (inp or 0) + (out or 0) + return total + # ── Лимиты ───────────────────────────────────────────────────────────────── def _get_limit(instance: str) -> int: with _db_lock: @@ -221,34 +414,92 @@ def _get_usage(instance: str, month: str = None) -> dict: if not row: return {"input": 0, "output": 0, "total": 0, "requests": 0} return {"input": row[0], "output": row[1], "total": row[0]+row[1], "requests": row[2]} -def _get_all_usage(month: str = None) -> list: +def _get_all_usage(month: str = None, department: str = None, status_filter: str = None) -> list: month = month or time.strftime("%Y-%m") limits = _get_all_limits() with _db_lock: rows = _conn.execute( "SELECT instance, input_tokens, output_tokens, requests FROM usage WHERE month=? " - "ORDER BY (input_tokens+output_tokens) DESC", (month,)).fetchall() + "ORDER BY (input_tokens+output_tokens) DESC, instance ASC", (month,)).fetchall() + dept_filter = department.strip().lower() if department else None + norm_status = status_filter.strip().lower() if status_filter else None result = [] for inst, inp, out, reqs in rows: + dept = _get_department(inst) + if dept_filter and dept != dept_filter: + continue total = inp + out limit = limits.get(inst, {}).get("limit", DEFAULT_MONTHLY) pct = round(total/limit*100, 1) if limit > 0 else None status = ("❌ exceeded" if limit>0 and total>=limit else "⚠️ warning" if limit>0 and pct>=80 else "✅ ok") - result.append({"instance": inst, "month": month, + if norm_status: + # Поддерживаем фильтр как по 'ok'/'warning'/'exceeded', так и по полному статусу + status_code = "exceeded" if "exceeded" in status else "warning" if "warning" in status else "ok" + if norm_status not in (status_code, status.lower()): + continue + result.append({"instance": inst, "department": dept, "month": month, "input_tokens": inp, "output_tokens": out, "total_tokens": total, "requests": reqs, "limit": limit if limit>0 else "unlimited", "used_pct": pct, "status": status}) return result +def _get_departments_summary(month: str = None) -> dict: + month = month or time.strftime("%Y-%m") + dept_limits = _get_all_department_limits() + with _db_lock: + rows = _conn.execute( + "SELECT instance, input_tokens, output_tokens, requests FROM usage WHERE month=?", + (month,)).fetchall() + summary: dict[str, dict] = {} + for dept_name, info in dept_limits.items(): + summary[dept_name] = { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + "requests": 0, + "instances": 0, + } + for inst, inp, out, reqs in rows: + dept = _get_department(inst) + entry = summary.setdefault(dept, { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0, + "requests": 0, + "instances": 0, + }) + entry["input_tokens"] += inp or 0 + entry["output_tokens"] += out or 0 + entry["total_tokens"] += (inp or 0) + (out or 0) + entry["requests"] += reqs or 0 + entry["instances"] += 1 + for dept, entry in summary.items(): + dlim = dept_limits.get(dept, {}).get("limit", 0) + total = entry["total_tokens"] + pct = round(total / dlim * 100, 1) if dlim > 0 else None + status = ("❌ exceeded" if dlim > 0 and total >= dlim + else "⚠️ warning" if dlim > 0 and pct >= 80 + else "✅ ok") + entry["limit"] = dlim if dlim > 0 else "unlimited" + entry["used_pct"] = pct + entry["status"] = status + return summary + def _check_quota(instance: str) -> tuple[bool, str]: limit = _get_limit(instance) - if limit <= 0: return True, "" - total = _get_usage(instance)["total"] - if total >= limit: - return False, f"Квота {instance} исчерпана: {total:,}/{limit:,} токенов. Обратись к администратору." + if limit > 0: + total = _get_usage(instance)["total"] + if total >= limit: + return False, f"Квота {instance} исчерпана: {total:,}/{limit:,} токенов. Обратись к администратору." + dept = _get_department(instance) + dept_limit = _get_department_limit(dept) + if dept_limit > 0: + dept_total = _get_department_usage(dept) + if dept_total >= dept_limit: + return False, f"Квота департамента {dept} исчерпана: {dept_total:,}/{dept_limit:,} токенов. Обратись к администратору." return True, "" def _reset_usage(instance: str, month: str = None): @@ -257,6 +508,101 @@ def _reset_usage(instance: str, month: str = None): _conn.execute("DELETE FROM usage WHERE instance=? AND month=?", (instance, month)) _conn.commit() +# ── SSE Usage Parser ─────────────────────────────────────────────────────── +class SSEUsageParser: + """Потоковый парсер SSE-событий Anthropic и OpenAI для учёта токенов без задержки стрима.""" + + def __init__(self): + self.input_tokens: int = 0 + self.output_tokens: int = 0 + self._current_event: str = "" + self._remainder: bytes = b"" + + def feed_chunk(self, chunk: bytes): + if not chunk: + return + data = self._remainder + chunk + lines = data.split(b"\n") + self._remainder = lines.pop() + if len(self._remainder) > MAX_SSE_LINE_BYTES: + self._remainder = self._remainder[-MAX_SSE_LINE_BYTES:] + for raw_line in lines: + self.feed_line(raw_line) + + def finish(self): + if self._remainder: + self.feed_line(self._remainder) + self._remainder = b"" + + def feed_line(self, raw_line: bytes | str): + line = raw_line.decode("utf-8", errors="replace") if isinstance(raw_line, bytes) else raw_line + line = line.strip("\r\n").strip() + if not line: + self._current_event = "" + return + if line.startswith("event:"): + self._current_event = line[6:].strip() + return + if not line.startswith("data:"): + return + data_str = line[5:].strip() + if not data_str or data_str == "[DONE]": + return + try: + payload = json.loads(data_str) + except Exception: + return + if not isinstance(payload, dict): + return + + ev_type = payload.get("type") or self._current_event + + # Anthropic: message_start -> message.usage.input_tokens + if ev_type == "message_start": + msg = payload.get("message") + if isinstance(msg, dict): + u = msg.get("usage") + if isinstance(u, dict): + if u.get("input_tokens") is not None: + self.input_tokens = int(u["input_tokens"]) + if u.get("output_tokens") is not None and self.output_tokens == 0: + self.output_tokens = int(u["output_tokens"]) + return + + # Anthropic: message_delta -> usage.output_tokens + if ev_type == "message_delta": + u = payload.get("usage") + if isinstance(u, dict): + if u.get("input_tokens") is not None and int(u["input_tokens"]) > 0: + self.input_tokens = int(u["input_tokens"]) + if u.get("output_tokens") is not None: + self.output_tokens = int(u["output_tokens"]) + return + + # OpenAI / совместимые SSE: data: {... "usage": {"prompt_tokens": ..., "completion_tokens": ...}} + u = payload.get("usage") + if isinstance(u, dict): + if u.get("prompt_tokens") is not None: + self.input_tokens = int(u["prompt_tokens"]) + elif u.get("input_tokens") is not None: + self.input_tokens = int(u["input_tokens"]) + if u.get("completion_tokens") is not None: + self.output_tokens = int(u["completion_tokens"]) + elif u.get("output_tokens") is not None: + self.output_tokens = int(u["output_tokens"]) + + +def _parse_sse_usage(raw_body: bytes | str) -> tuple[int, int]: + parser = SSEUsageParser() + if isinstance(raw_body, bytes): + parser.feed_chunk(raw_body) + parser.finish() + else: + for line in raw_body.splitlines(): + parser.feed_line(line) + return parser.input_tokens, parser.output_tokens + + # ── HTTP Handler ──────────────────────────────────────────────────────────── _SEC_HEADERS = { "X-Content-Type-Options": "nosniff", @@ -298,19 +644,32 @@ def _get_instance(self) -> str | None: key = auth[7:] if not key: return None key_hash = hashlib.sha256(key.encode()).hexdigest() - # Constant-time поиск — не раскрывает наличие ключа через время + # Constant-time поиск под RLock — не раскрывает наличие ключа через время + with _keys_lock: + items = list(_quota_keys.items()) result = None - for h, name in _quota_keys.items(): + for h, name in items: if hmac.compare_digest(key_hash, h): result = name return result def _read_body(self) -> bytes | None: - """Читает тело с лимитом размера.""" - length = int(self.headers.get("Content-Length", 0)) + """Читает тело с лимитом размера и валидацией Content-Length.""" + raw_len = self.headers.get("Content-Length", "0") + try: + length = int(raw_len) + except (ValueError, TypeError): + self.close_connection = True + self._json(400, {"error": "Invalid Content-Length"}) + return None + if length < 0: + self.close_connection = True + self._json(400, {"error": "Invalid Content-Length"}) + return None if length > MAX_BODY_BYTES: self.close_connection = True - return None # тело слишком большое + self._json(413, {"error": "Request body too large (max 1MB)"}) + return None return self.rfile.read(length) if length else b"" def _client_ip(self) -> str: @@ -322,8 +681,10 @@ def do_GET(self): qs = parse_qs(urlparse(self.path).query) if path == "/quota/health": + with _keys_lock: + instances = sorted(set(_quota_keys.values())) self._json(200, {"status": "ok", - "instances": sorted(set(_quota_keys.values())), + "instances": instances, "limits": _get_all_limits()}) return @@ -334,10 +695,31 @@ def do_GET(self): _audit("REPORT_DENIED", self._client_ip()) self._json(403, {"error": "Unauthorized"}); return month = qs.get("month", [time.strftime("%Y-%m")])[0] + department = qs.get("department", [None])[0] + status_filter = qs.get("status", [None])[0] + try: + page_limit = int(qs["limit"][0]) if "limit" in qs and qs["limit"][0] != "" else None + page_offset = int(qs.get("offset", ["0"])[0]) + if page_offset < 0 or (page_limit is not None and page_limit < 0): + raise ValueError("negative pagination") + except ValueError: + self._json(400, {"error": "Некорректные параметры limit/offset"}); return + + all_usage = _get_all_usage(month, department=department, status_filter=status_filter) + total_instances = len(all_usage) + if page_limit is not None: + paged_usage = all_usage[page_offset:page_offset + page_limit] + elif page_offset: + paged_usage = all_usage[page_offset:] + else: + paged_usage = all_usage + _audit("REPORT", "admin", detail=f"month={month}") self._json(200, {"month": month, - "usage": _get_all_usage(month), - "limits": _get_all_limits()}) + "usage": paged_usage, + "total_instances": total_instances, + "limits": _get_all_limits(), + "departments": _get_departments_summary(month)}) return self._proxy("GET") @@ -354,7 +736,7 @@ def do_POST(self): self._json(403, {"error": "Unauthorized"}); return body_raw = self._read_body() if body_raw is None: - self._json(413, {"error": "Request too large"}); return + return try: body = json.loads(body_raw) instance = str(body["instance"]).lower().strip() @@ -362,7 +744,7 @@ def do_POST(self): except (KeyError, ValueError, TypeError): self._json(400, {"error": "Нужны поля: instance (str), limit (int)"}); return # Валидация имени инстанса - if not instance.replace("-", "").replace("_", "").isalnum(): + if not _INSTANCE_RE.match(instance): self._json(400, {"error": "Недопустимое имя инстанса"}); return _set_limit(instance, limit, by="admin-api") _audit("SET_LIMIT", "admin", target=instance, detail=f"limit={limit}") @@ -372,6 +754,81 @@ def do_POST(self): "current_usage": current}) return + if path == "/quota/register-key": + if not _admin_rl.allow(self._client_ip()): + self._json(429, {"error": "Rate limit exceeded"}); return + if not self._is_admin(): + _audit("REGISTER_KEY_DENIED", self._client_ip()) + self._json(403, {"error": "Unauthorized"}); return + body_raw = self._read_body() + if body_raw is None: + return + try: + body = json.loads(body_raw) + instance = str(body["instance"]).lower().strip() + quota_key = body.get("quota_key") or body.get("key") + key_hash = body.get("key_hash") + department = str(body.get("department", "default")).lower().strip() or "default" + limit_val = int(body["limit"]) if body.get("limit") is not None else None + except (KeyError, ValueError, TypeError): + self._json(400, {"error": "Нужны поля: instance (str), quota_key или key_hash (str)"}); return + if not _INSTANCE_RE.match(instance): + self._json(400, {"error": "Недопустимое имя инстанса"}); return + if not _INSTANCE_RE.match(department): + self._json(400, {"error": "Недопустимое имя департамента"}); return + if not quota_key and not key_hash: + self._json(400, {"error": "Нужно поле quota_key или key_hash"}); return + if not quota_key and key_hash and not _KEY_HASH_RE.match(str(key_hash).strip()): + self._json(400, {"error": "Некорректный key_hash (ожидается 64 hex-символа SHA256)"}); return + try: + saved_hash = _register_key(instance, raw_key=str(quota_key) if quota_key else None, + key_hash=str(key_hash) if key_hash else None, + department=department, limit=limit_val, by="admin-api") + except ValueError as e: + self._json(400, {"error": str(e)}); return + _audit("REGISTER_KEY", "admin", target=instance, detail=f"dept={department} limit={limit_val}") + self._json(200, {"ok": True, "instance": instance, "key_hash": saved_hash, + "department": department, "limit": _get_limit(instance)}) + return + + if path == "/quota/reload-keys": + if not _admin_rl.allow(self._client_ip()): + self._json(429, {"error": "Rate limit exceeded"}); return + if not self._is_admin(): + _audit("RELOAD_KEYS_DENIED", self._client_ip()) + self._json(403, {"error": "Unauthorized"}); return + _load_keys() + with _keys_lock: + instances = sorted(set(_quota_keys.values())) + total_keys = len(_quota_keys) + _audit("RELOAD_KEYS", "admin", detail=f"total_keys={total_keys}") + self._json(200, {"ok": True, "instances": instances, "total_keys": total_keys}) + return + + if path == "/quota/set-department-limit": + if not _admin_rl.allow(self._client_ip()): + self._json(429, {"error": "Rate limit exceeded"}); return + if not self._is_admin(): + _audit("SET_DEPT_LIMIT_DENIED", self._client_ip()) + self._json(403, {"error": "Unauthorized"}); return + body_raw = self._read_body() + if body_raw is None: + return + try: + body = json.loads(body_raw) + department = str(body["department"]).lower().strip() + limit = int(body["limit"]) + except (KeyError, ValueError, TypeError): + self._json(400, {"error": "Нужны поля: department (str), limit (int)"}); return + if not _INSTANCE_RE.match(department): + self._json(400, {"error": "Недопустимое имя департамента"}); return + _set_department_limit(department, limit, by="admin-api") + _audit("SET_DEPT_LIMIT", "admin", target=department, detail=f"limit={limit}") + current = _get_department_usage(department) + self._json(200, {"ok": True, "department": department, "limit": limit, + "current_usage": current}) + return + if path == "/quota/reset": if not _admin_rl.allow(self._client_ip()): self._json(429, {"error": "Rate limit exceeded"}); return @@ -380,7 +837,7 @@ def do_POST(self): self._json(403, {"error": "Unauthorized"}); return body_raw = self._read_body() if body_raw is None: - self._json(413, {"error": "Request too large"}); return + return try: body = json.loads(body_raw) instance = str(body["instance"]).lower().strip() @@ -405,72 +862,75 @@ def _proxy(self, method: str): if not _proxy_rl.allow(instance): self._json(429, {"error": "Rate limit exceeded for instance"}); return - ok, msg = _check_quota(instance) - if not ok: - _audit("QUOTA_BLOCKED", instance) - self._json(429, {"type": "error", "error": "quota_exceeded", "message": msg}); return - body_raw = self._read_body() if body_raw is None: - self._json(413, {"error": "Request body too large (max 1MB)"}); return - - headers = { - "anthropic-version": self.headers.get("anthropic-version", "2023-06-01"), - "content-type": self.headers.get("content-type", "application/json"), - } + return - if _is_cliproxy: - # CLIProxyAPI: всегда x-api-key, без OAuth-бет, с User-Agent для обхода cloaking - headers["x-api-key"] = _upstream_key - headers["User-Agent"] = "claude-cli/2.1.44 (external, sdk-cli)" - if v := self.headers.get("anthropic-beta"): - headers["anthropic-beta"] = v - elif _upstream_key.startswith("sk-ant-oat"): - # Прямой Anthropic с OAuth-ключом - headers["Authorization"] = f"Bearer {_upstream_key}" - oauth_betas = ["claude-code-20250219", "oauth-2025-04-20"] - existing = self.headers.get("anthropic-beta", "") - existing_list = [b.strip() for b in existing.split(",") if b.strip()] - all_betas = list(dict.fromkeys(oauth_betas + existing_list)) - headers["anthropic-beta"] = ",".join(all_betas) - else: - # Прямой Anthropic с обычным API-ключом - headers["x-api-key"] = _upstream_key - if v := self.headers.get("anthropic-beta"): - headers["anthropic-beta"] = v + with _quota_admission_guard(instance): + ok, msg = _check_quota(instance) + if not ok: + _audit("QUOTA_BLOCKED", instance) + self._json(429, {"type": "error", "error": "quota_exceeded", "message": msg}); return + + headers = { + "anthropic-version": self.headers.get("anthropic-version", "2023-06-01"), + "content-type": self.headers.get("content-type", "application/json"), + } + + if _is_cliproxy: + # CLIProxyAPI: всегда x-api-key, без OAuth-бет, с User-Agent для обхода cloaking + headers["x-api-key"] = _upstream_key + headers["User-Agent"] = "claude-cli/2.1.44 (external, sdk-cli)" + if v := self.headers.get("anthropic-beta"): + headers["anthropic-beta"] = v + elif _upstream_key.startswith("sk-ant-oat"): + # Прямой Anthropic с OAuth-ключом + headers["Authorization"] = f"Bearer {_upstream_key}" + oauth_betas = ["claude-code-20250219", "oauth-2025-04-20"] + existing = self.headers.get("anthropic-beta", "") + existing_list = [b.strip() for b in existing.split(",") if b.strip()] + all_betas = list(dict.fromkeys(oauth_betas + existing_list)) + headers["anthropic-beta"] = ",".join(all_betas) + else: + # Прямой Anthropic с обычным API-ключом + headers["x-api-key"] = _upstream_key + if v := self.headers.get("anthropic-beta"): + headers["anthropic-beta"] = v - try: - req = Request(UPSTREAM_URL + self.path, data=body_raw or None, - headers=headers, method=method) - with urlopen(req, timeout=120) as r: - resp_body, resp_status = r.read(), r.status - resp_ct = r.headers.get("Content-Type", "application/json") - except HTTPError as e: - resp_body, resp_status, resp_ct = e.read(), e.code, e.headers.get("Content-Type", "application/json") - except URLError as e: - print(f"[quota] UPSTREAM ERROR {instance}: {e}", flush=True) - self._json(502, {"type": "error", "error": "upstream_unavailable", - "message": f"Anthropic API недоступен: {e.reason}"}); return - except Exception as e: - print(f"[quota] UNEXPECTED ERROR {instance}: {e}", flush=True) - self._json(502, {"type": "error", "error": "proxy_error", - "message": "Внутренняя ошибка proxy"}); return - - # Считаем токены из ответа - try: - usage = json.loads(resp_body).get("usage", {}) - inp, out = usage.get("input_tokens", 0), usage.get("output_tokens", 0) - if inp or out: - _add_usage(instance, inp, out) - total = _get_usage(instance)["total"] - limit = _get_limit(instance) - pct = f"{total/limit*100:.0f}%" if limit > 0 else "∞" - lvl = "⚠️" if limit>0 and total/limit>=0.8 else " " - print(f"[quota] {lvl} {instance}: +{inp+out} → {total:,}/{limit:,} ({pct})", flush=True) - except Exception as e: - # Streaming responses (text/event-stream) are not JSON — skip silently - if "event-stream" not in resp_ct: - print(f"[quota] WARNING: failed to parse usage from response for {instance}: {e}", flush=True) + try: + req = Request(UPSTREAM_URL + self.path, data=body_raw or None, + headers=headers, method=method) + with urlopen(req, timeout=120) as r: + resp_status = r.status + resp_ct = r.headers.get("Content-Type", "application/json") + if "event-stream" in resp_ct.lower(): + self._stream_sse(r, resp_status, resp_ct, instance) + return + resp_body = r.read() + except HTTPError as e: + resp_body, resp_status, resp_ct = e.read(), e.code, e.headers.get("Content-Type", "application/json") + except URLError as e: + print(f"[quota] UPSTREAM ERROR {instance}: {e}", flush=True) + self._json(502, {"type": "error", "error": "upstream_unavailable", + "message": f"Anthropic API недоступен: {e.reason}"}); return + except Exception as e: + print(f"[quota] UNEXPECTED ERROR {instance}: {e}", flush=True) + self._json(502, {"type": "error", "error": "proxy_error", + "message": "Внутренняя ошибка proxy"}); return + + # Считаем токены из ответа (JSON или буферизованный text/event-stream) + try: + if "event-stream" in resp_ct.lower(): + inp, out = _parse_sse_usage(resp_body) + else: + usage = json.loads(resp_body).get("usage", {}) + inp = usage.get("input_tokens") if usage.get("input_tokens") is not None else usage.get("prompt_tokens", 0) + out = usage.get("output_tokens") if usage.get("output_tokens") is not None else usage.get("completion_tokens", 0) + if inp or out: + self._record_usage_log(instance, int(inp), int(out)) + except Exception as e: + if "event-stream" not in resp_ct.lower(): + print(f"[quota] WARNING: failed to parse usage from response for {instance}: {e}", flush=True) self.send_response(resp_status) self.send_header("Content-Type", resp_ct) @@ -480,9 +940,44 @@ def _proxy(self, method: str): self.end_headers() self.wfile.write(resp_body) + def _record_usage_log(self, instance: str, inp: int, out: int): + _add_usage(instance, inp, out) + total = _get_usage(instance)["total"] + limit = _get_limit(instance) + pct = f"{total/limit*100:.0f}%" if limit > 0 else "∞" + lvl = "⚠️" if limit>0 and total/limit>=0.8 else " " + print(f"[quota] {lvl} {instance}: +{inp+out} → {total:,}/{limit:,} ({pct})", flush=True) + + def _stream_sse(self, upstream_resp, resp_status: int, resp_ct: str, instance: str): + """Потоковая передача SSE клиенту без буферизации с подсчётом токенов на лету.""" + parser = SSEUsageParser() + try: + self.send_response(resp_status) + self.send_header("Content-Type", resp_ct) + for k, v in _SEC_HEADERS.items(): + self.send_header(k, v) + self.end_headers() + while True: + line = upstream_resp.readline(MAX_SSE_LINE_BYTES) + if not line: + break + parser.feed_chunk(line) + try: + self.wfile.write(line) + self.wfile.flush() + except (BrokenPipeError, ConnectionResetError): + break + parser.finish() + finally: + self.close_connection = True + inp, out = parser.input_tokens, parser.output_tokens + if inp or out: + self._record_usage_log(instance, int(inp), int(out)) + if __name__ == "__main__": - instances = sorted(set(_quota_keys.values())) + with _keys_lock: + instances = sorted(set(_quota_keys.values())) upstream_mode = "CLIProxyAPI" if _is_cliproxy else "Anthropic API (direct)" print(f"✅ Quota-proxy hardened :{PORT}", flush=True) print(f" Upstream: {UPSTREAM_URL} ({upstream_mode})", flush=True) @@ -490,7 +985,7 @@ def _proxy(self, method: str): print(f" DB: {DB_PATH} (WAL mode)", flush=True) print(f" Max body: {MAX_BODY_BYTES//1024}KB", flush=True) - server = HTTPServer(("0.0.0.0", PORT), Handler) + server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler) def _shutdown(signum, frame): print(f"[quota] Получен сигнал {signum}, завершаем...", flush=True) diff --git a/resource-monitor/monitor.py b/resource-monitor/monitor.py index 171ec98..6806823 100644 --- a/resource-monitor/monitor.py +++ b/resource-monitor/monitor.py @@ -17,13 +17,14 @@ import sqlite3 import shutil import threading +import concurrent.futures import hmac import hashlib import logging -from http.server import HTTPServer, BaseHTTPRequestHandler +from http.server import HTTPServer, ThreadingHTTPServer, BaseHTTPRequestHandler from http.client import HTTPConnection from datetime import datetime, timezone -from urllib.parse import urlparse +from urllib.parse import urlparse, parse_qs # ── Configuration ──────────────────────────────────────────────────────────── @@ -217,9 +218,17 @@ def get_quota_report() -> dict: # ── SQLite ──────────────────────────────────────────────────────────────────── +def _get_db_conn() -> sqlite3.Connection: + conn = sqlite3.connect(MONITOR_DB, timeout=10.0) + conn.execute("PRAGMA busy_timeout=10000") + return conn + + def init_db(): os.makedirs(os.path.dirname(MONITOR_DB) or ".", exist_ok=True) - conn = sqlite3.connect(MONITOR_DB) + conn = _get_db_conn() + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA journal_size_limit=1048576") conn.execute(""" CREATE TABLE IF NOT EXISTS snapshots ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -230,6 +239,7 @@ def init_db(): ram_limit_mb REAL NOT NULL ) """) + conn.execute("CREATE INDEX IF NOT EXISTS idx_snapshots_container_ts ON snapshots (container, timestamp)") conn.commit() conn.close() @@ -238,7 +248,7 @@ def write_snapshots(metrics: list): """Write a list of metric dicts to SQLite.""" if not metrics: return - conn = sqlite3.connect(MONITOR_DB) + conn = _get_db_conn() try: conn.executemany( "INSERT INTO snapshots (timestamp, container, cpu_pct, ram_mb, ram_limit_mb) " @@ -261,14 +271,27 @@ def write_snapshots(metrics: list): # ── Alerts ──────────────────────────────────────────────────────────────────── +def _extract_instance_from_key(key: str) -> str: + """Extract instance or container name from alert key (e.g., 'cpu:corp-user-1' -> 'corp-user-1').""" + if ":" in key: + return key.split(":", 1)[1] + return key + + def send_alert(key: str, message: str): """Send alert to broker if cooldown has passed.""" now = time.time() + instance = _extract_instance_from_key(key) with _lock: last = _active_alerts.get(key, {}).get("fired_at", 0) if now - last < ALERT_COOLDOWN: return # still in cooldown - _active_alerts[key] = {"fired_at": now, "message": message, "key": key} + _active_alerts[key] = { + "fired_at": now, + "message": message, + "key": key, + "instance": instance, + } log.warning(f"ALERT [{key}]: {message}") if not BROKER_KEY_MONITOR: @@ -342,25 +365,35 @@ def poll_once(): metrics = [] containers = list_containers() - for c in containers: + + def _fetch_one(c: dict) -> dict | None: try: stats = get_container_stats(c["id"]) cpu_pct = parse_cpu_pct(stats) ram_mb, ram_limit_mb = parse_ram(stats) blkio = parse_blkio(stats) - m = { + mem_pct = (ram_mb / ram_limit_mb * 100.0) if ram_limit_mb > 0 else 0.0 + return { "timestamp": ts, "id": c["short_id"], "name": c["name"], "cpu_pct": round(cpu_pct, 2), "ram_mb": round(ram_mb, 2), "ram_limit_mb": round(ram_limit_mb, 2), + "mem_pct": round(mem_pct, 2), "blkio": blkio, } - metrics.append(m) except Exception as e: log.warning(f"Stats for {c['name']} failed: {e}") + return None + + if containers: + max_workers = min(32, len(containers) or 1) + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: + for m in pool.map(_fetch_one, containers): + if m is not None: + metrics.append(m) disk = get_disk_usage() quota = get_quota_report() @@ -419,13 +452,66 @@ def send_json(handler, data: dict, status: int = 200): handler.wfile.write(body) +def _container_mem_pct(c: dict) -> float: + """Compute or retrieve memory usage percentage for a container metric dict.""" + if "mem_pct" in c and c["mem_pct"] is not None: + return float(c["mem_pct"]) + ram_mb = float(c.get("ram_mb", c.get("mem_mb", 0.0)) or 0.0) + limit_mb = float(c.get("ram_limit_mb", c.get("mem_limit_mb", 0.0)) or 0.0) + if limit_mb > 0: + return round(ram_mb / limit_mb * 100.0, 2) + return 0.0 + + +def _is_container_unhealthy(c: dict) -> bool: + """Determine if a container is unhealthy by status flag or resource threshold.""" + if c.get("healthy") is False or c.get("unhealthy") is True: + return True + if c.get("status") in ("unhealthy", "dead", "exited", "restarting"): + return True + if c.get("health") in ("unhealthy", "failing"): + return True + cpu_pct = float(c.get("cpu_pct", 0.0) or 0.0) + mem_pct = _container_mem_pct(c) + if cpu_pct > ALERT_CPU_PCT or mem_pct > ALERT_RAM_PCT: + return True + return False + + +def _alert_matches_instance(alert: dict, instance: str) -> bool: + """Check whether an alert dictionary matches the requested instance filter.""" + target = instance.strip() + if not target: + return True + candidates = {target, f"corp-{target}"} + if target.startswith("corp-"): + candidates.add(target[5:]) + + alert_inst = str(alert.get("instance", "") or alert.get("container", "")) + if alert_inst in candidates: + return True + + key = str(alert.get("key", "")) + if ":" in key and key.split(":", 1)[1] in candidates: + return True + if any(cand in key for cand in candidates): + return True + + msg = str(alert.get("message", "")) + if any(cand in msg for cand in candidates): + return True + return False + + class MonitorHandler(BaseHTTPRequestHandler): def log_message(self, fmt, *args): log.debug(f"HTTP {fmt % args}") def do_GET(self): - path = self.path.split("?")[0] + parsed = urlparse(self.path) + path = parsed.path + qs = parse_qs(parsed.query) # /health — no auth required if path == "/health": @@ -436,13 +522,86 @@ def do_GET(self): return if path == "/metrics": + search_q = (qs.get("search", [None])[0] or "").strip().lower() + sort_by = (qs.get("sort", [None])[0] or "").strip().lower() + order_q = (qs.get("order", [None])[0] or "").strip().lower() + raw_limit = qs.get("limit", [None])[0] + raw_offset = qs.get("offset", ["0"])[0] + + try: + limit_val = int(raw_limit) if raw_limit is not None else None + offset_val = max(0, int(raw_offset)) if raw_offset is not None else 0 + except ValueError: + send_json(self, {"error": "invalid limit or offset"}, 400) + return + with _lock: - payload = { - "containers": list(_last_metrics.values()), - "disk": dict(_last_disk), - "quota": dict(_last_quota), - "polled_at": datetime.now(timezone.utc).isoformat(), - } + raw_containers = [dict(c) for c in _last_metrics.values()] + disk_copy = dict(_last_disk) + quota_copy = dict(_last_quota) + + # Ensure mem_pct is populated on each container entry + for c in raw_containers: + if "mem_pct" not in c: + c["mem_pct"] = _container_mem_pct(c) + + # Filter by substring in container name + if search_q: + filtered = [ + c for c in raw_containers + if search_q in str(c.get("name", "")).lower() + ] + else: + filtered = raw_containers + + # Compute summary across all matched containers (pre-pagination) + total_containers = len(filtered) + total_mem_mb = round( + sum(float(c.get("ram_mb", c.get("mem_mb", 0.0)) or 0.0) for c in filtered), + 2, + ) + avg_cpu_pct = ( + round( + sum(float(c.get("cpu_pct", 0.0) or 0.0) for c in filtered) / total_containers, + 2, + ) + if total_containers > 0 + else 0.0 + ) + unhealthy_count = sum(1 for c in filtered if _is_container_unhealthy(c)) + + # Sort containers if requested + if sort_by == "mem_pct": + reverse = False if order_q == "asc" else True + filtered.sort(key=lambda c: (_container_mem_pct(c), c.get("name", "")), reverse=reverse) + elif sort_by == "cpu_pct": + reverse = False if order_q == "asc" else True + filtered.sort(key=lambda c: (float(c.get("cpu_pct", 0.0) or 0.0), c.get("name", "")), reverse=reverse) + elif sort_by == "name": + reverse = True if order_q == "desc" else False + filtered.sort(key=lambda c: str(c.get("name", "")), reverse=reverse) + + # Apply pagination + if limit_val is not None: + safe_limit = max(1, limit_val) + paged_containers = filtered[offset_val:offset_val + safe_limit] + elif offset_val > 0: + paged_containers = filtered[offset_val:] + else: + paged_containers = filtered + + payload = { + "containers": paged_containers, + "total_containers": total_containers, + "summary": { + "total_mem_mb": total_mem_mb, + "avg_cpu_pct": avg_cpu_pct, + "unhealthy_count": unhealthy_count, + }, + "disk": disk_copy, + "quota": quota_copy, + "polled_at": datetime.now(timezone.utc).isoformat(), + } send_json(self, payload) elif path.startswith("/metrics/"): @@ -454,10 +613,43 @@ def do_GET(self): else: send_json(self, m) - elif path == "/alerts/active": + elif path in ("/alerts", "/alerts/active"): + instance_filter = (qs.get("instance", [None])[0] or "").strip() + raw_limit = qs.get("limit", [None])[0] + raw_offset = qs.get("offset", ["0"])[0] + + try: + limit_val = int(raw_limit) if raw_limit is not None else None + offset_val = max(0, int(raw_offset)) if raw_offset is not None else 0 + except ValueError: + send_json(self, {"error": "invalid limit or offset"}, 400) + return + with _lock: - alerts = list(_active_alerts.values()) - send_json(self, {"alerts": alerts, "count": len(alerts)}) + all_alerts = [dict(a) for a in _active_alerts.values()] + + if instance_filter: + filtered_alerts = [ + a for a in all_alerts + if _alert_matches_instance(a, instance_filter) + ] + else: + filtered_alerts = all_alerts + + total_alerts = len(filtered_alerts) + if limit_val is not None: + safe_limit = max(1, limit_val) + paged_alerts = filtered_alerts[offset_val:offset_val + safe_limit] + elif offset_val > 0: + paged_alerts = filtered_alerts[offset_val:] + else: + paged_alerts = filtered_alerts + + send_json(self, { + "alerts": paged_alerts, + "count": len(paged_alerts), + "total": total_alerts, + }) else: send_json(self, {"error": "not found"}, 404) @@ -473,7 +665,7 @@ def main(): t = threading.Thread(target=polling_loop, daemon=True, name="poller") t.start() - server = HTTPServer(("0.0.0.0", MONITOR_PORT), MonitorHandler) + server = ThreadingHTTPServer(("0.0.0.0", MONITOR_PORT), MonitorHandler) log.info(f"HTTP API listening on 0.0.0.0:{MONITOR_PORT}") try: server.serve_forever() diff --git a/scripts/add-user.sh b/scripts/add-user.sh index 96c32a9..79c26d7 100644 --- a/scripts/add-user.sh +++ b/scripts/add-user.sh @@ -41,9 +41,21 @@ fi TEMPLATE="instances/_template" TARGET="instances/$NAME" +NAME_UPPER=$(echo "$NAME" | tr '[:lower:]-' '[:upper:]_') [ -d "$TARGET" ] && echo "❌ Инстанс '$NAME' уже существует" && exit 1 +for existing_dir in instances/*/; do + [ -d "$existing_dir" ] || continue + existing_name=$(basename "$existing_dir") + [ "$existing_name" = "_template" ] || [ "$existing_name" = "admin" ] && continue + existing_upper=$(echo "$existing_name" | tr '[:lower:]-' '[:upper:]_') + if [ "$existing_upper" = "$NAME_UPPER" ] && [ "$existing_name" != "$NAME" ]; then + echo "❌ Коллизия имени '$NAME' с существующим инстансом '$existing_name' (общий env-суффикс $NAME_UPPER)" + exit 1 + fi +done + if [ ! -d "$TEMPLATE/workspace" ]; then echo "❌ Директория шаблона не найдена: $TEMPLATE/workspace" exit 1 @@ -96,7 +108,7 @@ cp "$INSTANCE_WORKSPACE/SOUL.md" "$TARGET/hermes_data/SOUL.md" BROKER_KEY=$(python3 -c "import secrets; print(secrets.token_urlsafe(32))") QUOTA_KEY=$(python3 -c "import secrets; print('quota-${NAME}-' + secrets.token_urlsafe(24))") GATEWAY_TOKEN=$(python3 -c "import secrets; print(secrets.token_urlsafe(32))") -NAME_UPPER=$(echo "$NAME" | tr '[:lower:]' '[:upper:]') +NAME_UPPER=$(echo "$NAME" | tr '[:lower:]-' '[:upper:]_') # .env инстанса (compose env_file) — ключи для скиллов и брокера cat > "$TARGET/.env" << EOF diff --git a/scripts/benchmark-1000-users.py b/scripts/benchmark-1000-users.py new file mode 100755 index 0000000..8af0277 --- /dev/null +++ b/scripts/benchmark-1000-users.py @@ -0,0 +1,612 @@ +#!/usr/bin/env python3 +""" +ENT1K-3: 1000-User Concurrency & Scale Benchmark CLI for CortexForge. + +Spins up an isolated temporary environment with: + - Mock LLM Upstream (JSON + Anthropic SSE + OpenAI SSE) + - Live quota-proxy (ThreadingHTTPServer + SQLite WAL) + - Live message-broker (ThreadingHTTPServer + SQLite WAL) + +Registers 1,000 users/agents across 10 departments (dept-0..dept-9), executes high-concurrency +workloads across a thread pool, measures throughput (RPS) and latency percentiles +(p50, p95, p99 in ms), prints a formatted benchmark table, and optionally writes +a structured JSON report via --output. + +Python 3.12+ stdlib only. +""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import importlib.util +import io +import json +import math +import os +import pathlib +import shutil +import sqlite3 +import sys +import tempfile +import threading +import time +import urllib.error +import urllib.request +from typing import Any + +REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent +PROXY_PATH = REPO_ROOT / "quota-proxy" / "proxy.py" +BROKER_DIR = REPO_ROOT / "broker" + + +class HighConcurrencyHTTPServer(ThreadingHTTPServer): + """ThreadingHTTPServer tuned for high-concurrency benchmark bursts.""" + + request_queue_size = 512 + allow_reuse_address = True + daemon_threads = True + + +class BenchmarkUpstreamHandler(BaseHTTPRequestHandler): + """Fast local mock upstream supporting JSON, Anthropic SSE, and OpenAI SSE.""" + + def log_message(self, fmt, *args): + pass + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + raw = self.rfile.read(length) if length else b"{}" + try: + payload = json.loads(raw) + except Exception: + payload = {} + + mode = payload.get("mock_mode", "json") + in_tok = int(payload.get("in_tok", 20)) + out_tok = int(payload.get("out_tok", 30)) + + if mode == "sse_anthropic": + data = ( + "event: message_start\n" + f'data: {{"type":"message_start","message":{{"id":"msg_b","usage":{{"input_tokens":{in_tok},"output_tokens":1}}}}}}\n\n' + "event: content_block_delta\n" + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}\n\n' + "event: message_delta\n" + f'data: {{"type":"message_delta","delta":{{"stop_reason":"end_turn"}},"usage":{{"output_tokens":{out_tok}}}}}\n\n' + "event: message_stop\n" + 'data: {"type":"message_stop"}\n\n' + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.end_headers() + self.wfile.write(data) + self.wfile.flush() + return + + if mode == "sse_openai": + data = ( + 'data: {"id":"chatcmpl-b","choices":[{"index":0,"delta":{"content":"ok"}}]}\n\n' + f'data: {{"id":"chatcmpl-b","choices":[],"usage":{{"prompt_tokens":{in_tok},"completion_tokens":{out_tok},"total_tokens":{in_tok + out_tok}}}}}\n\n' + "data: [DONE]\n\n" + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + self.wfile.write(data) + self.wfile.flush() + return + + body = json.dumps( + { + "id": "msg_bench", + "type": "message", + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + } + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +def _percentile(sorted_vals: list[float], pct: float) -> float: + if not sorted_vals: + return 0.0 + if len(sorted_vals) == 1: + return round(sorted_vals[0], 2) + rank = (pct / 100.0) * (len(sorted_vals) - 1) + low = int(math.floor(rank)) + high = int(math.ceil(rank)) + if low == high: + return round(sorted_vals[low], 2) + weight = rank - low + return round(sorted_vals[low] * (1.0 - weight) + sorted_vals[high] * weight, 2) + + +def _summarize_stage( + name: str, + latencies_ms: list[float], + elapsed_sec: float, + errors: int, + logical_ops: int | None = None, +) -> dict[str, Any]: + s = sorted(latencies_ms) + ops = logical_ops if logical_ops is not None else len(latencies_ms) + rps = round(ops / elapsed_sec, 1) if elapsed_sec > 0 else 0.0 + return { + "stage": name, + "operations": ops, + "requests": len(latencies_ms), + "elapsed_sec": round(elapsed_sec, 3), + "rps": rps, + "p50_ms": _percentile(s, 50.0), + "p95_ms": _percentile(s, 95.0), + "p99_ms": _percentile(s, 99.0), + "errors": errors, + "status": "PASS" if errors == 0 else "FAIL", + } + + +@contextmanager +def _suppress_service_stdout(): + """Suppress verbose per-request service prints during benchmark execution.""" + saved_stdout = sys.stdout + sys.stdout = io.StringIO() + try: + yield + finally: + sys.stdout = saved_stdout + + +def _load_module(name: str, path: pathlib.Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +def run_benchmark(num_users: int = 1000, workers: int = 64) -> dict[str, Any]: + saved_env = dict(os.environ) + tmpdir = tempfile.mkdtemp(prefix="cortex_bench_1000_") + quota_db = os.path.join(tmpdir, "quota_bench.db") + broker_db = os.path.join(tmpdir, "broker_bench.db") + admin_token = "adm-bench-token-0123456789abcdef" + broker_admin_key = "adm-broker-bench-secret" + + upstream_srv = HighConcurrencyHTTPServer(("127.0.0.1", 0), BenchmarkUpstreamHandler) + upstream_port = upstream_srv.server_address[1] + threading.Thread(target=upstream_srv.serve_forever, daemon=True).start() + + proxy_srv = None + broker_srv = None + proxy_mod = None + broker_mod = None + + stages: list[dict[str, Any]] = [] + t_suite_start = time.monotonic() + + try: + with _suppress_service_stdout(): + # 1. Start quota-proxy + os.environ["QUOTA_DB"] = quota_db + os.environ["UPSTREAM_URL"] = f"http://127.0.0.1:{upstream_port}" + os.environ["UPSTREAM_API_KEY"] = "clip-bench-upstream-key" + os.environ["QUOTA_ADMIN_TOKEN"] = admin_token + os.environ["QUOTA_DEFAULT_MONTHLY"] = "1000000" + + proxy_mod = _load_module("quota_proxy_bench", PROXY_PATH) + proxy_mod._admin_rl = proxy_mod.RateLimiter(rate=100_000.0, capacity=100_000) + proxy_mod._proxy_rl = proxy_mod.RateLimiter(rate=100_000.0, capacity=100_000) + proxy_mod.Handler.log_message = lambda self, fmt, *a: None + + proxy_srv = HighConcurrencyHTTPServer(("127.0.0.1", 0), proxy_mod.Handler) + proxy_port = proxy_srv.server_address[1] + threading.Thread(target=proxy_srv.serve_forever, daemon=True).start() + + # 2. Start message-broker + if str(BROKER_DIR) not in sys.path: + sys.path.insert(0, str(BROKER_DIR)) + import broker as _broker_mod + + broker_mod = _broker_mod + broker_mod.BROKER_DB = broker_db + broker_mod.MAX_AGE_SEC = 86400 + broker_mod.RETENTION_SECONDS = 86400 + broker_mod._min_active_ttl = 86400.0 + broker_mod._last_purge_ts = 0.0 + os.environ["BROKER_DB"] = broker_db + os.environ["BROKER_KEY_ADMIN"] = broker_admin_key + broker_mod._init_db() + broker_mod._load_keys() + broker_mod._send_rl = broker_mod.RateLimiter(rate=100_000.0, capacity=100_000) + broker_mod.Handler.log_message = lambda self, fmt, *a: None + + broker_srv = HighConcurrencyHTTPServer(("127.0.0.1", 0), broker_mod.Handler) + broker_port = broker_srv.server_address[1] + threading.Thread(target=broker_srv.serve_forever, daemon=True).start() + + proxy_base = f"http://127.0.0.1:{proxy_port}" + broker_base = f"http://127.0.0.1:{broker_port}" + + # ── Stage 1: Quota-Proxy Key Registration (1,000 users) ────── + def qp_register(i: int) -> tuple[float, bool]: + t0 = time.monotonic() + payload = json.dumps( + { + "instance": f"user-{i}", + "quota_key": f"quota-user-{i}-bench", + "department": f"dept-{(i - 1) % 10}", + "limit": 1_000_000, + } + ).encode("utf-8") + req = urllib.request.Request( + f"{proxy_base}/quota/register-key", + data=payload, + headers={ + "Authorization": f"Bearer {admin_token}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + _ = resp.read() + ok = resp.status == 200 + except Exception: + ok = False + return (time.monotonic() - t0) * 1000.0, ok + + t_stage = time.monotonic() + with ThreadPoolExecutor(max_workers=workers) as pool: + res_qp_reg = list(pool.map(qp_register, range(1, num_users + 1))) + el_qp_reg = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"quota-proxy: Register {num_users} users", + [r[0] for r in res_qp_reg], + el_qp_reg, + sum(1 for r in res_qp_reg if not r[1]), + ) + ) + + # ── Stage 2: Quota-Proxy Concurrent Requests (JSON + SSE) ──── + def qp_call(i: int) -> tuple[float, bool]: + t0 = time.monotonic() + if i % 3 == 1: + mode = "json" + elif i % 3 == 2: + mode = "sse_anthropic" + else: + mode = "sse_openai" + payload = json.dumps( + {"mock_mode": mode, "in_tok": 25, "out_tok": 25} + ).encode("utf-8") + req = urllib.request.Request( + f"{proxy_base}/v1/messages", + data=payload, + headers={ + "x-api-key": f"quota-user-{i}-bench", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + _ = resp.read() + ok = resp.status == 200 + except Exception: + ok = False + return (time.monotonic() - t0) * 1000.0, ok + + t_stage = time.monotonic() + with ThreadPoolExecutor(max_workers=workers) as pool: + res_qp_call = list(pool.map(qp_call, range(1, num_users + 1))) + el_qp_call = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"quota-proxy: {num_users} concurrent requests (JSON+SSE)", + [r[0] for r in res_qp_call], + el_qp_call, + sum(1 for r in res_qp_call if not r[1]), + ) + ) + + # ── Stage 3: Quota-Proxy Paginated Report (10 pages) ───────── + num_pages = max(1, math.ceil(num_users / 100)) + + def qp_report_page(page_idx: int) -> tuple[float, bool]: + t0 = time.monotonic() + offset = page_idx * 100 + req = urllib.request.Request( + f"{proxy_base}/quota/report?limit=100&offset={offset}", + headers={"Authorization": f"Bearer {admin_token}"}, + method="GET", + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + data = json.loads(resp.read().decode("utf-8")) + ok = resp.status == 200 and data.get("total_instances") == num_users + except Exception: + ok = False + return (time.monotonic() - t0) * 1000.0, ok + + t_stage = time.monotonic() + with ThreadPoolExecutor(max_workers=min(workers, num_pages)) as pool: + res_qp_rep = list(pool.map(qp_report_page, range(num_pages))) + el_qp_rep = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"quota-proxy: Paginated report ({num_pages} pages)", + [r[0] for r in res_qp_rep], + el_qp_rep, + sum(1 for r in res_qp_rep if not r[1]), + ) + ) + + # ── Stage 4: Message-Broker Register 1,000 Agents ──────────── + def brk_register(i: int) -> tuple[float, bool]: + t0 = time.monotonic() + payload = json.dumps( + {"instance": f"user-{i}", "broker_key": f"brk-user-{i}-bench"} + ).encode("utf-8") + req = urllib.request.Request( + f"{broker_base}/admin/register-key", + data=payload, + headers={ + "Authorization": f"Bearer {broker_admin_key}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + _ = resp.read() + ok = resp.status == 200 + except Exception: + ok = False + return (time.monotonic() - t0) * 1000.0, ok + + t_stage = time.monotonic() + with ThreadPoolExecutor(max_workers=workers) as pool: + res_brk_reg = list(pool.map(brk_register, range(1, num_users + 1))) + el_brk_reg = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"message-broker: Register {num_users} agents", + [r[0] for r in res_brk_reg], + el_brk_reg, + sum(1 for r in res_brk_reg if not r[1]), + ) + ) + + # ── Stage 5: Message-Broker 1,000 Concurrent P2P Messages ──── + priorities = ("low", "normal", "high", "urgent") + + def brk_send(i: int) -> tuple[float, bool]: + t0 = time.monotonic() + target = f"user-{(i % num_users) + 1}" + payload = json.dumps( + { + "to": target, + "message": f"bench-msg-{i}", + "priority": priorities[i % 4], + } + ).encode("utf-8") + req = urllib.request.Request( + f"{broker_base}/send", + data=payload, + headers={ + "Authorization": f"Bearer brk-user-{i}-bench", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + _ = resp.read() + ok = resp.status == 200 + except Exception: + ok = False + return (time.monotonic() - t0) * 1000.0, ok + + t_stage = time.monotonic() + with ThreadPoolExecutor(max_workers=workers) as pool: + res_brk_send = list(pool.map(brk_send, range(1, num_users + 1))) + el_brk_send = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"message-broker: {num_users} concurrent P2P messages", + [r[0] for r in res_brk_send], + el_brk_send, + sum(1 for r in res_brk_send if not r[1]), + ) + ) + + # ── Stage 6: Message-Broker Admin Broadcast (*) to 1,000 ───── + t_stage = time.monotonic() + bcast_payload = json.dumps( + { + "to": "*", + "text": "Enterprise-wide benchmark broadcast", + "priority": "urgent", + } + ).encode("utf-8") + bcast_req = urllib.request.Request( + f"{broker_base}/broadcast", + data=bcast_payload, + headers={ + "Authorization": f"Bearer {broker_admin_key}", + "Content-Type": "application/json", + }, + method="POST", + ) + try: + with urllib.request.urlopen(bcast_req, timeout=15) as resp: + bcast_body = json.loads(resp.read().decode("utf-8")) + bcast_ok = ( + resp.status == 200 + and bcast_body.get("delivered_count") == num_users + ) + except Exception: + bcast_ok = False + el_bcast = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"message-broker: Broadcast (*) to {num_users} agents", + [el_bcast * 1000.0], + el_bcast, + 0 if bcast_ok else 1, + logical_ops=num_users, + ) + ) + + # ── Stage 7: Message-Broker Paginated /inbox/all ───────────── + def brk_inbox_page(page_idx: int) -> tuple[float, bool]: + t0 = time.monotonic() + offset = page_idx * 100 + req = urllib.request.Request( + f"{broker_base}/inbox/all?limit=100&offset={offset}", + headers={"Authorization": f"Bearer {broker_admin_key}"}, + method="GET", + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + data = json.loads(resp.read().decode("utf-8")) + ok = resp.status == 200 and data.get("total_inboxes") == num_users + except Exception: + ok = False + return (time.monotonic() - t0) * 1000.0, ok + + t_stage = time.monotonic() + with ThreadPoolExecutor(max_workers=min(workers, num_pages)) as pool: + res_brk_all = list(pool.map(brk_inbox_page, range(num_pages))) + el_brk_all = time.monotonic() - t_stage + stages.append( + _summarize_stage( + f"message-broker: Paginated /inbox/all ({num_pages} pages)", + [r[0] for r in res_brk_all], + el_brk_all, + sum(1 for r in res_brk_all if not r[1]), + ) + ) + + total_elapsed = round(time.monotonic() - t_suite_start, 3) + total_ops = sum(s["operations"] for s in stages) + total_errors = sum(s["errors"] for s in stages) + overall_rps = round(total_ops / total_elapsed, 1) if total_elapsed > 0 else 0.0 + + return { + "benchmark": "CortexForge ENT1K-3 1000-User Scale Benchmark", + "users": num_users, + "workers": workers, + "total_operations": total_ops, + "total_elapsed_sec": total_elapsed, + "overall_rps": overall_rps, + "total_errors": total_errors, + "status": "PASS" if total_errors == 0 else "FAIL", + "stages": stages, + } + finally: + if proxy_srv: + proxy_srv.shutdown() + proxy_srv.server_close() + if broker_srv: + broker_srv.shutdown() + broker_srv.server_close() + upstream_srv.shutdown() + upstream_srv.server_close() + if proxy_mod and getattr(proxy_mod, "_conn", None): + try: + proxy_mod._conn.close() + except Exception: + pass + os.environ.clear() + os.environ.update(saved_env) + shutil.rmtree(tmpdir, ignore_errors=True) + + +def print_report_table(report: dict[str, Any]) -> None: + title = ( + f"CortexForge 1000-User Benchmark | Users: {report['users']} | " + f"Workers: {report['workers']} | Status: {report['status']}" + ) + sep = "=" * 112 + thin = "-" * 112 + print(f"\n{sep}") + print(f" {title}") + print(sep) + header = ( + f"{'Stage':<48} | {'Ops':>6} | {'Time(s)':>7} | {'RPS':>9} | " + f"{'p50(ms)':>8} | {'p95(ms)':>8} | {'p99(ms)':>8} | {'Err':>4}" + ) + print(header) + print(thin) + for st in report["stages"]: + print( + f"{st['stage']:<48} | {st['operations']:>6} | {st['elapsed_sec']:>7.3f} | " + f"{st['rps']:>9.1f} | {st['p50_ms']:>8.2f} | {st['p95_ms']:>8.2f} | " + f"{st['p99_ms']:>8.2f} | {st['errors']:>4}" + ) + print(thin) + print( + f"{'TOTAL / SUMMARY':<48} | {report['total_operations']:>6} | " + f"{report['total_elapsed_sec']:>7.3f} | {report['overall_rps']:>9.1f} | " + f"{'-':>8} | {'-':>8} | {'-':>8} | {report['total_errors']:>4}" + ) + print(f"{sep}\n") + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="CortexForge 1000-User Scale & Concurrency Benchmark (ENT1K-3)" + ) + parser.add_argument( + "--users", + type=int, + default=1000, + help="Number of concurrent users/agents to benchmark (default: 1000)", + ) + parser.add_argument( + "--workers", + type=int, + default=64, + help="ThreadPoolExecutor max_workers concurrency (default: 64)", + ) + parser.add_argument( + "--output", + type=str, + default=None, + help="Optional path to save the JSON benchmark report", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = parse_args(argv) + if args.users < 1 or args.workers < 1: + print("Error: --users and --workers must be >= 1", file=sys.stderr) + return 2 + + report = run_benchmark(num_users=args.users, workers=args.workers) + print_report_table(report) + + if args.output: + out_path = pathlib.Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text( + json.dumps(report, indent=2, ensure_ascii=False) + "\n", + encoding="utf-8", + ) + print(f"Saved JSON benchmark report to: {out_path}") + + return 0 if report["total_errors"] == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/bulk-users.py b/scripts/bulk-users.py new file mode 100755 index 0000000..263c2ff --- /dev/null +++ b/scripts/bulk-users.py @@ -0,0 +1,779 @@ +#!/usr/bin/env python3 +""" +Bulk Provisioning & Fleet Management CLI (ENT1K-2) for CortexForge. + +Supports onboarding and managing up to 1,000+ isolated Hermes Agent instances +with atomic .env updates, compose shard profiles, port collision prevention, +and symlink-safe template copying. Python 3.12+ stdlib only. +""" +from __future__ import annotations + +import argparse +import csv +import fcntl +import json +import os +import pathlib +import re +import secrets +import shutil +import stat +import sys +import tempfile +from dataclasses import dataclass +from typing import Any + +NAME_RE = re.compile(r"^[a-z][a-z0-9_-]{1,31}$") +BOT_TOKEN_RE = re.compile(r"^[0-9]+:[A-Za-z0-9_-]+$") +TG_ID_RE = re.compile(r"^[0-9]+(?:,[0-9]+)*$") +DEPT_RE = re.compile(r"^[A-Za-z0-9_.-]{1,64}$") +PORT_MAPPING_RE = re.compile(r'127\.0\.0\.1:(\d+):9119') +SERVICE_HEADER_RE = re.compile(r"^ assistant-([a-z0-9_-]+):\s*$", re.MULTILINE) + + +@dataclass(frozen=True) +class UserSpec: + name: str + bot_token: str + full_name: str + tg_id: str + department: str = "" + quota_limit: int = 1_000_000 + ui_port: int | None = None + + +def _ensure_not_symlink(path: pathlib.Path, description: str = "path") -> None: + if path.is_symlink(): + raise ValueError(f"Symlink rejected for {description}: {path}") + + +def _verify_tree_no_symlinks(root: pathlib.Path) -> None: + _ensure_not_symlink(root, "template root") + for dirpath, dirnames, filenames in os.walk(root, followlinks=False): + base = pathlib.Path(dirpath) + for d in dirnames: + _ensure_not_symlink(base / d, "template directory") + for f in filenames: + _ensure_not_symlink(base / f, "template file") + + +def _safe_copy_file(src: pathlib.Path, dst: pathlib.Path, mode: int | None = None) -> None: + _ensure_not_symlink(src, "source file") + if dst.exists() or dst.is_symlink(): + _ensure_not_symlink(dst, "destination file") + shutil.copy2(src, dst, follow_symlinks=False) + if mode is not None: + os.chmod(dst, mode) + + +def _safe_write_file(path: pathlib.Path, content: str, mode: int = 0o600) -> None: + _ensure_not_symlink(path, "target file") + parent = path.parent + _ensure_not_symlink(parent, "parent directory") + fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=str(parent)) + tmp_path = pathlib.Path(tmp_name) + try: + os.fchmod(fd, mode) + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(content) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp_path, path) + finally: + if tmp_path.exists(): + tmp_path.unlink(missing_ok=True) + + +def _collect_migration_stems(root: pathlib.Path) -> list[str]: + mig_dir = root / "migrations" + if not mig_dir.is_dir(): + return [] + return sorted(p.stem for p in mig_dir.glob("[0-9]*.py") if p.is_file() and not p.is_symlink()) + + +def _parse_env_file(path: pathlib.Path) -> dict[str, str]: + result: dict[str, str] = {} + if not path.is_file() or path.is_symlink(): + return result + for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines(): + line = raw_line.strip() + if not line or line.startswith("#") or "=" not in line: + continue + k, v = line.split("=", 1) + val = v.strip() + if " #" in val: + val = val.split(" #", 1)[0].strip() + result[k.strip()] = val + return result + + +def _load_specs_from_file(file_path: pathlib.Path) -> list[dict[str, Any]]: + _ensure_not_symlink(file_path, "input file") + if not file_path.is_file(): + raise ValueError(f"Input file not found: {file_path}") + + suffix = file_path.suffix.lower() + raw_text = file_path.read_text(encoding="utf-8") + if suffix == ".json": + data = json.loads(raw_text) + if isinstance(data, dict) and "users" in data: + data = data["users"] + if not isinstance(data, list): + raise ValueError("JSON input must be a list of user objects or {'users': [...]}") + return data + if suffix == ".csv": + reader = csv.DictReader(raw_text.splitlines()) + return [dict(row) for row in reader] + raise ValueError(f"Unsupported file format '{suffix}'. Use .csv or .json") + + +def _validate_and_normalize_specs(raw_rows: list[dict[str, Any]]) -> list[UserSpec]: + if not raw_rows: + raise ValueError("Input file contains no user records") + + specs: list[UserSpec] = [] + seen_names: set[str] = set() + seen_env_suffixes: dict[str, str] = {} + seen_ports: set[int] = set() + + for idx, row in enumerate(raw_rows, start=1): + if not isinstance(row, dict): + raise ValueError(f"Record #{idx}: expected object/dict") + + name = str(row.get("name") or "").strip() + if name == "admin": + raise ValueError(f"Record #{idx}: name 'admin' is reserved") + if not NAME_RE.match(name): + raise ValueError( + f"Record #{idx}: invalid name '{name}' " + "(expected ^[a-z][a-z0-9_-]{1,31}$)" + ) + if name in seen_names: + raise ValueError(f"Record #{idx}: duplicate user name '{name}' in batch") + seen_names.add(name) + + env_suffix = name.upper().replace("-", "_") + if env_suffix in seen_env_suffixes: + raise ValueError( + f"Record #{idx}: name '{name}' collides on normalized env suffix " + f"'{env_suffix}' with '{seen_env_suffixes[env_suffix]}'" + ) + seen_env_suffixes[env_suffix] = name + + bot_token = str(row.get("bot_token") or "").strip() + if not bot_token or not BOT_TOKEN_RE.match(bot_token): + raise ValueError( + f"Record #{idx} ({name}): invalid bot_token (expected digits:alphanumeric)" + ) + + tg_id = str(row.get("tg_id") or "").strip() + if not tg_id or not TG_ID_RE.match(tg_id): + raise ValueError( + f"Record #{idx} ({name}): invalid tg_id '{tg_id}' (expected numeric Telegram ID)" + ) + + full_name = str(row.get("full_name") or "").strip() or name + if "\n" in full_name or "\r" in full_name: + raise ValueError(f"Record #{idx} ({name}): full_name must not contain newlines") + + department = str(row.get("department") or "").strip() + if department and not DEPT_RE.match(department): + raise ValueError( + f"Record #{idx} ({name}): invalid department '{department}'" + ) + + raw_quota = row.get("quota_limit") + if raw_quota is None or str(raw_quota).strip() == "": + quota_limit = 1_000_000 + else: + try: + quota_limit = int(str(raw_quota).strip()) + except ValueError as exc: + raise ValueError(f"Record #{idx} ({name}): non-integer quota_limit") from exc + if quota_limit <= 0: + raise ValueError(f"Record #{idx} ({name}): quota_limit must be > 0") + + raw_port = row.get("ui_port") + if raw_port is None or str(raw_port).strip() == "": + ui_port = None + else: + try: + ui_port = int(str(raw_port).strip()) + except ValueError as exc: + raise ValueError(f"Record #{idx} ({name}): non-integer ui_port") from exc + if ui_port < 1024 or ui_port > 65535 or ui_port == 18789: + raise ValueError( + f"Record #{idx} ({name}): ui_port {ui_port} out of allowed range (1024-65535, != 18789)" + ) + if ui_port in seen_ports: + raise ValueError( + f"Record #{idx} ({name}): duplicate ui_port {ui_port} in batch" + ) + seen_ports.add(ui_port) + + specs.append( + UserSpec( + name=name, + bot_token=bot_token, + full_name=full_name, + tg_id=tg_id, + department=department, + quota_limit=quota_limit, + ui_port=ui_port, + ) + ) + + return specs + + +def _parse_existing_override(override_path: pathlib.Path) -> tuple[dict[str, str], dict[str, int]]: + """Return (service_blocks_by_user, port_by_user) from existing docker-compose.override.yml.""" + if not override_path.is_file(): + return {}, {} + _ensure_not_symlink(override_path, "docker-compose.override.yml") + text = override_path.read_text(encoding="utf-8") + + # Strip trailing top-level networks: block if present so we can regenerate cleanly + net_marker = re.search(r"^networks:\s*$", text, re.MULTILINE) + services_text = text[: net_marker.start()] if net_marker else text + + matches = list(SERVICE_HEADER_RE.finditer(services_text)) + blocks: dict[str, str] = {} + ports: dict[str, int] = {} + for i, m in enumerate(matches): + uname = m.group(1) + start = m.start() + end = matches[i + 1].start() if i + 1 < len(matches) else len(services_text) + block = services_text[start:end].rstrip() + blocks[uname] = block + pm = PORT_MAPPING_RE.search(block) + if pm: + ports[uname] = int(pm.group(1)) + return blocks, ports + + +def _render_service_block( + name: str, + ui_port: int, + shard_id: int, + department: str = "", +) -> str: + name_upper = name.upper().replace("-", "_") + shard_name = f"shard-{shard_id}" + shard_net = f"corp-{shard_name}" + dept_label = f"\n corp.assistant.department: {department}" if department else "" + return f""" assistant-{name}: + build: + context: . + dockerfile: instances/Dockerfile.user + image: cortex-forge-user:latest + container_name: corp-{name} + restart: unless-stopped + profiles: + - "{shard_name}" + - "all-users" + env_file: + - instances/{name}/.env + volumes: + - ./instances/{name}/hermes_data:/opt/data + - ./instances/{name}/workspace:/opt/data/workspace + - ./shared/skills:/shared/skills + - ./shared/docs:/shared/docs + - ./shared/compliance-data:/shared/compliance-data:ro + environment: + - ANTHROPIC_API_KEY=${{QUOTA_KEY_{name_upper}}} + - OPENAI_API_KEY=${{OPENAI_API_KEY:-}} + - SERVICE_API_KEY=${{SERVICE_API_KEY}} + - BROKER_URL=http://message-broker:8080 + - BROKER_KEY=${{BROKER_KEY_{name_upper}}} + - HERMES_UID=1000 + - HERMES_GID=1000 + - HERMES_DASHBOARD=1 + ports: + - "127.0.0.1:{ui_port}:9119" + networks: + - corp-internal + - corp-outbound + - {shard_net} + security_opt: + - no-new-privileges:true + deploy: + resources: + limits: + cpus: "2.0" + memory: 1G + depends_on: + quota-proxy: + condition: service_healthy + message-broker: + condition: service_healthy + labels: + corp.assistant.role: personal + corp.assistant.user: {name} + corp.assistant.shard: {shard_name}{dept_label}""" + + +def _build_override_content(service_blocks: dict[str, str]) -> str: + if not service_blocks: + return "services: {}\n" + + shards_seen: set[int] = set() + for block in service_blocks.values(): + for m in re.finditer(r"corp-shard-(\d+)", block): + shards_seen.add(int(m.group(1))) + + lines = ["services:"] + for uname in service_blocks: + lines.append(service_blocks[uname]) + + if shards_seen: + lines.append("") + lines.append("networks:") + for sid in sorted(shards_seen): + lines.append(f" corp-shard-{sid}:") + lines.append(" driver: bridge") + lines.append(" internal: true") + + return "\n".join(lines) + "\n" + + +def _provision_instance_files( + root: pathlib.Path, + spec: UserSpec, + broker_key: str, + quota_key: str, + gateway_token: str, + migration_stems: list[str], +) -> None: + template_dir = root / "instances" / "_template" + template_ws = template_dir / "workspace" + template_cfg = template_dir / "hermes_config.yaml.template" + + target_dir = root / "instances" / spec.name + _ensure_not_symlink(target_dir, f"instance directory {spec.name}") + target_ws = target_dir / "workspace" + target_hermes = target_dir / "hermes_data" + + os.makedirs(target_hermes / "ssh", exist_ok=True) + os.makedirs(target_ws / "skills", exist_ok=True) + os.makedirs(target_ws / "memory", exist_ok=True) + + # Copy template workspace files/dirs (excluding memory/) without following symlinks + for item in template_ws.iterdir(): + _ensure_not_symlink(item, "template workspace item") + if item.name == "memory": + continue + dest = target_ws / item.name + if item.is_dir(): + shutil.copytree( + item, + dest, + symlinks=False, + dirs_exist_ok=True, + copy_function=lambda s, d: _safe_copy_file(pathlib.Path(s), pathlib.Path(d)), + ) + elif item.is_file(): + _safe_copy_file(item, dest) + + # Substitute placeholders in USER.md, IDENTITY.md, SOUL.md + for fname in ("USER.md", "IDENTITY.md", "SOUL.md"): + fpath = target_ws / fname + if fpath.is_file(): + _ensure_not_symlink(fpath, fname) + text = fpath.read_text(encoding="utf-8") + text = text.replace("{{FULL_NAME}}", spec.full_name).replace("{{NAME}}", spec.name) + _safe_write_file(fpath, text, mode=0o644) + + # Record applied migrations so fresh instances don't re-run historical migrations + applied_file = target_ws / ".migrations_applied" + _safe_write_file( + applied_file, + ("\n".join(migration_stems) + "\n") if migration_stems else "", + mode=0o644, + ) + + # Copy SOUL.md into hermes_data/SOUL.md with symlink protection + soul_src = target_ws / "SOUL.md" + soul_dst = target_hermes / "SOUL.md" + if soul_src.is_file(): + _safe_copy_file(soul_src, soul_dst, mode=0o600) + + # Copy hermes_config.yaml.template -> hermes_data/config.yaml (0600) + _safe_copy_file(template_cfg, target_hermes / "config.yaml", mode=0o600) + + # Create instances//.env (0600) + dept_comment = f"# DEPARTMENT={spec.department}\n" if spec.department else "" + instance_env = ( + f"# Instance secrets for {spec.name} — DO NOT commit to git!\n" + f"{dept_comment}" + f"BROKER_URL=http://message-broker:8080\n" + f"BROKER_KEY={broker_key}\n" + ) + _safe_write_file(target_dir / ".env", instance_env, mode=0o600) + + # Create instances//hermes_data/.env (0600) + hermes_env = ( + f"# Hermes secrets for {spec.name} — DO NOT commit to git!\n" + f"TELEGRAM_BOT_TOKEN={spec.bot_token}\n" + f"TELEGRAM_ALLOWED_USERS={spec.tg_id}\n" + f"HERMES_GATEWAY_TOKEN={gateway_token}\n" + f"ANTHROPIC_API_KEY={quota_key}\n" + f"# OPENAI_API_KEY=\n" + ) + _safe_write_file(target_hermes / ".env", hermes_env, mode=0o600) + + +def _atomic_append_root_env(root_env: pathlib.Path, entries: list[str]) -> None: + if not entries: + return + _ensure_not_symlink(root_env, "root .env") + lock_path = root_env.parent / ".env.bulk.lock" + with open(lock_path, "a", encoding="utf-8") as lock_file: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + try: + current = root_env.read_text(encoding="utf-8") if root_env.exists() else "" + if current and not current.endswith("\n"): + current += "\n" + updated = current + "\n".join(entries) + "\n" + _safe_write_file(root_env, updated, mode=0o600) + finally: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + + +def cmd_import(args: argparse.Namespace) -> int: + root = pathlib.Path(args.root).resolve() + if args.shard_size <= 0: + print("❌ --shard-size must be >= 1", file=sys.stderr) + return 1 + if args.base_port < 1024 or args.base_port > 65535 or args.base_port == 18789: + print("❌ --base-port must be between 1024 and 65535 (excluding 18789)", file=sys.stderr) + return 1 + + root_env = root / ".env" + if not root_env.is_file() or root_env.is_symlink(): + print(f"❌ Root .env not found or is a symlink: {root_env}", file=sys.stderr) + return 1 + + template_dir = root / "instances" / "_template" + if not (template_dir / "workspace").is_dir() or not (template_dir / "hermes_config.yaml.template").is_file(): + print(f"❌ Template directory incomplete: {template_dir}", file=sys.stderr) + return 1 + + try: + _verify_tree_no_symlinks(template_dir) + raw_rows = _load_specs_from_file(pathlib.Path(args.file)) + specs = _validate_and_normalize_specs(raw_rows) + except (ValueError, json.JSONDecodeError) as exc: + print(f"❌ Validation error: {exc}", file=sys.stderr) + return 1 + + instances_dir = root / "instances" + os.makedirs(instances_dir, exist_ok=True) + + existing_env_suffixes: dict[str, str] = {} + for d in sorted(instances_dir.iterdir(), key=lambda p: p.name): + if not d.is_dir() or d.name in ("_template", "admin"): + continue + existing_env_suffixes[d.name.upper().replace("-", "_")] = d.name + + # Pre-check existing target directories for symlinks and normalized env-suffix collisions + for spec in specs: + candidate = instances_dir / spec.name + if candidate.is_symlink(): + print(f"❌ Validation error: target instance path is a symlink: {candidate}", file=sys.stderr) + return 1 + env_suffix = spec.name.upper().replace("-", "_") + existing_name = existing_env_suffixes.get(env_suffix) + if existing_name is not None and existing_name != spec.name: + print( + f"❌ Validation error: instance '{spec.name}' collides on normalized env suffix " + f"'{env_suffix}' with existing instance '{existing_name}'", + file=sys.stderr, + ) + return 1 + + override_path = root / "docker-compose.override.yml" + try: + service_blocks, existing_ports = _parse_existing_override(override_path) + except ValueError as exc: + print(f"❌ Override error: {exc}", file=sys.stderr) + return 1 + + used_ports: set[int] = set(existing_ports.values()) | {18789} + # Reserve explicit ports from non-existing specs first and check collisions with existing ports + for spec in specs: + target_dir = instances_dir / spec.name + if target_dir.exists(): + continue + if spec.ui_port is not None: + if spec.ui_port in used_ports: + print( + f"❌ Port collision: ui_port {spec.ui_port} for '{spec.name}' is already in use", + file=sys.stderr, + ) + return 1 + used_ports.add(spec.ui_port) + + next_port = args.base_port + created: list[dict[str, Any]] = [] + skipped: list[str] = [] + env_lines_to_append: list[str] = [] + migration_stems = _collect_migration_stems(root) + + existing_user_count = len(service_blocks) + + for spec in specs: + target_dir = instances_dir / spec.name + if target_dir.exists(): + skipped.append(spec.name) + continue + + if spec.ui_port is not None: + port = spec.ui_port + else: + while next_port in used_ports: + next_port += 1 + port = next_port + used_ports.add(port) + next_port += 1 + + user_index = existing_user_count + len(created) + shard_id = user_index // args.shard_size + + broker_key = secrets.token_urlsafe(32) + quota_key = f"quota-{spec.name}-" + secrets.token_urlsafe(24) + gateway_token = secrets.token_urlsafe(32) + name_upper = spec.name.upper().replace("-", "_") + + if not args.dry_run: + _provision_instance_files( + root=root, + spec=spec, + broker_key=broker_key, + quota_key=quota_key, + gateway_token=gateway_token, + migration_stems=migration_stems, + ) + + env_lines_to_append.extend( + [ + f"BROKER_KEY_{name_upper}={broker_key}", + f"QUOTA_KEY_{name_upper}={quota_key}", + f"QUOTA_LIMIT_{name_upper}={spec.quota_limit}", + ] + ) + if spec.department: + env_lines_to_append.append(f"QUOTA_DEPT_{name_upper}={spec.department}") + service_blocks[spec.name] = _render_service_block( + name=spec.name, + ui_port=port, + shard_id=shard_id, + department=spec.department, + ) + created.append( + { + "name": spec.name, + "ui_port": port, + "shard": f"shard-{shard_id}", + "quota_limit": spec.quota_limit, + "department": spec.department, + } + ) + + if not args.dry_run and created: + _atomic_append_root_env(root_env, env_lines_to_append) + override_content = _build_override_content(service_blocks) + _safe_write_file(override_path, override_content, mode=0o644) + + summary = { + "dry_run": bool(args.dry_run), + "created_count": len(created), + "skipped_count": len(skipped), + "shard_size": args.shard_size, + "created": created, + "skipped": skipped, + } + print(json.dumps(summary, ensure_ascii=False, indent=2)) + return 0 + + +def _collect_fleet_records(root: pathlib.Path, include_secrets: bool = False) -> list[dict[str, Any]]: + instances_dir = root / "instances" + root_env = _parse_env_file(root / ".env") + override_path = root / "docker-compose.override.yml" + service_blocks, override_ports = _parse_existing_override(override_path) + + records: list[dict[str, Any]] = [] + if not instances_dir.is_dir(): + return records + + for d in sorted(instances_dir.iterdir(), key=lambda p: p.name): + if not d.is_dir() or d.is_symlink() or d.name in ("_template", "admin"): + continue + name = d.name + name_upper = name.upper().replace("-", "_") + inst_env_path = d / ".env" + hermes_env_path = d / "hermes_data" / ".env" + hermes_cfg_path = d / "hermes_data" / "config.yaml" + hermes_soul_path = d / "hermes_data" / "SOUL.md" + ws_user_path = d / "workspace" / "USER.md" + + inst_env = _parse_env_file(inst_env_path) + hermes_env = _parse_env_file(hermes_env_path) + + # Extract department comment if present + department = "" + if inst_env_path.is_file() and not inst_env_path.is_symlink(): + for line in inst_env_path.read_text(encoding="utf-8", errors="replace").splitlines(): + if line.startswith("# DEPARTMENT="): + department = line.split("=", 1)[1].strip() + break + + full_name = name + if ws_user_path.is_file() and not ws_user_path.is_symlink(): + m = re.search(r"^-\s*(?:Имя|Name):\s*(.+)$", ws_user_path.read_text(encoding="utf-8", errors="replace"), re.MULTILINE) + if m: + full_name = m.group(1).strip() + + block = service_blocks.get(name, "") + shard = "unassigned" + sm = re.search(r"corp\.assistant\.shard:\s*(shard-\d+)", block) + if sm: + shard = sm.group(1) + else: + pm = re.search(r'-\s*"(shard-\d+)"', block) + if pm: + shard = pm.group(1) + + raw_limit = root_env.get(f"QUOTA_LIMIT_{name_upper}", "1000000") + try: + quota_limit = int(raw_limit) + except ValueError: + quota_limit = 1_000_000 + + checks = { + "instance_env": inst_env_path.is_file() and not inst_env_path.is_symlink(), + "hermes_env": hermes_env_path.is_file() and not hermes_env_path.is_symlink(), + "hermes_config": hermes_cfg_path.is_file() and not hermes_cfg_path.is_symlink(), + "hermes_soul": hermes_soul_path.is_file() and not hermes_soul_path.is_symlink(), + } + healthy = all(checks.values()) + + rec: dict[str, Any] = { + "name": name, + "full_name": full_name, + "tg_id": hermes_env.get("TELEGRAM_ALLOWED_USERS", ""), + "department": department, + "ui_port": override_ports.get(name), + "shard": shard, + "quota_limit": quota_limit, + "healthy": healthy, + "checks": checks, + } + if include_secrets: + rec["secrets"] = { + "bot_token": hermes_env.get("TELEGRAM_BOT_TOKEN", ""), + "hermes_gateway_token": hermes_env.get("HERMES_GATEWAY_TOKEN", ""), + "broker_key": inst_env.get("BROKER_KEY", root_env.get(f"BROKER_KEY_{name_upper}", "")), + "quota_key": hermes_env.get("ANTHROPIC_API_KEY", root_env.get(f"QUOTA_KEY_{name_upper}", "")), + } + records.append(rec) + + return records + + +def cmd_export(args: argparse.Namespace) -> int: + root = pathlib.Path(args.root).resolve() + out_path = pathlib.Path(args.output) + records = _collect_fleet_records(root, include_secrets=bool(args.include_secrets)) + payload = { + "instance_count": len(records), + "include_secrets": bool(args.include_secrets), + "instances": records, + } + mode = 0o600 if args.include_secrets else 0o644 + os.makedirs(out_path.parent if out_path.parent != pathlib.Path("") else pathlib.Path("."), exist_ok=True) + _safe_write_file(out_path, json.dumps(payload, ensure_ascii=False, indent=2) + "\n", mode=mode) + os.chmod(out_path, mode) + print(json.dumps({"exported": len(records), "output": str(out_path)}, ensure_ascii=False)) + return 0 + + +def cmd_status(args: argparse.Namespace) -> int: + root = pathlib.Path(args.root).resolve() + records = _collect_fleet_records(root, include_secrets=False) + shards: dict[str, int] = {} + healthy_count = 0 + degraded_instances: list[dict[str, Any]] = [] + ports_seen: dict[int, list[str]] = {} + + for rec in records: + shard = rec["shard"] + shards[shard] = shards.get(shard, 0) + 1 + if rec["healthy"]: + healthy_count += 1 + else: + missing = [k for k, ok in rec["checks"].items() if not ok] + degraded_instances.append({"name": rec["name"], "missing": missing}) + port = rec.get("ui_port") + if isinstance(port, int): + ports_seen.setdefault(port, []).append(rec["name"]) + + port_collisions = {str(p): names for p, names in ports_seen.items() if len(names) > 1} + status_report = { + "total_instances": len(records), + "healthy_instances": healthy_count, + "degraded_instances": len(degraded_instances), + "shards": shards, + "port_collisions": port_collisions, + "degraded_details": degraded_instances, + } + print(json.dumps(status_report, ensure_ascii=False, indent=2)) + return 0 if not degraded_instances and not port_collisions else 1 + + +def build_parser() -> argparse.ArgumentParser: + default_root = pathlib.Path.cwd() + if not (default_root / "instances").exists(): + default_root = pathlib.Path(__file__).resolve().parent.parent + + parser = argparse.ArgumentParser( + prog="bulk-users.py", + description="Bulk user provisioning & fleet management for CortexForge (1000+ instances)", + ) + parser.add_argument( + "--root", + default=str(default_root), + help="Repository root directory (default: current directory or repo root)", + ) + sub = parser.add_subparsers(dest="command", required=True) + + p_import = sub.add_parser("import", help="Bulk import users from CSV or JSON") + p_import.add_argument("--file", required=True, help="Path to users.csv or users.json") + p_import.add_argument("--dry-run", action="store_true", help="Validate and preview without writing files") + p_import.add_argument("--shard-size", type=int, default=100, help="Instances per compose profile shard (default: 100)") + p_import.add_argument("--base-port", type=int, default=18790, help="Starting localhost dashboard port (default: 18790)") + p_import.add_argument("--root", default=argparse.SUPPRESS, help="Repository root directory") + p_import.set_defaults(func=cmd_import) + + p_export = sub.add_parser("export", help="Export fleet metadata manifest to JSON") + p_export.add_argument("--output", required=True, help="Output manifest JSON file path") + p_export.add_argument("--include-secrets", action="store_true", help="Include secret tokens (writes 0600 file)") + p_export.add_argument("--root", default=argparse.SUPPRESS, help="Repository root directory") + p_export.set_defaults(func=cmd_export) + + p_status = sub.add_parser("status", help="Inspect fleet health, shard distribution, and file integrity") + p_status.add_argument("--root", default=argparse.SUPPRESS, help="Repository root directory") + p_status.set_defaults(func=cmd_status) + + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() + args = parser.parse_args(argv) + return args.func(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/provision-hermes-data.sh b/scripts/provision-hermes-data.sh index 6c08601..35b05d2 100755 --- a/scripts/provision-hermes-data.sh +++ b/scripts/provision-hermes-data.sh @@ -33,7 +33,7 @@ provision() { # .env для Hermes — из instance .env (отклоняем симлинки; || true защищает от pipefail) if [ ! -e "$DIR/hermes_data/.env" ] && [ ! -L "$DIR/hermes_data/.env" ] && [ -f "$DIR/.env" ] && [ ! -L "$DIR/.env" ]; then local tg_token tg_users gw ak oai name_upper - name_upper=$(echo "$NAME" | tr '[:lower:]' '[:upper:]') + name_upper=$(echo "$NAME" | tr '[:lower:]-' '[:upper:]_') tg_token=$(grep -E '^TELEGRAM_BOT_TOKEN=' "$DIR/.env" | head -1 | cut -d= -f2- || true) tg_users=$(grep -E '^TELEGRAM_ALLOWED_USERS=' "$DIR/.env" | head -1 | cut -d= -f2- || true) [ -z "$tg_users" ] && tg_users=$(grep -E '^TELEGRAM_ALLOW_FROM=' "$DIR/.env" | head -1 | cut -d= -f2- || true) diff --git a/scripts/remove-user.sh b/scripts/remove-user.sh index b225b11..75b9c9b 100644 --- a/scripts/remove-user.sh +++ b/scripts/remove-user.sh @@ -43,7 +43,7 @@ rm -rf "$TARGET" echo "✅ Инстанс '$NAME' удалён" # Удаляем ключи из глобального .env -NAME_UPPER=$(echo "$NAME" | tr '[:lower:]' '[:upper:]') +NAME_UPPER=$(echo "$NAME" | tr '[:lower:]-' '[:upper:]_') if [ -f ".env" ]; then # Удаляем строки BROKER_KEY_*, QUOTA_KEY_*, QUOTA_LIMIT_* для этого пользователя python3 -c " diff --git a/scripts/sync-instances.sh b/scripts/sync-instances.sh index 944db05..5b63666 100755 --- a/scripts/sync-instances.sh +++ b/scripts/sync-instances.sh @@ -51,22 +51,33 @@ for INST_DIR in "$INSTANCES"/*/; do fi done -# ─── 3. Контейнеры живы? ─────────────────────────────────────────────────── +# ─── 3. Контейнеры живы? (O(1) docker ps -a для 1000 инстансов) ──────────── +RUNNING_MAP=$(docker ps -a --format '{{.Names}}={{.State}}' 2>/dev/null || true) +STOPPED_SERVICES=() + for INST_DIR in "$INSTANCES"/*/; do NAME=$(basename "$INST_DIR") [ "$NAME" = "_template" ] && continue [ "$NAME" = "$ADMIN_INSTANCE" ] && continue # admin manages itself CONTAINER="corp-$NAME" - STATUS=$(docker inspect --format '{{.State.Status}}' "$CONTAINER" 2>/dev/null || echo "not_found") + STATUS=$(echo "$RUNNING_MAP" | grep -E "^${CONTAINER}=" | head -1 | cut -d= -f2- || true) + [ -z "$STATUS" ] && STATUS="not_found" if [ "$STATUS" != "running" ]; then - # Пробуем поднять (cd в /infra/ чтобы compose нашёл и base, и override) - (cd "$INFRA" && docker compose up -d "assistant-$NAME" 2>&1) || true - CHANGES+=("[$NAME] контейнер был '$STATUS' — попытка запуска") + STOPPED_SERVICES+=("assistant-$NAME") + CHANGES+=("[$NAME] контейнер был '$STATUS' — добавлен в очередь запуска") fi done +if [ ${#STOPPED_SERVICES[@]} -gt 0 ]; then + # Запускаем упавшие сервисы батчами по 50, чтобы не перегрузить демон при 1000 инстансах + ( + cd "$INFRA" + printf '%s\n' "${STOPPED_SERVICES[@]}" | xargs -n 50 docker compose up -d 2>&1 + ) || true +fi + # ─── Итог ────────────────────────────────────────────────────────────────── if [ ${#CHANGES[@]} -eq 0 ]; then echo "ALL_OK" diff --git a/service-agent/server.py b/service-agent/server.py index e100b01..1132382 100644 --- a/service-agent/server.py +++ b/service-agent/server.py @@ -22,18 +22,21 @@ import logging import threading import subprocess -from http.server import HTTPServer, BaseHTTPRequestHandler +from http.server import HTTPServer, ThreadingHTTPServer, BaseHTTPRequestHandler from datetime import datetime, timezone from pathlib import Path +from urllib.parse import urlparse, parse_qs # ── Configuration ───────────────────────────────────────────────────────────── -SERVICE_PORT = int(os.environ.get("SERVICE_PORT", "8090")) -SERVICE_API_KEY = os.environ.get("SERVICE_API_KEY", "") -_service_api_hash = hashlib.sha256(SERVICE_API_KEY.encode()).hexdigest() if SERVICE_API_KEY else "" -SKILL_TIMEOUT = int(os.environ.get("SKILL_TIMEOUT", "120")) -MAX_CONCURRENT_TASKS = int(os.environ.get("MAX_CONCURRENT_TASKS", "5")) -USAGE_DB = os.environ.get("USAGE_DB", "/data/usage.db") +SERVICE_PORT = int(os.environ.get("SERVICE_PORT", "8090")) +SERVICE_API_KEY = os.environ.get("SERVICE_API_KEY", "") +_service_api_hash = hashlib.sha256(SERVICE_API_KEY.encode()).hexdigest() if SERVICE_API_KEY else "" +SKILL_TIMEOUT = int(os.environ.get("SKILL_TIMEOUT", "120")) +MAX_CONCURRENT_SKILLS = int(os.environ.get("MAX_CONCURRENT_SKILLS", os.environ.get("MAX_CONCURRENT_TASKS", "5"))) +MAX_CONCURRENT_TASKS = MAX_CONCURRENT_SKILLS # backward compatibility alias +SKILL_RATE_LIMIT = int(os.environ.get("SKILL_RATE_LIMIT", "30")) +USAGE_DB = os.environ.get("USAGE_DB", "/data/usage.db") MAX_BODY_BYTES = 1 * 1024 * 1024 # 1 MB — защита от oversized requests SKILLS_DIR = Path(os.environ.get("SKILLS_DIR", "/app/skills")) @@ -46,7 +49,54 @@ log = logging.getLogger("service-agent") # Semaphore to cap concurrent skill executions -_task_semaphore = threading.Semaphore(MAX_CONCURRENT_TASKS) +_task_semaphore = threading.Semaphore(MAX_CONCURRENT_SKILLS) + + +# ── Per-caller rate limiting (leaky bucket) ─────────────────────────────────── + +_CALLER_RE = re.compile(r"^[a-zA-Z0-9_.:-]{1,64}$") + + +class CallerRateLimiter: + """Leaky bucket per caller: max capacity tokens, refills at rate/sec.""" + + def __init__(self, rate_per_min: int): + self.set_limit(rate_per_min) + self._buckets: dict[str, tuple[float, float]] = {} + self._lock = threading.Lock() + + def set_limit(self, rate_per_min: int): + self.capacity = max(1, int(rate_per_min)) + self.rate = self.capacity / 60.0 + + def reset(self): + with self._lock: + self._buckets.clear() + + def allow(self, caller: str) -> bool: + if SKILL_RATE_LIMIT <= 0: + return True + # Sync capacity if SKILL_RATE_LIMIT was updated dynamically + if self.capacity != SKILL_RATE_LIMIT: + self.set_limit(SKILL_RATE_LIMIT) + now = time.monotonic() + raw_key = (caller or "unknown").strip() or "unknown" + key = raw_key if _CALLER_RE.match(raw_key) else "invalid-caller" + with self._lock: + if len(self._buckets) > 10000: + self._buckets = { + k: v for k, v in self._buckets.items() if now - v[1] <= 120 + } + tokens, last = self._buckets.get(key, (float(self.capacity), now)) + tokens = min(float(self.capacity), tokens + (now - last) * self.rate) + if tokens < 1.0: + self._buckets[key] = (tokens, now) + return False + self._buckets[key] = (tokens - 1.0, now) + return True + + +_caller_limiter = CallerRateLimiter(SKILL_RATE_LIMIT) # Explicit skill whitelist — populated at startup and on reload. # run_skill() validates skill_id against this set before any exec call. @@ -56,9 +106,17 @@ # ── SQLite usage tracking ───────────────────────────────────────────────────── +def _get_db_conn() -> sqlite3.Connection: + conn = sqlite3.connect(USAGE_DB, timeout=10.0) + conn.execute("PRAGMA busy_timeout=10000") + return conn + + def init_db(): os.makedirs(os.path.dirname(USAGE_DB) or ".", exist_ok=True) - conn = sqlite3.connect(USAGE_DB) + conn = _get_db_conn() + conn.execute("PRAGMA journal_mode=WAL") + conn.execute("PRAGMA journal_size_limit=1048576") conn.execute(""" CREATE TABLE IF NOT EXISTS calls ( id INTEGER PRIMARY KEY AUTOINCREMENT, @@ -71,13 +129,15 @@ def init_db(): output_len INTEGER NOT NULL ) """) + conn.execute("CREATE INDEX IF NOT EXISTS idx_calls_caller ON calls (caller)") + conn.execute("CREATE INDEX IF NOT EXISTS idx_calls_skill ON calls (skill)") conn.commit() conn.close() def record_call(caller: str, skill: str, params: dict, success: bool, duration_ms: int, output_len: int): - conn = sqlite3.connect(USAGE_DB) + conn = _get_db_conn() try: conn.execute( "INSERT INTO calls (timestamp, caller, skill, params_json, success, duration_ms, output_len) " @@ -97,10 +157,26 @@ def record_call(caller: str, skill: str, params: dict, conn.close() -def get_usage_stats() -> dict: - conn = sqlite3.connect(USAGE_DB) +def get_usage_stats( + caller: str | None = None, + skill: str | None = None, + limit: int | None = None, + offset: int = 0, +) -> dict: + conn = _get_db_conn() try: - rows = conn.execute(""" + where_clauses = [] + where_params: list = [] + if caller: + where_clauses.append("caller = ?") + where_params.append(caller) + if skill: + where_clauses.append("skill = ?") + where_params.append(skill) + + where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else "" + + skill_rows = conn.execute(f""" SELECT skill, COUNT(*) AS total_calls, SUM(success) AS successful, @@ -108,20 +184,95 @@ def get_usage_stats() -> dict: SUM(output_len) AS total_output_bytes, MAX(timestamp) AS last_called FROM calls + {where_sql} GROUP BY skill - """).fetchall() + ORDER BY total_calls DESC, skill ASC + """, where_params).fetchall() + + caller_rows = conn.execute(f""" + SELECT caller, + COUNT(*) AS total_calls, + SUM(success) AS successful, + AVG(duration_ms) AS avg_duration_ms, + SUM(output_len) AS total_output_bytes, + MAX(timestamp) AS last_called + FROM calls + {where_sql} + GROUP BY caller + ORDER BY total_calls DESC, caller ASC + """, where_params).fetchall() + + total_calls = conn.execute( + f"SELECT COUNT(*) FROM calls {where_sql}", where_params + ).fetchone()[0] + + by_skill_all = [ + { + "skill": r[0], + "total_calls": r[1], + "successful": r[2] or 0, + "avg_duration_ms": round(r[3] or 0, 1), + "total_output_bytes": r[4] or 0, + "last_called": r[5], + } + for r in skill_rows + ] + + by_caller_all = [ + { + "caller": r[0], + "total_calls": r[1], + "successful": r[2] or 0, + "avg_duration_ms": round(r[3] or 0, 1), + "total_output_bytes": r[4] or 0, + "last_called": r[5], + } + for r in caller_rows + ] + + safe_offset = max(0, int(offset or 0)) + if limit is not None: + safe_limit = max(1, int(limit)) + by_skill = by_skill_all[safe_offset:safe_offset + safe_limit] + by_caller = by_caller_all[safe_offset:safe_offset + safe_limit] + elif safe_offset > 0: + by_skill = by_skill_all[safe_offset:] + by_caller = by_caller_all[safe_offset:] + else: + by_skill = by_skill_all + by_caller = by_caller_all + + call_limit = max(1, int(limit)) if limit is not None else 100 + recent_rows = conn.execute(f""" + SELECT id, timestamp, caller, skill, success, duration_ms, output_len + FROM calls + {where_sql} + ORDER BY id DESC + LIMIT ? OFFSET ? + """, [*where_params, call_limit, safe_offset]).fetchall() + + calls_list = [ + { + "id": r[0], + "timestamp": r[1], + "caller": r[2], + "skill": r[3], + "success": bool(r[4]), + "duration_ms": r[5], + "output_len": r[6], + } + for r in recent_rows + ] + return { - "by_skill": [ - { - "skill": r[0], - "total_calls": r[1], - "successful": r[2], - "avg_duration_ms": round(r[3] or 0, 1), - "total_output_bytes": r[4], - "last_called": r[5], - } - for r in rows - ] + "by_skill": by_skill, + "by_caller": by_caller, + "calls": calls_list, + "total_calls": total_calls, + "total_skills": len(by_skill_all), + "total_callers": len(by_caller_all), + "limit": limit, + "offset": safe_offset, } finally: conn.close() @@ -386,12 +537,18 @@ def send_json(handler, data: dict, status: int = 200): handler.wfile.write(body) -def read_body(handler) -> bytes | None: - length = int(handler.headers.get("Content-Length", "0")) +def read_body(handler) -> tuple[bytes | None, int | None]: + raw_len = handler.headers.get("Content-Length", "0") + try: + length = int(raw_len) + except (TypeError, ValueError): + return None, 400 + if length < 0: + return None, 400 if length > MAX_BODY_BYTES: handler.close_connection = True - return None - return handler.rfile.read(length) if length else b"" + return None, 413 + return (handler.rfile.read(length) if length else b""), None # ── HTTP Handler ────────────────────────────────────────────────────────────── @@ -408,9 +565,11 @@ def do_GET(self): if path == "/v1/health": send_json(self, { - "status": "ok", - "skills_dir": str(SKILLS_DIR), - "max_concurrent": MAX_CONCURRENT_TASKS, + "status": "ok", + "skills_dir": str(SKILLS_DIR), + "max_concurrent": MAX_CONCURRENT_SKILLS, + "max_concurrent_skills": MAX_CONCURRENT_SKILLS, + "skill_rate_limit": SKILL_RATE_LIMIT, }) return @@ -423,7 +582,27 @@ def do_GET(self): send_json(self, {"skills": list(skills.values()), "count": len(skills)}) elif path == "/v1/usage": - send_json(self, get_usage_stats()) + parsed = urlparse(self.path) + qs = parse_qs(parsed.query) + caller_filter = qs.get("caller", [None])[0] + skill_filter = qs.get("skill", [None])[0] + raw_limit = qs.get("limit", [None])[0] + raw_offset = qs.get("offset", ["0"])[0] + try: + limit_val = int(raw_limit) if raw_limit is not None else None + offset_val = int(raw_offset) if raw_offset is not None else 0 + except ValueError: + send_json(self, {"status": "error", "error": "Invalid limit or offset"}, 400) + return + send_json( + self, + get_usage_stats( + caller=caller_filter, + skill=skill_filter, + limit=limit_val, + offset=offset_val, + ), + ) else: send_json(self, {"status": "error", "error": "not found"}, 404) @@ -437,8 +616,11 @@ def do_POST(self): return if path == "/v1/run": - raw = read_body(self) - if raw is None: + raw, err_code = read_body(self) + if err_code == 400: + send_json(self, {"status": "error", "error": "Invalid Content-Length"}, 400) + return + if err_code == 413 or raw is None: send_json(self, {"status": "error", "error": "Request body too large"}, 413) return try: @@ -448,9 +630,13 @@ def do_POST(self): return skill_id = body.get("skill", "").strip() - caller = body.get("caller", "unknown") + caller = str(body.get("caller", "unknown") or "unknown").strip() or "unknown" params = body.get("params", {}) + if not _CALLER_RE.match(caller): + send_json(self, {"status": "error", "error": "Invalid caller name"}, 400) + return + if not skill_id: send_json(self, {"status": "error", "error": "Missing 'skill' field"}, 400) return @@ -460,6 +646,20 @@ def do_POST(self): send_json(self, {"status": "error", "error": "Invalid skill name"}, 400) return + # Per-caller rate limiting (protects pool from a runaway bot) + if not _caller_limiter.allow(caller): + log.warning("Rate limit exceeded for caller=%s skill=%s", caller, skill_id) + send_json( + self, + { + "status": "error", + "error": f"Rate limit exceeded for caller '{caller}' ({SKILL_RATE_LIMIT} req/min)", + "caller": caller, + }, + 429, + ) + return + # Whitelist check: skill must be in ALLOWED_SKILLS (populated at startup) with _skills_lock: allowed_now = frozenset(_ALLOWED_SKILLS) @@ -479,14 +679,17 @@ def do_POST(self): # ── Entry Point ─────────────────────────────────────────────────────────────── def main(): - log.info(f"Service Agent starting — port={SERVICE_PORT} max_concurrent={MAX_CONCURRENT_TASKS}") + log.info( + f"Service Agent starting — port={SERVICE_PORT} " + f"max_concurrent={MAX_CONCURRENT_SKILLS} rate_limit={SKILL_RATE_LIMIT}/min" + ) log.info(f"Skills dir: {SKILLS_DIR}") init_db() skills_set = reload_skills() # populates _ALLOWED_SKILLS whitelist log.info(f"Loaded {len(skills_set)} skill(s): {', '.join(sorted(skills_set)) or 'none'}") - server = HTTPServer(("0.0.0.0", SERVICE_PORT), ServiceHandler) + server = ThreadingHTTPServer(("0.0.0.0", SERVICE_PORT), ServiceHandler) log.info(f"Listening on 0.0.0.0:{SERVICE_PORT}") try: server.serve_forever() diff --git a/tests/test_broker_enterprise.py b/tests/test_broker_enterprise.py new file mode 100644 index 0000000..c905a4a --- /dev/null +++ b/tests/test_broker_enterprise.py @@ -0,0 +1,625 @@ +#!/usr/bin/env python3 +""" +Enterprise scale & feature tests for CortexForge Message Broker (ENT1K-4). + +Covers: + 1. Dynamic key registration (POST /admin/register-key) & hot-reload (POST /admin/reload-keys) + 2. Underscore vs hyphen instance identity preservation (BROKER_KEY_TEAM_ONE vs BROKER_KEY_USER_1) + 3. Safe SQLite schema migration (ALTER TABLE messages ADD COLUMN priority ...) + 4. Message priorities (low, normal, high, urgent) and custom ttl_sec / TTL expiry + 5. Admin broadcast / multicast (POST /broadcast with list and "*") + 6. Inbox pagination & filtering (GET /inbox?peek=1&limit=&offset=&since_id=) + 7. Paginated GET /inbox/all (limit, offset, total_inboxes) + 8. High-concurrency workload across hundreds of agents on ThreadingHTTPServer +""" + +import concurrent.futures +import json +import os +import sqlite3 +import sys +import tempfile +import threading +import time +import unittest +import urllib.error +import urllib.request +from http.server import ThreadingHTTPServer + +# Add broker directory to sys.path without mutating global env at import time +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "broker")) + +import broker + +K_ADMIN = "adm" +K_ALICE = "alc" +K_BOB = "bob" + + +class TestBrokerEnterprise(unittest.TestCase): + + server = None + server_thread = None + port = None + _tmpdir = None + _test_db = None + _orig_db = None + _orig_rl = None + _saved_env = None + + @classmethod + def setUpClass(cls): + cls._saved_env = { + k: v for k, v in os.environ.items() if k.startswith("BROKER_") + } + cls._orig_db = broker.BROKER_DB + cls._orig_rl = broker._send_rl + + cls._tmpdir = tempfile.mkdtemp() + cls._test_db = os.path.join(cls._tmpdir, "test_broker_enterprise.db") + + broker.BROKER_DB = cls._test_db + os.environ["BROKER_DB"] = cls._test_db + os.environ["BROKER_KEY_ADMIN"] = K_ADMIN + os.environ["BROKER_KEY_ALICE"] = K_ALICE + os.environ["BROKER_KEY_BOB"] = K_BOB + + broker._init_db() + broker._load_keys() + broker._send_rl = broker.RateLimiter(rate=1000.0, capacity=5000) + + cls.server = ThreadingHTTPServer(("127.0.0.1", 0), broker.Handler) + cls.port = cls.server.server_address[1] + cls.server_thread = threading.Thread( + target=cls.server.serve_forever, daemon=True + ) + cls.server_thread.start() + + @classmethod + def tearDownClass(cls): + if cls.server: + cls.server.shutdown() + cls.server.server_close() + + # Restore environment variables so subsequent test suites remain isolated + for k in list(os.environ.keys()): + if k.startswith("BROKER_"): + os.environ.pop(k, None) + if cls._saved_env: + os.environ.update(cls._saved_env) + + if cls._saved_env and "BROKER_DB" in cls._saved_env: + broker.BROKER_DB = cls._saved_env["BROKER_DB"] + elif cls._orig_db is not None: + broker.BROKER_DB = cls._orig_db + if cls._orig_rl is not None: + broker._send_rl = cls._orig_rl + broker._load_keys() + + if cls._test_db: + for f in [cls._test_db, cls._test_db + "-wal", cls._test_db + "-shm"]: + try: + os.remove(f) + except FileNotFoundError: + pass + if cls._tmpdir: + try: + os.rmdir(cls._tmpdir) + except OSError: + pass + + def setUp(self): + broker.MAX_AGE_SEC = 86400 + broker.RETENTION_SECONDS = 86400 + broker._min_active_ttl = 86400.0 + broker._last_purge_ts = 0.0 + + def _url(self, path: str) -> str: + return f"http://127.0.0.1:{self.port}{path}" + + def _request(self, method: str, path: str, data=None, key: str | None = None): + headers = {"Content-Type": "application/json"} + if key: + headers["Authorization"] = f"Bearer {key}" + body = json.dumps(data).encode() if data is not None else None + req = urllib.request.Request( + self._url(path), data=body, headers=headers, method=method + ) + try: + with urllib.request.urlopen(req, timeout=10) as resp: + return resp.status, json.loads(resp.read().decode()) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read().decode()) + + # ── 1. Dynamic Key Management & Hot-Reload ───────────────────────────── + + def test_register_key_and_rotate_on_the_fly(self): + """Admin can register and rotate user-500 key without restarting broker.""" + # Non-admin cannot register keys + code, _ = self._request( + "POST", + "/admin/register-key", + {"instance": "user-500", "broker_key": "k500-v1"}, + key=K_ALICE, + ) + self.assertEqual(code, 403) + + # Unauthenticated cannot register keys + code, _ = self._request( + "POST", + "/admin/register-key", + {"instance": "user-500", "broker_key": "k500-v1"}, + ) + self.assertEqual(code, 403) + + # Admin registers user-500 + code, body = self._request( + "POST", + "/admin/register-key", + {"instance": "user-500", "broker_key": "k500-v1"}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + self.assertTrue(body["ok"]) + self.assertEqual(body["instance"], "user-500") + + # user-500 can immediately send and receive messages + code, send_body = self._request( + "POST", + "/send", + {"to": "alice", "message": "hello from user-500"}, + key="k500-v1", + ) + self.assertEqual(code, 200) + self.assertEqual(send_body["from"], "user-500") + + # Rotate user-500 key to v2 + code, _ = self._request( + "POST", + "/admin/register-key", + {"instance": "user-500", "broker_key": "k500-v2"}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + + # Old key v1 is rejected + code, _ = self._request("GET", "/inbox", key="k500-v1") + self.assertEqual(code, 403) + + # New key v2 works immediately + code, inbox_body = self._request("GET", "/inbox", key="k500-v2") + self.assertEqual(code, 200) + self.assertIn("inbox", inbox_body) + + # Reject invalid instance names ("*", "a,b", "-bad") with 400 Bad Request + for bad_name in ("*", "a,b", "-bad", "bad space", "a" * 65): + code, _ = self._request( + "POST", + "/admin/register-key", + {"instance": bad_name, "broker_key": "k-bad"}, + key=K_ADMIN, + ) + self.assertEqual(code, 400, f"Expected 400 for invalid instance '{bad_name}'") + + def test_rate_limiter_eviction_and_content_length_validation(self): + """RateLimiter evicts stale buckets above 10000 entries and invalid Content-Length returns 400.""" + rl = broker.RateLimiter(rate=1.0, capacity=5) + now = time.monotonic() + for i in range(10005): + rl._b[f"stale-{i}"] = (4.0, now - 200.0) + rl._b["fresh-user"] = (4.0, now - 10.0) + self.assertTrue(rl.allow("active-user")) + self.assertLess(len(rl._b), 10) + self.assertIn("fresh-user", rl._b) + self.assertIn("active-user", rl._b) + + # Negative / malformed Content-Length returns 400 + import http.client + conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=5) + conn.putrequest("POST", "/send") + conn.putheader("Authorization", f"Bearer {K_ALICE}") + conn.putheader("Content-Type", "application/json") + conn.putheader("Content-Length", "not-an-int") + conn.endheaders() + resp = conn.getresponse() + resp.read() + conn.close() + self.assertEqual(resp.status, 400) + + def test_reload_keys_from_env_and_sqlite(self): + """POST /admin/reload-keys merges keys from os.environ and SQLite broker_keys.""" + os.environ["BROKER_KEY_USER_777"] = "k-777" + try: + code, body = self._request( + "POST", + "/admin/reload-keys", + {}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + self.assertTrue(body["ok"]) + self.assertIn("user-777", body["instances"]) + + # Verify user-777 can authenticate + code, _ = self._request("GET", "/inbox", key="k-777") + self.assertEqual(code, 200) + finally: + os.environ.pop("BROKER_KEY_USER_777", None) + self._request("POST", "/admin/reload-keys", {}, key=K_ADMIN) + + def test_underscore_and_hyphen_env_identities_preserved(self): + """BROKER_KEY_TEAM_ONE preserves 'team_one' identity while BROKER_KEY_USER_1 routes 'user-1'.""" + os.environ["BROKER_KEY_TEAM_ONE"] = "k-team-1" + os.environ["BROKER_KEY_USER_1"] = "k-usr-1" + try: + code, body = self._request( + "POST", + "/admin/reload-keys", + {}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + self.assertIn("team_one", body["instances"]) + self.assertIn("user-1", body["instances"]) + + # 1. Send message to 'team_one' (underscore identity) and read from its inbox + code, send_res = self._request( + "POST", + "/send", + {"to": "team_one", "message": "hello team_one"}, + key=K_ALICE, + ) + self.assertEqual(code, 200) + self.assertEqual(send_res["to"], "team_one") + + code, team_inbox = self._request("GET", "/inbox", key="k-team-1") + self.assertEqual(code, 200) + self.assertTrue( + any(m["message"] == "hello team_one" for m in team_inbox["inbox"]) + ) + + # 2. Send message to 'user-1' (hyphen identity) and read from its inbox + code, send_u1 = self._request( + "POST", + "/send", + {"to": "user-1", "message": "hello user-1"}, + key=K_ALICE, + ) + self.assertEqual(code, 200) + self.assertEqual(send_u1["to"], "user-1") + + code, u1_inbox = self._request("GET", "/inbox", key="k-usr-1") + self.assertEqual(code, 200) + self.assertTrue( + any(m["message"] == "hello user-1" for m in u1_inbox["inbox"]) + ) + finally: + os.environ.pop("BROKER_KEY_TEAM_ONE", None) + os.environ.pop("BROKER_KEY_USER_1", None) + self._request("POST", "/admin/reload-keys", {}, key=K_ADMIN) + + # ── 2. Safe Migration on Legacy Schema ───────────────────────────────── + + def test_legacy_db_migration_adds_priority_column(self): + """_init_db safely migrates an older messages table without priority column.""" + legacy_dir = tempfile.mkdtemp() + legacy_db = os.path.join(legacy_dir, "legacy.db") + try: + conn = sqlite3.connect(legacy_db) + conn.execute(""" + CREATE TABLE messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + sender TEXT NOT NULL, + recipient TEXT NOT NULL, + body TEXT NOT NULL, + ts REAL NOT NULL + ) + """) + conn.execute( + "INSERT INTO messages (sender, recipient, body, ts) VALUES ('alice', 'bob', 'legacy msg', ?)", + (time.time(),), + ) + conn.commit() + conn.close() + + prev_db = broker.BROKER_DB + broker.BROKER_DB = legacy_db + try: + broker._init_db() + conn = sqlite3.connect(legacy_db) + cols = {r[1] for r in conn.execute("PRAGMA table_info(messages)").fetchall()} + self.assertIn("priority", cols) + self.assertIn("delivered", cols) + self.assertIn("created_at", cols) + row = conn.execute("SELECT body, priority FROM messages WHERE recipient='bob'").fetchone() + conn.close() + self.assertEqual(row[0], "legacy msg") + self.assertEqual(row[1], "normal") + finally: + broker.BROKER_DB = prev_db + broker._load_keys() + finally: + for f in [legacy_db, legacy_db + "-wal", legacy_db + "-shm"]: + try: + os.remove(f) + except FileNotFoundError: + pass + try: + os.rmdir(legacy_dir) + except OSError: + pass + + # ── 3. Priorities & Custom TTL Expiry ────────────────────────────────── + + def test_message_priorities(self): + """POST /send supports low, normal, high, urgent priorities and rejects invalid ones.""" + self._request("DELETE", "/inbox", key=K_BOB) + + for prio in ("low", "normal", "high", "urgent"): + code, body = self._request( + "POST", + "/send", + {"to": "bob", "message": f"prio-{prio}", "priority": prio}, + key=K_ALICE, + ) + self.assertEqual(code, 200) + self.assertEqual(body["priority"], prio) + + # Invalid priority returns 400 + code, _ = self._request( + "POST", + "/send", + {"to": "bob", "message": "bad prio", "priority": "critical-invalid"}, + key=K_ALICE, + ) + self.assertEqual(code, 400) + + code, inbox_body = self._request("GET", "/inbox", key=K_BOB) + self.assertEqual(code, 200) + priorities = [m["priority"] for m in inbox_body["inbox"] if m["message"].startswith("prio-")] + self.assertEqual(priorities, ["low", "normal", "high", "urgent"]) + + def test_ttl_expiry(self): + """Messages expire via custom ttl_sec or short RETENTION_SECONDS without delay.""" + self._request("DELETE", "/inbox", key=K_BOB) + + # Send with custom short ttl_sec=0.3 + code, _ = self._request( + "POST", + "/send", + {"to": "bob", "message": "ephemeral-ttl", "ttl_sec": 0.3}, + key=K_ALICE, + ) + self.assertEqual(code, 200) + + # Immediately visible + code, body = self._request("GET", "/inbox?peek=1", key=K_BOB) + self.assertEqual(code, 200) + self.assertTrue(any(m["message"] == "ephemeral-ttl" for m in body["inbox"])) + + # Wait for expiry + time.sleep(0.45) + + # Should be purged & filtered out immediately + code, body = self._request("GET", "/inbox", key=K_BOB) + self.assertEqual(code, 200) + self.assertFalse(any(m["message"] == "ephemeral-ttl" for m in body["inbox"])) + + # Also verify global RETENTION_SECONDS <= 5 triggers immediate purge + broker.RETENTION_SECONDS = 1 + self._request( + "POST", + "/send", + {"to": "bob", "message": "short-retention-msg"}, + key=K_ALICE, + ) + time.sleep(1.15) + code, body = self._request("GET", "/inbox", key=K_BOB) + self.assertEqual(code, 200) + self.assertFalse(any(m["message"] == "short-retention-msg" for m in body["inbox"])) + + # ── 4. Admin Broadcast / Multicast ───────────────────────────────────── + + def test_broadcast_specific_and_wildcard(self): + """Admin can broadcast to a specific list of instances or '*' (all registered).""" + for u in ("bcast-1", "bcast-2"): + self._request( + "POST", + "/admin/register-key", + {"instance": u, "broker_key": f"k-{u}"}, + key=K_ADMIN, + ) + self._request("DELETE", "/inbox", key=f"k-{u}") + self._request("DELETE", "/inbox", key=K_ALICE) + + # Non-admin cannot broadcast + code, _ = self._request( + "POST", + "/broadcast", + {"to": "*", "text": "unauthorized broadcast", "priority": "high"}, + key=K_ALICE, + ) + self.assertEqual(code, 403) + + # Admin multicast to specific list + code, body = self._request( + "POST", + "/broadcast", + {"to": ["bcast-1", "bcast-2"], "text": "targeted maintenance", "priority": "high"}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + self.assertTrue(body["ok"]) + self.assertEqual(body["delivered_count"], 2) + + for u in ("bcast-1", "bcast-2"): + code, inbox = self._request("GET", "/inbox?peek=1", key=f"k-{u}") + self.assertEqual(code, 200) + self.assertTrue( + any(m["message"] == "targeted maintenance" and m["priority"] == "high" for m in inbox["inbox"]) + ) + + # Alice was not in the target list + _, alice_inbox = self._request("GET", "/inbox?peek=1", key=K_ALICE) + self.assertFalse( + any(m["message"] == "targeted maintenance" for m in alice_inbox["inbox"]) + ) + + # Admin wildcard broadcast to "*" + code, body = self._request( + "POST", + "/broadcast", + {"to": "*", "text": "global alert to all", "priority": "urgent"}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + self.assertGreaterEqual(body["delivered_count"], 4) + + _, alice_inbox = self._request("GET", "/inbox?peek=1", key=K_ALICE) + self.assertTrue( + any(m["message"] == "global alert to all" and m["priority"] == "urgent" for m in alice_inbox["inbox"]) + ) + + # ── 5. Inbox Pagination, since_id, peek=1, and /inbox/all ───────────── + + def test_inbox_peek_limit_offset_since_id(self): + """GET /inbox supports peek=1 (without marking delivered=1), limit, offset, since_id.""" + self._request("DELETE", "/inbox", key=K_BOB) + + msg_ids = [] + for i in range(6): + code, res = self._request( + "POST", + "/send", + {"to": "bob", "message": f"page-msg-{i}", "priority": "normal"}, + key=K_ALICE, + ) + self.assertEqual(code, 200) + msg_ids.append(res["id"]) + + # 1. Peek mode: messages returned with delivered=False and DB stays delivered=0 + code, peek_body = self._request( + "GET", "/inbox?peek=1&limit=2&offset=1", key=K_BOB + ) + self.assertEqual(code, 200) + self.assertEqual(peek_body["count"], 2) + self.assertEqual([m["message"] for m in peek_body["inbox"]], ["page-msg-1", "page-msg-2"]) + self.assertTrue(all(m["delivered"] is False for m in peek_body["inbox"])) + + conn = sqlite3.connect(self._test_db) + delivered_flags = [ + r[0] for r in conn.execute("SELECT delivered FROM messages WHERE recipient='bob'").fetchall() + ] + conn.close() + self.assertEqual(delivered_flags, [0, 0, 0, 0, 0, 0]) + + # 2. since_id filtering + code, since_body = self._request( + "GET", f"/inbox?peek=1&since_id={msg_ids[3]}", key=K_BOB + ) + self.assertEqual(code, 200) + self.assertEqual([m["message"] for m in since_body["inbox"]], ["page-msg-4", "page-msg-5"]) + + # 3. Non-peek read marks fetched messages as delivered=1 in DB + code, read_body = self._request( + "GET", "/inbox?limit=3&offset=0", key=K_BOB + ) + self.assertEqual(code, 200) + self.assertEqual(read_body["count"], 3) + self.assertTrue(all(m["delivered"] is True for m in read_body["inbox"])) + + conn = sqlite3.connect(self._test_db) + delivered_after = [ + r[0] + for r in conn.execute( + "SELECT delivered FROM messages WHERE recipient='bob' ORDER BY id ASC" + ).fetchall() + ] + conn.close() + self.assertEqual(delivered_after, [1, 1, 1, 0, 0, 0]) + + def test_inbox_all_pagination_and_total(self): + """GET /inbox/all supports limit, offset, and returns total_inboxes.""" + self._request( + "POST", + "/broadcast", + {"to": "*", "text": "populate all inboxes"}, + key=K_ADMIN, + ) + + code, full = self._request("GET", "/inbox/all", key=K_ADMIN) + self.assertEqual(code, 200) + total = full["total_inboxes"] + self.assertGreaterEqual(total, 3) + self.assertEqual(len(full["inboxes"]), total) + + code, page = self._request( + "GET", "/inbox/all?limit=2&offset=1", key=K_ADMIN + ) + self.assertEqual(code, 200) + self.assertEqual(page["total_inboxes"], total) + self.assertEqual(len(page["inboxes"]), 2) + self.assertEqual(page["limit"], 2) + self.assertEqual(page["offset"], 1) + + # ── 6. High-Concurrency Workload (Hundreds of Concurrent Agents) ─────── + + def test_high_concurrency_agents_send_and_read(self): + """Simulate 150 concurrent agents registering, sending, and reading in parallel.""" + num_agents = 150 + + # Register 150 agents dynamically via admin endpoint + for i in range(num_agents): + code, _ = self._request( + "POST", + "/admin/register-key", + {"instance": f"agent-{i}", "broker_key": f"ka-{i}"}, + key=K_ADMIN, + ) + self.assertEqual(code, 200) + + def agent_worker(idx: int) -> tuple[int, int]: + sender_key = f"ka-{idx}" + target_idx = (idx + 1) % num_agents + target_name = f"agent-{target_idx}" + send_code, _ = self._request( + "POST", + "/send", + { + "to": target_name, + "message": f"concurrent-payload-from-{idx}", + "priority": "high" if idx % 2 == 0 else "normal", + }, + key=sender_key, + ) + read_code, _ = self._request( + "GET", + "/inbox?peek=1&limit=10", + key=sender_key, + ) + return send_code, read_code + + with concurrent.futures.ThreadPoolExecutor(max_workers=50) as pool: + futures = [pool.submit(agent_worker, i) for i in range(num_agents)] + results = [f.result() for f in concurrent.futures.as_completed(futures)] + + for send_code, read_code in results: + self.assertEqual(send_code, 200) + self.assertEqual(read_code, 200) + + # Verify each of the spot-checked agents received their message without any SQLite lock errors + spot_check_indices = [0, 17, 42, 99, 149] + for idx in spot_check_indices: + code, inbox_body = self._request( + "GET", "/inbox", key=f"ka-{idx}" + ) + self.assertEqual(code, 200) + expected_sender = f"agent-{(idx - 1) % num_agents}" + self.assertTrue( + any(m["from"] == expected_sender for m in inbox_body["inbox"]), + f"agent-{idx} missing message from {expected_sender}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_bulk_provisioning.py b/tests/test_bulk_provisioning.py new file mode 100644 index 0000000..1de3560 --- /dev/null +++ b/tests/test_bulk_provisioning.py @@ -0,0 +1,307 @@ +#!/usr/bin/env python3 +""" +Tests for ENT1K-2 Bulk Provisioning & Fleet Management (scripts/bulk-users.py). +""" +from __future__ import annotations + +import csv +import importlib.util +import io +import json +import os +import pathlib +import shutil +import stat +import tempfile +import time +import unittest +from contextlib import redirect_stderr, redirect_stdout + +REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent +SCRIPT_PATH = REPO_ROOT / "scripts" / "bulk-users.py" + + +def _load_bulk_module(): + spec = importlib.util.spec_from_file_location("bulk_users", SCRIPT_PATH) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +bulk_users = _load_bulk_module() + + +class TestBulkProvisioning(unittest.TestCase): + def setUp(self) -> None: + self.tmp_dir = pathlib.Path(tempfile.mkdtemp(prefix="cortex-bulk-test-")) + # Prepare minimal repo layout with real template files + tpl_src = REPO_ROOT / "instances" / "_template" + tpl_dst = self.tmp_dir / "instances" / "_template" + shutil.copytree(tpl_src, tpl_dst, symlinks=False) + + mig_dir = self.tmp_dir / "migrations" + mig_dir.mkdir(parents=True, exist_ok=True) + (mig_dir / "001_init.py").write_text("DESCRIPTION = 'init'\n", encoding="utf-8") + (mig_dir / "024_latest.py").write_text("DESCRIPTION = 'latest'\n", encoding="utf-8") + + root_env = self.tmp_dir / ".env" + root_env.write_text( + "ANTHROPIC_API_KEY=sk-ant-test-key\n" + "QUOTA_ADMIN_TOKEN=admin-test-token-0123456789abcdef\n", + encoding="utf-8", + ) + os.chmod(root_env, 0o600) + + def tearDown(self) -> None: + shutil.rmtree(self.tmp_dir, ignore_errors=True) + + def _run_cli(self, argv: list[str]) -> tuple[int, str, str]: + out_buf = io.StringIO() + err_buf = io.StringIO() + with redirect_stdout(out_buf), redirect_stderr(err_buf): + rc = bulk_users.main(["--root", str(self.tmp_dir), *argv]) + return rc, out_buf.getvalue(), err_buf.getvalue() + + def test_bulk_import_120_users_sharding_and_ports(self) -> None: + """Import 120 users in one call; verify speed, shards, ports, and 0600 files.""" + csv_path = self.tmp_dir / "users_120.csv" + with open(csv_path, "w", encoding="utf-8", newline="") as f: + writer = csv.DictWriter( + f, + fieldnames=["name", "bot_token", "full_name", "tg_id", "department", "quota_limit"], + ) + writer.writeheader() + for i in range(1, 121): + writer.writerow( + { + "name": f"emp-{i:04d}", + "bot_token": f"700{i:04d}:AAtestBotToken_{i:04d}", + "full_name": f"Employee #{i}", + "tg_id": str(1000000 + i), + "department": "eng" if i <= 60 else "ops", + "quota_limit": str(500_000 if i % 2 == 0 else 1_000_000), + } + ) + + t0 = time.monotonic() + rc, out, err = self._run_cli( + ["import", "--file", str(csv_path), "--shard-size", "50", "--base-port", "18790"] + ) + elapsed = time.monotonic() - t0 + self.assertEqual(rc, 0, f"Unexpected stderr: {err}") + self.assertLess(elapsed, 10.0, f"Bulk import of 120 users took too long: {elapsed:.2f}s") + + summary = json.loads(out) + self.assertEqual(summary["created_count"], 120) + self.assertEqual(summary["skipped_count"], 0) + + # Verify unique ports 18790..18909 + ports = [item["ui_port"] for item in summary["created"]] + self.assertEqual(len(ports), len(set(ports))) + self.assertEqual(ports[0], 18790) + self.assertEqual(ports[-1], 18909) + + # Verify shard distribution (50 per shard -> shard-0: 50, shard-1: 50, shard-2: 20) + shards = [item["shard"] for item in summary["created"]] + self.assertEqual(shards.count("shard-0"), 50) + self.assertEqual(shards.count("shard-1"), 50) + self.assertEqual(shards.count("shard-2"), 20) + + # Spot-check instance files and permissions + sample_dir = self.tmp_dir / "instances" / "emp-0001" + self.assertTrue((sample_dir / "hermes_data" / "ssh").is_dir()) + self.assertTrue((sample_dir / "workspace" / "skills").is_dir()) + self.assertTrue((sample_dir / "workspace" / "memory").is_dir()) + + for rel in (".env", "hermes_data/.env", "hermes_data/config.yaml", "hermes_data/SOUL.md"): + p = sample_dir / rel + self.assertTrue(p.is_file(), f"Missing {rel}") + mode = stat.S_IMODE(p.stat().st_mode) + self.assertEqual(mode, 0o600, f"Expected 0600 on {rel}, got {oct(mode)}") + + user_md = (sample_dir / "workspace" / "USER.md").read_text(encoding="utf-8") + self.assertIn("Employee #1", user_md) + self.assertNotIn("{{FULL_NAME}}", user_md) + + migs = (sample_dir / "workspace" / ".migrations_applied").read_text(encoding="utf-8").splitlines() + self.assertEqual(migs, ["001_init", "024_latest"]) + + # Verify root .env has all 120 users' keys and 0600 permissions + root_env_path = self.tmp_dir / ".env" + self.assertEqual(stat.S_IMODE(root_env_path.stat().st_mode), 0o600) + root_env_text = root_env_path.read_text(encoding="utf-8") + self.assertIn("BROKER_KEY_EMP_0001=", root_env_text) + self.assertIn("QUOTA_KEY_EMP_0120=quota-emp-0120-", root_env_text) + self.assertIn("QUOTA_LIMIT_EMP_0002=500000", root_env_text) + self.assertIn("QUOTA_DEPT_EMP_0001=eng", root_env_text) + self.assertIn("QUOTA_DEPT_EMP_0120=ops", root_env_text) + + # Verify docker-compose.override.yml contains shard profiles and shard bridge networks + override_text = (self.tmp_dir / "docker-compose.override.yml").read_text(encoding="utf-8") + self.assertIn("assistant-emp-0001:", override_text) + self.assertIn("assistant-emp-0120:", override_text) + self.assertIn('- "shard-0"', override_text) + self.assertIn('- "shard-1"', override_text) + self.assertIn('- "shard-2"', override_text) + self.assertIn("corp-shard-0:", override_text) + self.assertIn("corp-shard-1:", override_text) + self.assertIn("corp-shard-2:", override_text) + + def test_idempotency_and_incremental_import(self) -> None: + """Re-importing existing users skips them safely while adding new ones.""" + batch1 = [ + {"name": "user-a", "bot_token": "701:AAtokenA", "full_name": "User A", "tg_id": "111"}, + {"name": "user-b", "bot_token": "702:AAtokenB", "full_name": "User B", "tg_id": "222"}, + ] + json_path = self.tmp_dir / "batch.json" + json_path.write_text(json.dumps(batch1), encoding="utf-8") + + rc1, out1, _ = self._run_cli(["import", "--file", str(json_path)]) + self.assertEqual(rc1, 0) + self.assertEqual(json.loads(out1)["created_count"], 2) + + orig_env_a = (self.tmp_dir / "instances" / "user-a" / "hermes_data" / ".env").read_text(encoding="utf-8") + + batch2 = batch1 + [ + {"name": "user-c", "bot_token": "703:AAtokenC", "full_name": "User C", "tg_id": "333"}, + ] + json_path.write_text(json.dumps({"users": batch2}), encoding="utf-8") + + rc2, out2, _ = self._run_cli(["import", "--file", str(json_path)]) + self.assertEqual(rc2, 0) + res2 = json.loads(out2) + self.assertEqual(res2["created_count"], 1) + self.assertEqual(res2["skipped"], ["user-a", "user-b"]) + self.assertEqual(res2["created"][0]["name"], "user-c") + self.assertEqual(res2["created"][0]["ui_port"], 18792) + + # Ensure user-a secrets were not overwritten + after_env_a = (self.tmp_dir / "instances" / "user-a" / "hermes_data" / ".env").read_text(encoding="utf-8") + self.assertEqual(orig_env_a, after_env_a) + + def test_status_and_export_commands(self) -> None: + """Verify status integrity checks and export secret redaction / 0600 mode.""" + users = [ + { + "name": "alice-1", + "bot_token": "801:AAaliceSecret", + "full_name": "Alice One", + "tg_id": "9001", + "department": "finance", + "quota_limit": 750000, + }, + { + "name": "bob-2", + "bot_token": "802:AAbobSecret", + "full_name": "Bob Two", + "tg_id": "9002", + "department": "legal", + "quota_limit": 1250000, + }, + ] + src_file = self.tmp_dir / "users.json" + src_file.write_text(json.dumps(users), encoding="utf-8") + self.assertEqual(self._run_cli(["import", "--file", str(src_file), "--shard-size", "1"])[0], 0) + + # Status when healthy + rc_st, out_st, _ = self._run_cli(["status"]) + self.assertEqual(rc_st, 0) + status_data = json.loads(out_st) + self.assertEqual(status_data["total_instances"], 2) + self.assertEqual(status_data["healthy_instances"], 2) + self.assertEqual(status_data["shards"], {"shard-0": 1, "shard-1": 1}) + + # Export without secrets + manifest_pub = self.tmp_dir / "manifest-public.json" + rc_ex, _, _ = self._run_cli(["export", "--output", str(manifest_pub)]) + self.assertEqual(rc_ex, 0) + pub_data = json.loads(manifest_pub.read_text(encoding="utf-8")) + self.assertEqual(pub_data["instance_count"], 2) + self.assertNotIn("secrets", pub_data["instances"][0]) + self.assertNotIn("AAaliceSecret", manifest_pub.read_text(encoding="utf-8")) + + # Export with secrets (must have 0600 permissions) + manifest_sec = self.tmp_dir / "manifest-secret.json" + rc_ex_sec, _, _ = self._run_cli(["export", "--output", str(manifest_sec), "--include-secrets"]) + self.assertEqual(rc_ex_sec, 0) + self.assertEqual(stat.S_IMODE(manifest_sec.stat().st_mode), 0o600) + sec_data = json.loads(manifest_sec.read_text(encoding="utf-8")) + self.assertEqual(sec_data["instances"][0]["secrets"]["bot_token"], "801:AAaliceSecret") + + # Break one file and verify status returns non-zero and flags degraded instance + (self.tmp_dir / "instances" / "bob-2" / "hermes_data" / "config.yaml").unlink() + rc_deg, out_deg, _ = self._run_cli(["status"]) + self.assertEqual(rc_deg, 1) + deg_data = json.loads(out_deg) + self.assertEqual(deg_data["degraded_instances"], 1) + self.assertEqual(deg_data["degraded_details"][0]["name"], "bob-2") + self.assertIn("hermes_config", deg_data["degraded_details"][0]["missing"]) + + def test_rejects_invalid_records_and_symlinks(self) -> None: + """Invalid records or symlinks must be rejected before modifying state.""" + invalid_batches = [ + [{"name": "admin", "bot_token": "123:AAok", "full_name": "Admin", "tg_id": "111"}], + [{"name": "../evil", "bot_token": "123:AAok", "full_name": "Evil", "tg_id": "111"}], + [{"name": "UpperCase", "bot_token": "123:AAok", "full_name": "Bad", "tg_id": "111"}], + [{"name": "valid-1", "bot_token": "not-a-token", "full_name": "Bad Tok", "tg_id": "111"}], + [{"name": "valid-1", "bot_token": "123:AAok", "full_name": "Bad TG", "tg_id": "abc"}], + [{"name": "valid-1", "bot_token": "123:AAok", "full_name": "Port 18789", "tg_id": "111", "ui_port": 18789}], + [ + {"name": "user-1", "bot_token": "123:AAok1", "full_name": "U1", "tg_id": "111"}, + {"name": "user_1", "bot_token": "124:AAok2", "full_name": "U2", "tg_id": "222"}, + ], + ] + bad_file = self.tmp_dir / "bad.json" + for batch in invalid_batches: + bad_file.write_text(json.dumps(batch), encoding="utf-8") + rc, _, err = self._run_cli(["import", "--file", str(bad_file)]) + self.assertNotEqual(rc, 0, f"Batch should have failed: {batch}") + self.assertIn("Validation error", err) + + # Collision with an already-existing instance directory on normalized env suffix (user-1 vs user_1) + first_file = self.tmp_dir / "first_u1.json" + first_file.write_text( + json.dumps([{"name": "user-1", "bot_token": "123:AAok1", "full_name": "U1", "tg_id": "111"}]), + encoding="utf-8", + ) + self.assertEqual(self._run_cli(["import", "--file", str(first_file)])[0], 0) + + second_file = self.tmp_dir / "second_u1.json" + second_file.write_text( + json.dumps([{"name": "user_1", "bot_token": "124:AAok2", "full_name": "U2", "tg_id": "222"}]), + encoding="utf-8", + ) + rc_col, _, err_col = self._run_cli(["import", "--file", str(second_file)]) + self.assertNotEqual(rc_col, 0) + self.assertIn("collides on normalized env suffix", err_col) + + # Symlink attack inside template workspace must be rejected + good_file = self.tmp_dir / "good.json" + good_file.write_text( + json.dumps([{"name": "safe-user", "bot_token": "123:AAok", "full_name": "Safe", "tg_id": "111"}]), + encoding="utf-8", + ) + symlink_in_tpl = self.tmp_dir / "instances" / "_template" / "workspace" / "evil_link.md" + symlink_in_tpl.symlink_to("/etc/passwd") + try: + rc_sym, _, err_sym = self._run_cli(["import", "--file", str(good_file)]) + self.assertNotEqual(rc_sym, 0) + self.assertIn("Symlink rejected", err_sym) + self.assertFalse((self.tmp_dir / "instances" / "safe-user").exists()) + finally: + symlink_in_tpl.unlink(missing_ok=True) + + # Symlink at target instance path must also be rejected + target_sym = self.tmp_dir / "instances" / "safe-user" + target_sym.symlink_to(self.tmp_dir) + try: + rc_tsym, _, err_tsym = self._run_cli(["import", "--file", str(good_file)]) + self.assertNotEqual(rc_tsym, 0) + self.assertIn("symlink", err_tsym.lower()) + finally: + target_sym.unlink(missing_ok=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_quota_enterprise.py b/tests/test_quota_enterprise.py new file mode 100644 index 0000000..4e0825f --- /dev/null +++ b/tests/test_quota_enterprise.py @@ -0,0 +1,504 @@ +#!/usr/bin/env python3 +""" +Enterprise tests for quota-proxy (ENT1K-1: 1000-user scale readiness). + +Covers: +1. SSE Anthropic streaming token accounting (message_start + message_delta) +2. SSE OpenAI streaming token accounting (usage chunk with prompt_tokens + completion_tokens) +3. Fragmented SSE chunk parsing (SSEUsageParser / _parse_sse_usage) +4. Dynamic key registration & rotation without restart (POST /quota/register-key) +5. Hot-reload of env + SQLite keys (POST /quota/reload-keys) +6. Hierarchical department quota enforcement (POST /quota/set-department-limit) +7. Report filtering (?department=, ?status=) and pagination (?limit=&offset=) +8. High-concurrency parallel requests across hundreds of users on ThreadingHTTPServer +""" + +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import importlib.util +import json +import os +import tempfile +import threading +import unittest +import urllib.error +import urllib.request + + +class MockUpstreamHandler(BaseHTTPRequestHandler): + """Mock upstream supporting JSON, Anthropic SSE, and OpenAI SSE responses.""" + + def log_message(self, fmt, *args): + pass + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + raw = self.rfile.read(length) if length else b"{}" + try: + payload = json.loads(raw) + except Exception: + payload = {} + + mode = payload.get("mock_mode", "json") + if mode == "sse_anthropic": + in_tok = int(payload.get("in_tok", 120)) + out_tok = int(payload.get("out_tok", 45)) + events = [ + 'event: message_start\n', + f'data: {{"type":"message_start","message":{{"id":"msg_1","usage":{{"input_tokens":{in_tok},"output_tokens":1}}}}}}\n\n', + 'event: content_block_delta\n', + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}\n\n', + 'event: message_delta\n', + f'data: {{"type":"message_delta","delta":{{"stop_reason":"end_turn"}},"usage":{{"output_tokens":{out_tok}}}}}\n\n', + 'event: message_stop\n', + 'data: {"type":"message_stop"}\n\n', + ] + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.end_headers() + for ev in events: + self.wfile.write(ev.encode("utf-8")) + self.wfile.flush() + return + + if mode == "sse_openai": + in_tok = int(payload.get("in_tok", 80)) + out_tok = int(payload.get("out_tok", 30)) + events = [ + 'data: {"id":"chatcmpl-1","choices":[{"index":0,"delta":{"content":"Hi"}}]}\n\n', + f'data: {{"id":"chatcmpl-1","choices":[],"usage":{{"prompt_tokens":{in_tok},"completion_tokens":{out_tok},"total_tokens":{in_tok + out_tok}}}}}\n\n', + 'data: [DONE]\n\n', + ] + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + for ev in events: + self.wfile.write(ev.encode("utf-8")) + self.wfile.flush() + return + + in_tok = int(payload.get("in_tok", 10)) + out_tok = int(payload.get("out_tok", 20)) + body = json.dumps({ + "id": "msg_json", + "type": "message", + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + }).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +class TestQuotaEnterprise(unittest.TestCase): + upstream_server = None + proxy_server = None + proxy_mod = None + tmpdir = None + db_path = None + admin_token = "adm-test-token" + + @classmethod + def setUpClass(cls): + cls._saved_env = dict(os.environ) + cls.tmpdir = tempfile.mkdtemp(prefix="quota_ent_") + cls.db_path = os.path.join(cls.tmpdir, "quota.db") + + # Start mock upstream server first + cls.upstream_server = ThreadingHTTPServer(("127.0.0.1", 0), MockUpstreamHandler) + cls.upstream_port = cls.upstream_server.server_address[1] + cls.upstream_thread = threading.Thread( + target=cls.upstream_server.serve_forever, daemon=True + ) + cls.upstream_thread.start() + + # Configure environment before loading quota-proxy module + os.environ["QUOTA_DB"] = cls.db_path + os.environ["UPSTREAM_URL"] = f"http://127.0.0.1:{cls.upstream_port}" + os.environ["UPSTREAM_API_KEY"] = "clip-mock-upstream-key" + os.environ["QUOTA_ADMIN_TOKEN"] = cls.admin_token + os.environ["QUOTA_DEFAULT_MONTHLY"] = "1000000" + os.environ["QUOTA_KEY_USER_1"] = "quota-user-1-initkey" + + proxy_path = os.path.join( + os.path.dirname(__file__), "..", "quota-proxy", "proxy.py" + ) + spec = importlib.util.spec_from_file_location("quota_proxy_ent", proxy_path) + cls.proxy_mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(cls.proxy_mod) + + # Raise rate limits for high-concurrency enterprise tests + cls.proxy_mod._admin_rl = cls.proxy_mod.RateLimiter(rate=10000.0, capacity=10000) + cls.proxy_mod._proxy_rl = cls.proxy_mod.RateLimiter(rate=10000.0, capacity=10000) + + cls.proxy_server = ThreadingHTTPServer(("127.0.0.1", 0), cls.proxy_mod.Handler) + cls.proxy_port = cls.proxy_server.server_address[1] + cls.proxy_thread = threading.Thread( + target=cls.proxy_server.serve_forever, daemon=True + ) + cls.proxy_thread.start() + + @classmethod + def tearDownClass(cls): + if cls.proxy_server: + cls.proxy_server.shutdown() + cls.proxy_server.server_close() + if cls.upstream_server: + cls.upstream_server.shutdown() + cls.upstream_server.server_close() + if cls.proxy_mod and getattr(cls.proxy_mod, "_conn", None): + cls.proxy_mod._conn.close() + for suffix in ("", "-wal", "-shm"): + p = cls.db_path + suffix + if os.path.exists(p): + try: + os.remove(p) + except OSError: + pass + if cls.tmpdir and os.path.exists(cls.tmpdir): + try: + os.rmdir(cls.tmpdir) + except OSError: + pass + if getattr(cls, "_saved_env", None) is not None: + os.environ.clear() + os.environ.update(cls._saved_env) + + def _url(self, path: str) -> str: + return f"http://127.0.0.1:{self.proxy_port}{path}" + + def _admin_req(self, method: str, path: str, body: dict = None) -> tuple[int, dict]: + headers = { + "Authorization": f"Bearer {self.admin_token}", + "Content-Type": "application/json", + } + data = json.dumps(body).encode("utf-8") if body is not None else None + req = urllib.request.Request(self._url(path), data=data, headers=headers, method=method) + try: + with urllib.request.urlopen(req, timeout=10) as r: + return r.status, json.loads(r.read().decode("utf-8")) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read().decode("utf-8")) + + def _proxy_req(self, quota_key: str, payload: dict) -> tuple[int, str, bytes]: + headers = { + "x-api-key": quota_key, + "Content-Type": "application/json", + } + data = json.dumps(payload).encode("utf-8") + req = urllib.request.Request( + self._url("/v1/messages"), data=data, headers=headers, method="POST" + ) + try: + with urllib.request.urlopen(req, timeout=10) as r: + return r.status, r.headers.get("Content-Type", ""), r.read() + except urllib.error.HTTPError as e: + return e.code, e.headers.get("Content-Type", ""), e.read() + + # ── 1. SSE Anthropic Streaming Accounting ──────────────────────────── + + def test_01_sse_anthropic_streaming_usage_accounting(self): + """Anthropic text/event-stream increments input & output tokens from message_start/message_delta.""" + self._admin_req("POST", "/quota/reset", {"instance": "user-1"}) + status, ct, raw = self._proxy_req( + "quota-user-1-initkey", + {"mock_mode": "sse_anthropic", "in_tok": 210, "out_tok": 90}, + ) + self.assertEqual(status, 200) + self.assertIn("text/event-stream", ct) + self.assertIn(b"event: message_start", raw) + self.assertIn(b"event: message_delta", raw) + + usage = self.proxy_mod._get_usage("user-1") + self.assertEqual(usage["input"], 210) + self.assertEqual(usage["output"], 90) + self.assertEqual(usage["total"], 300) + self.assertEqual(usage["requests"], 1) + + # ── 2. SSE OpenAI Streaming Accounting ─────────────────────────────── + + def test_02_sse_openai_streaming_usage_accounting(self): + """OpenAI-compatible SSE stream records prompt_tokens and completion_tokens.""" + self._admin_req("POST", "/quota/reset", {"instance": "user-1"}) + status, ct, raw = self._proxy_req( + "quota-user-1-initkey", + {"mock_mode": "sse_openai", "in_tok": 140, "out_tok": 60}, + ) + self.assertEqual(status, 200) + self.assertIn("text/event-stream", ct) + self.assertIn(b"[DONE]", raw) + + usage = self.proxy_mod._get_usage("user-1") + self.assertEqual(usage["input"], 140) + self.assertEqual(usage["output"], 60) + self.assertEqual(usage["total"], 200) + self.assertEqual(usage["requests"], 1) + + # ── 3. Fragmented SSE Chunk Parser ─────────────────────────────────── + + def test_03_sse_parser_fragmented_chunks(self): + """SSEUsageParser handles arbitrary byte chunk boundaries without losing usage.""" + parser = self.proxy_mod.SSEUsageParser() + stream = ( + b"event: message_start\r\n" + b'data: {"type":"message_start","message":{"usage":{"input_tokens":333,"output_tokens":1}}}\r\n\r\n' + b"event: message_delta\r\n" + b'data: {"type":"message_delta","usage":{"output_tokens":77}}\r\n\r\n' + ) + # Feed 7 bytes at a time to split lines across chunks + for i in range(0, len(stream), 7): + parser.feed_chunk(stream[i:i + 7]) + parser.finish() + self.assertEqual(parser.input_tokens, 333) + self.assertEqual(parser.output_tokens, 77) + + # ── 4. Dynamic Key Registration & Rotation Without Restart ─────────── + + def test_04_dynamic_register_key_and_rotation(self): + """POST /quota/register-key registers and rotates keys in SQLite + memory immediately.""" + code, resp = self._admin_req( + "POST", + "/quota/register-key", + { + "instance": "user-500", + "quota_key": "quota-user-500-first", + "limit": 50000, + "department": "engineering", + }, + ) + self.assertEqual(code, 200) + self.assertTrue(resp["ok"]) + self.assertEqual(resp["instance"], "user-500") + self.assertEqual(resp["department"], "engineering") + self.assertEqual(resp["limit"], 50000) + + # Immediately usable without restart + status, _, _ = self._proxy_req( + "quota-user-500-first", {"mock_mode": "json", "in_tok": 25, "out_tok": 75} + ) + self.assertEqual(status, 200) + self.assertEqual(self.proxy_mod._get_usage("user-500")["total"], 100) + + # Rotate key for user-500: old key must be invalidated, new key must work + code, _ = self._admin_req( + "POST", + "/quota/register-key", + { + "instance": "user-500", + "quota_key": "quota-user-500-rotated", + "department": "engineering", + }, + ) + self.assertEqual(code, 200) + + # Reject invalid instance name or invalid key_hash with 400 + bad_inst_code, _ = self._admin_req( + "POST", + "/quota/register-key", + {"instance": "../bad", "quota_key": "quota-bad"}, + ) + self.assertEqual(bad_inst_code, 400) + bad_hash_code, _ = self._admin_req( + "POST", + "/quota/register-key", + {"instance": "user-501", "key_hash": "not-a-sha256-hex"}, + ) + self.assertEqual(bad_hash_code, 400) + + old_status, _, _ = self._proxy_req( + "quota-user-500-first", {"mock_mode": "json", "in_tok": 10, "out_tok": 10} + ) + self.assertEqual(old_status, 401) + + new_status, _, _ = self._proxy_req( + "quota-user-500-rotated", {"mock_mode": "json", "in_tok": 50, "out_tok": 50} + ) + self.assertEqual(new_status, 200) + self.assertEqual(self.proxy_mod._get_usage("user-500")["total"], 200) + + # ── 5. Hot-Reload Keys from Env & SQLite ───────────────────────────── + + def test_05_reload_keys_endpoint(self): + """POST /quota/reload-keys reloads keys added to os.environ and SQLite.""" + os.environ["QUOTA_KEY_USER_777"] = "quota-user-777-hotreload" + try: + code, body = self._admin_req("POST", "/quota/reload-keys", {}) + self.assertEqual(code, 200) + self.assertTrue(body["ok"]) + self.assertIn("user-777", body["instances"]) + self.assertIn("user-500", body["instances"]) + + status, _, _ = self._proxy_req( + "quota-user-777-hotreload", {"mock_mode": "json", "in_tok": 15, "out_tok": 5} + ) + self.assertEqual(status, 200) + finally: + os.environ.pop("QUOTA_KEY_USER_777", None) + + # ── 6. Department Quota Enforcement ────────────────────────────────── + + def test_06_department_quota_enforcement(self): + """Department limit blocks all instances in the department once collective usage >= limit.""" + for idx in (1, 2): + self._admin_req( + "POST", + "/quota/register-key", + { + "instance": f"eng-{idx}", + "quota_key": f"quota-eng-{idx}-key", + "limit": 10000, + "department": "ml-research", + }, + ) + self._admin_req("POST", "/quota/reset", {"instance": f"eng-{idx}"}) + + self._admin_req( + "POST", + "/quota/register-key", + { + "instance": "sales-1", + "quota_key": "quota-sales-1-key", + "limit": 10000, + "department": "sales", + }, + ) + self._admin_req("POST", "/quota/reset", {"instance": "sales-1"}) + + # Set ml-research department limit to 500 tokens (while individual limits are 10,000) + code, body = self._admin_req( + "POST", + "/quota/set-department-limit", + {"department": "ml-research", "limit": 500}, + ) + self.assertEqual(code, 200) + self.assertEqual(body["limit"], 500) + + # eng-1 uses 300 tokens -> allowed + s1, _, _ = self._proxy_req( + "quota-eng-1-key", {"mock_mode": "json", "in_tok": 200, "out_tok": 100} + ) + self.assertEqual(s1, 200) + + # eng-2 uses 250 tokens -> allowed (total becomes 550 >= 500) + s2, _, _ = self._proxy_req( + "quota-eng-2-key", {"mock_mode": "json", "in_tok": 150, "out_tok": 100} + ) + self.assertEqual(s2, 200) + + # Next request from eng-1 must be blocked by department limit (429) + s3, _, raw_err = self._proxy_req( + "quota-eng-1-key", {"mock_mode": "json", "in_tok": 10, "out_tok": 10} + ) + self.assertEqual(s3, 429) + err_payload = json.loads(raw_err.decode("utf-8")) + self.assertEqual(err_payload["error"], "quota_exceeded") + self.assertIn("департамента ml-research", err_payload["message"]) + + # sales-1 in another department is unaffected + s4, _, _ = self._proxy_req( + "quota-sales-1-key", {"mock_mode": "json", "in_tok": 100, "out_tok": 100} + ) + self.assertEqual(s4, 200) + + # ── 7. Report Filtering & Pagination ───────────────────────────────── + + def test_07_report_filtering_and_pagination(self): + """GET /quota/report supports department, status, limit, offset while preserving schema.""" + # Filter by department=ml-research + code, rep = self._admin_req("GET", "/quota/report?department=ml-research") + self.assertEqual(code, 200) + self.assertIn("month", rep) + self.assertIn("usage", rep) + self.assertIn("limits", rep) + self.assertIn("total_instances", rep) + self.assertIn("departments", rep) + self.assertEqual(rep["total_instances"], 2) + self.assertEqual(len(rep["usage"]), 2) + for item in rep["usage"]: + self.assertEqual(item["department"], "ml-research") + self.assertEqual(rep["departments"]["ml-research"]["total_tokens"], 550) + self.assertEqual(rep["departments"]["ml-research"]["status"], "❌ exceeded") + + # Pagination: limit=1&offset=0 and limit=1&offset=1 + _, page1 = self._admin_req("GET", "/quota/report?department=ml-research&limit=1&offset=0") + _, page2 = self._admin_req("GET", "/quota/report?department=ml-research&limit=1&offset=1") + self.assertEqual(page1["total_instances"], 2) + self.assertEqual(len(page1["usage"]), 1) + self.assertEqual(len(page2["usage"]), 1) + self.assertNotEqual(page1["usage"][0]["instance"], page2["usage"][0]["instance"]) + + # ── 8. High-Concurrency 250-User Parallel Stress Test ──────────────── + + def test_08_high_concurrency_parallel_users(self): + """250 distinct users registering and sending concurrent requests with zero lock failures.""" + num_users = 250 + + def register_and_call(i: int): + inst = f"corp-u-{i:04d}" + key = f"quota-{inst}-secret-{i}" + c1, _ = self._admin_req( + "POST", + "/quota/register-key", + { + "instance": inst, + "quota_key": key, + "limit": 100000, + "department": f"dept-{i % 5}", + }, + ) + if c1 != 200: + return False + c2, _, _ = self._proxy_req( + key, {"mock_mode": "sse_anthropic", "in_tok": 10, "out_tok": 15} + ) + return c2 == 200 + + with ThreadPoolExecutor(max_workers=32) as pool: + results = list(pool.map(register_and_call, range(num_users))) + + self.assertTrue(all(results), "All 250 concurrent user flows must succeed") + + code, rep = self._admin_req("GET", "/quota/report?limit=50&offset=0") + self.assertEqual(code, 200) + self.assertGreaterEqual(rep["total_instances"], num_users) + self.assertEqual(len(rep["usage"]), 50) + + # ── 9. Concurrent Same-Instance Quota Admission Barrier ────────────── + + def test_concurrent_same_instance_quota_admission_barrier(self): + """8 concurrent requests on the same instance with 30-token limit: 1 gets 200, 7 get 429.""" + inst = "barrier-user-1" + key = "quota-barrier-user-1-k" + code, _ = self._admin_req( + "POST", + "/quota/register-key", + { + "instance": inst, + "quota_key": key, + "limit": 30, + "department": "barrier-dept", + }, + ) + self.assertEqual(code, 200) + self._admin_req("POST", "/quota/reset", {"instance": inst}) + + barrier = threading.Barrier(8) + + def worker(_idx: int) -> int: + barrier.wait(timeout=5) + status, _, _ = self._proxy_req( + key, {"mock_mode": "json", "in_tok": 15, "out_tok": 15} + ) + return status + + with ThreadPoolExecutor(max_workers=8) as pool: + statuses = list(pool.map(worker, range(8))) + + self.assertEqual(statuses.count(200), 1) + self.assertEqual(statuses.count(429), 7) + self.assertEqual(self.proxy_mod._get_usage(inst)["total"], 30) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_scale_1000_users.py b/tests/test_scale_1000_users.py new file mode 100644 index 0000000..66f5e85 --- /dev/null +++ b/tests/test_scale_1000_users.py @@ -0,0 +1,682 @@ +#!/usr/bin/env python3 +""" +ENT1K-3: 1000-User Scale & Concurrency Test Suite for CortexForge. + +Covers: +1. test_1000_users_quota_proxy_concurrent_and_sse: + - Registers 1,000 users (user-1..user-1000 across dept-0..dept-9) in a live quota-proxy + - Runs 1,000 concurrent requests across ThreadPoolExecutor(max_workers=64) mixing + standard JSON responses, Anthropic SSE streams, and OpenAI SSE streams + - Verifies 0 HTTP 500s and 0 sqlite3.OperationalError ("database is locked") + - Verifies exact token accounting across all 1,000 users and all 10 departments + - Verifies report pagination (GET /quota/report?limit=100&offset=...) and + department filtering (?department=dept-0) + - Verifies department quota enforcement (POST /quota/set-department-limit) + +2. test_1000_agents_broker_concurrent_messaging_and_broadcast: + - Registers 1,000 agents (user-1..user-1000) in a live message-broker + - Runs 1,100+ concurrent P2P messages across ThreadPoolExecutor(max_workers=64) + plus admin POST /broadcast to "*" (all 1,000 agents) + - Verifies GET /inbox/all?limit=100&offset=0 pagination (total_inboxes == 1000) + - Verifies GET /inbox?peek=1 non-destructive reads, priority ordering, and zero + lost messages or database lock errors + +3. test_bulk_provision_1000_users_manifest: + - Generates a 1,000-user manifest in a temporary repository root + - Executes scripts/bulk-users.py import --file ... --shard-size 100 + - Verifies all 1,000 instances, .env keys, and 10 compose shards (shard-0..shard-9) + are created rapidly with zero port collisions +""" + +from __future__ import annotations + +from concurrent.futures import ThreadPoolExecutor +from contextlib import redirect_stderr, redirect_stdout +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +import importlib.util +import io +import json +import os +import pathlib +import shutil +import sqlite3 +import stat +import sys +import tempfile +import threading +import time +import unittest +import urllib.error +import urllib.request + +REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent +PROXY_PATH = REPO_ROOT / "quota-proxy" / "proxy.py" +BROKER_DIR = REPO_ROOT / "broker" +BULK_SCRIPT_PATH = REPO_ROOT / "scripts" / "bulk-users.py" + + +class HighConcurrencyHTTPServer(ThreadingHTTPServer): + """ThreadingHTTPServer with a larger socket backlog for 64+ concurrent workers.""" + + request_queue_size = 512 + allow_reuse_address = True + daemon_threads = True + + +class MockScaleUpstreamHandler(BaseHTTPRequestHandler): + """Mock LLM upstream supporting JSON, Anthropic SSE, and OpenAI SSE streams.""" + + def log_message(self, fmt, *args): + pass + + def do_POST(self): + length = int(self.headers.get("Content-Length", 0)) + raw = self.rfile.read(length) if length else b"{}" + try: + payload = json.loads(raw) + except Exception: + payload = {} + + mode = payload.get("mock_mode", "json") + in_tok = int(payload.get("in_tok", 10)) + out_tok = int(payload.get("out_tok", 20)) + + if mode == "sse_anthropic": + stream_bytes = ( + "event: message_start\n" + f'data: {{"type":"message_start","message":{{"id":"msg_a","usage":{{"input_tokens":{in_tok},"output_tokens":1}}}}}}\n\n' + "event: content_block_delta\n" + 'data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"ok"}}\n\n' + "event: message_delta\n" + f'data: {{"type":"message_delta","delta":{{"stop_reason":"end_turn"}},"usage":{{"output_tokens":{out_tok}}}}}\n\n' + "event: message_stop\n" + 'data: {"type":"message_stop"}\n\n' + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "text/event-stream; charset=utf-8") + self.end_headers() + self.wfile.write(stream_bytes) + self.wfile.flush() + return + + if mode == "sse_openai": + stream_bytes = ( + 'data: {"id":"chatcmpl-scale","choices":[{"index":0,"delta":{"content":"ok"}}]}\n\n' + f'data: {{"id":"chatcmpl-scale","choices":[],"usage":{{"prompt_tokens":{in_tok},"completion_tokens":{out_tok},"total_tokens":{in_tok + out_tok}}}}}\n\n' + "data: [DONE]\n\n" + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + self.wfile.write(stream_bytes) + self.wfile.flush() + return + + body = json.dumps( + { + "id": "msg_json_scale", + "type": "message", + "usage": {"input_tokens": in_tok, "output_tokens": out_tok}, + } + ).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +def _load_module(name: str, path: pathlib.Path): + spec = importlib.util.spec_from_file_location(name, path) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +class TestScale1000Users(unittest.TestCase): + """End-to-end 1000-user concurrency and scale verification (ENT1K-3).""" + + # ── Test 1: 1000 Concurrent Users in Quota-Proxy (JSON + SSE) ───────── + + def test_1000_users_quota_proxy_concurrent_and_sse(self): + saved_env = dict(os.environ) + tmpdir = tempfile.mkdtemp(prefix="ent1k_quota_1000_") + db_path = os.path.join(tmpdir, "quota_1000.db") + admin_token = "adm-" * 8 + + upstream_srv = HighConcurrencyHTTPServer(("127.0.0.1", 0), MockScaleUpstreamHandler) + upstream_port = upstream_srv.server_address[1] + upstream_thread = threading.Thread(target=upstream_srv.serve_forever, daemon=True) + upstream_thread.start() + + proxy_srv = None + proxy_mod = None + try: + os.environ["QUOTA_DB"] = db_path + os.environ["UPSTREAM_URL"] = f"http://127.0.0.1:{upstream_port}" + os.environ["UPSTREAM_API_KEY"] = "mock" + os.environ["QUOTA_ADMIN_TOKEN"] = admin_token + os.environ["QUOTA_DEFAULT_MONTHLY"] = "1000000" + + proxy_mod = _load_module("quota_proxy_scale_1000", PROXY_PATH) + proxy_mod._admin_rl = proxy_mod.RateLimiter(rate=100_000.0, capacity=100_000) + proxy_mod._proxy_rl = proxy_mod.RateLimiter(rate=100_000.0, capacity=100_000) + proxy_mod.Handler.log_message = lambda self, fmt, *a: None + + proxy_srv = HighConcurrencyHTTPServer(("127.0.0.1", 0), proxy_mod.Handler) + proxy_port = proxy_srv.server_address[1] + proxy_thread = threading.Thread(target=proxy_srv.serve_forever, daemon=True) + proxy_thread.start() + + base_url = f"http://127.0.0.1:{proxy_port}" + + def admin_req(method: str, path: str, payload: dict | None = None) -> tuple[int, dict]: + headers = { + "Authorization": f"Bearer {admin_token}", + "Content-Type": "application/json", + } + data = json.dumps(payload).encode("utf-8") if payload is not None else None + req = urllib.request.Request( + f"{base_url}{path}", data=data, headers=headers, method=method + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + return resp.status, json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as err: + with err: + return err.code, json.loads(err.read().decode("utf-8")) + + def proxy_req(quota_key: str, payload: dict) -> tuple[int, bytes]: + headers = { + "x-api-key": quota_key, + "Content-Type": "application/json", + } + data = json.dumps(payload).encode("utf-8") + req = urllib.request.Request( + f"{base_url}/v1/messages", data=data, headers=headers, method="POST" + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + return resp.status, resp.read() + except urllib.error.HTTPError as err: + with err: + return err.code, err.read() + + # 1. Register 1,000 users (user-1..user-1000) across 10 departments (dept-0..dept-9) + num_users = 1000 + expected_user_tokens: dict[str, tuple[int, int, str]] = {} + expected_dept_tokens: dict[str, dict[str, int]] = { + f"dept-{d}": {"input": 0, "output": 0, "total": 0, "requests": 0, "instances": 0} + for d in range(10) + } + + for i in range(1, num_users + 1): + inst = f"user-{i}" + dept = f"dept-{(i - 1) % 10}" + if i % 3 == 1: + in_tok, out_tok = 12 + (i % 7), 18 + (i % 5) + elif i % 3 == 2: + in_tok, out_tok = 35 + (i % 11), 45 + (i % 9) + else: + in_tok, out_tok = 50 + (i % 13), 25 + (i % 7) + expected_user_tokens[inst] = (in_tok, out_tok, dept) + expected_dept_tokens[dept]["input"] += in_tok + expected_dept_tokens[dept]["output"] += out_tok + expected_dept_tokens[dept]["total"] += in_tok + out_tok + expected_dept_tokens[dept]["requests"] += 1 + expected_dept_tokens[dept]["instances"] += 1 + + def register_one(i: int) -> tuple[int, str]: + inst = f"user-{i}" + dept = f"dept-{(i - 1) % 10}" + key = f"q-{i}" + try: + code, _ = admin_req( + "POST", + "/quota/register-key", + { + "instance": inst, + "quota_key": key, + "department": dept, + "limit": 1_000_000, + }, + ) + return code, "" + except Exception as exc: + return 599, str(exc) + + with ThreadPoolExecutor(max_workers=64) as pool: + reg_results = list(pool.map(register_one, range(1, num_users + 1))) + + reg_errors = [r for r in reg_results if r[0] != 200] + self.assertEqual( + len(reg_errors), + 0, + f"Expected 0 registration failures across 1000 users, got: {reg_errors[:5]}", + ) + + # 2. Execute 1,000 concurrent requests (JSON + Anthropic SSE + OpenAI SSE) with 64 workers + def invoke_user(i: int) -> tuple[int, int, str]: + inst = f"user-{i}" + key = f"q-{i}" + in_tok, out_tok, _ = expected_user_tokens[inst] + if i % 3 == 1: + mode = "json" + elif i % 3 == 2: + mode = "sse_anthropic" + else: + mode = "sse_openai" + try: + code, raw = proxy_req( + key, {"mock_mode": mode, "in_tok": in_tok, "out_tok": out_tok} + ) + return i, code, raw.decode("utf-8", errors="replace") + except sqlite3.OperationalError as db_exc: + return i, 500, f"sqlite3.OperationalError: {db_exc}" + except Exception as exc: + return i, 599, str(exc) + + with redirect_stdout(io.StringIO()), ThreadPoolExecutor(max_workers=64) as pool: + call_results = list(pool.map(invoke_user, range(1, num_users + 1))) + + http_500s = [r for r in call_results if r[1] >= 500] + db_locked = [r for r in call_results if "database is locked" in r[2].lower()] + non_200s = [r for r in call_results if r[1] != 200] + + self.assertEqual(len(http_500s), 0, f"Expected 0 HTTP 500 errors, got {http_500s[:5]}") + self.assertEqual(len(db_locked), 0, f"Expected 0 database lock errors, got {db_locked[:5]}") + self.assertEqual(len(non_200s), 0, f"Expected all 1000 requests to return 200, got {non_200s[:5]}") + + # 3. Verify pagination GET /quota/report?limit=100&offset=... across all 1,000 instances + collected_usage: dict[str, dict] = {} + departments_summary: dict[str, dict] = {} + for page_idx in range(10): + offset = page_idx * 100 + code, report = admin_req("GET", f"/quota/report?limit=100&offset={offset}") + self.assertEqual(code, 200) + self.assertEqual(report["total_instances"], num_users) + self.assertEqual(len(report["usage"]), 100) + if page_idx == 0: + departments_summary = report["departments"] + for row in report["usage"]: + collected_usage[row["instance"]] = row + + self.assertEqual(len(collected_usage), num_users) + + # Verify exact token accounting for every single user (user-1..user-1000) + for inst, (exp_in, exp_out, exp_dept) in expected_user_tokens.items(): + self.assertIn(inst, collected_usage) + actual = collected_usage[inst] + self.assertEqual(actual["input_tokens"], exp_in, f"Mismatch input_tokens for {inst}") + self.assertEqual(actual["output_tokens"], exp_out, f"Mismatch output_tokens for {inst}") + self.assertEqual(actual["total_tokens"], exp_in + exp_out, f"Mismatch total_tokens for {inst}") + self.assertEqual(actual["requests"], 1, f"Mismatch requests for {inst}") + self.assertEqual(actual["department"], exp_dept, f"Mismatch department for {inst}") + + # Verify exact token accounting across all 10 departments (dept-0..dept-9) + for dept_name, exp_d in expected_dept_tokens.items(): + self.assertIn(dept_name, departments_summary) + act_d = departments_summary[dept_name] + self.assertEqual(act_d["input_tokens"], exp_d["input"], f"Dept {dept_name} input mismatch") + self.assertEqual(act_d["output_tokens"], exp_d["output"], f"Dept {dept_name} output mismatch") + self.assertEqual(act_d["total_tokens"], exp_d["total"], f"Dept {dept_name} total mismatch") + self.assertEqual(act_d["requests"], exp_d["requests"], f"Dept {dept_name} requests mismatch") + self.assertEqual(act_d["instances"], exp_d["instances"], f"Dept {dept_name} instances mismatch") + + # Verify department filtering ?department=dept-0 + code, dept0_report = admin_req("GET", "/quota/report?department=dept-0") + self.assertEqual(code, 200) + self.assertEqual(dept0_report["total_instances"], 100) + self.assertEqual(len(dept0_report["usage"]), 100) + for item in dept0_report["usage"]: + self.assertEqual(item["department"], "dept-0") + + # 4. Verify department quota limit enforcement (POST /quota/set-department-limit) + dept0_total = expected_dept_tokens["dept-0"]["total"] + code, limit_resp = admin_req( + "POST", + "/quota/set-department-limit", + {"department": "dept-0", "limit": dept0_total}, + ) + self.assertEqual(code, 200) + self.assertEqual(limit_resp["limit"], dept0_total) + + # user-1 belongs to dept-0 -> must now receive 429 quota_exceeded + blocked_code, blocked_raw = proxy_req( + "q-1", {"mock_mode": "json", "in_tok": 5, "out_tok": 5} + ) + self.assertEqual(blocked_code, 429) + blocked_payload = json.loads(blocked_raw.decode("utf-8")) + self.assertEqual(blocked_payload["error"], "quota_exceeded") + self.assertIn("dept-0", blocked_payload["message"]) + + # user-2 belongs to dept-1 -> unaffected, must still succeed with 200 + ok_code, _ = proxy_req( + "q-2", {"mock_mode": "json", "in_tok": 5, "out_tok": 5} + ) + self.assertEqual(ok_code, 200) + finally: + if proxy_srv: + proxy_srv.shutdown() + proxy_srv.server_close() + upstream_srv.shutdown() + upstream_srv.server_close() + if proxy_mod and getattr(proxy_mod, "_conn", None): + proxy_mod._conn.close() + os.environ.clear() + os.environ.update(saved_env) + shutil.rmtree(tmpdir, ignore_errors=True) + + # ── Test 2: 1000 Concurrent Agents in Broker + Broadcast ────────────── + + def test_1000_agents_broker_concurrent_messaging_and_broadcast(self): + if str(BROKER_DIR) not in sys.path: + sys.path.insert(0, str(BROKER_DIR)) + import broker + + saved_env = {k: v for k, v in os.environ.items() if k.startswith("BROKER_")} + orig_db = broker.BROKER_DB + orig_rl = broker._send_rl + orig_log = broker.Handler.log_message + orig_max_age = broker.MAX_AGE_SEC + orig_retention = broker.RETENTION_SECONDS + orig_min_ttl = broker._min_active_ttl + + tmpdir = tempfile.mkdtemp(prefix="ent1k_broker_1000_") + db_path = os.path.join(tmpdir, "broker_1000.db") + admin_key = "adm" + + srv = None + try: + for k in list(os.environ.keys()): + if k.startswith("BROKER_"): + os.environ.pop(k, None) + + broker.BROKER_DB = db_path + broker.MAX_AGE_SEC = 86400 + broker.RETENTION_SECONDS = 86400 + broker._min_active_ttl = 86400.0 + broker._last_purge_ts = 0.0 + os.environ["BROKER_DB"] = db_path + os.environ["BROKER_KEY_ADMIN"] = admin_key + + broker._init_db() + broker._load_keys() + broker._send_rl = broker.RateLimiter(rate=100_000.0, capacity=100_000) + broker.Handler.log_message = lambda self, fmt, *a: None + + srv = HighConcurrencyHTTPServer(("127.0.0.1", 0), broker.Handler) + port = srv.server_address[1] + threading.Thread(target=srv.serve_forever, daemon=True).start() + + base_url = f"http://127.0.0.1:{port}" + + def broker_req(method: str, path: str, payload: dict | None = None, key: str | None = None) -> tuple[int, dict]: + headers = {"Content-Type": "application/json"} + if key: + headers["Authorization"] = f"Bearer {key}" + data = json.dumps(payload).encode("utf-8") if payload is not None else None + req = urllib.request.Request( + f"{base_url}{path}", data=data, headers=headers, method=method + ) + try: + with urllib.request.urlopen(req, timeout=15) as resp: + return resp.status, json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as err: + with err: + return err.code, json.loads(err.read().decode("utf-8")) + + num_agents = 1000 + + # 1. Register 1,000 agents (user-1..user-1000) concurrently + def register_agent(i: int) -> int: + code, _ = broker_req( + "POST", + "/admin/register-key", + {"instance": f"user-{i}", "broker_key": f"brk-agent-{i}-key"}, + key=admin_key, + ) + return code + + with ThreadPoolExecutor(max_workers=64) as pool: + reg_codes = list(pool.map(register_agent, range(1, num_agents + 1))) + + self.assertTrue( + all(c == 200 for c in reg_codes), + "All 1,000 agent key registrations must succeed with HTTP 200", + ) + + # 2. Send 1,100 concurrent P2P messages across agents + # - 1,000 ring messages: user-i -> user-((i % 1000) + 1) + # - 100 extra priority messages: user-(i + 500) -> user-i for i in 1..100 + tasks: list[tuple[int, str, str, str]] = [] + priorities = ("low", "normal", "high", "urgent") + for i in range(1, num_agents + 1): + target = f"user-{(i % num_agents) + 1}" + tasks.append((i, target, f"ring-msg-from-{i}", priorities[i % 4])) + for i in range(1, 101): + tasks.append((i + 500, f"user-{i}", f"cross-high-msg-{i}", "high")) + + def send_task(item: tuple[int, str, str, str]) -> tuple[int, str]: + sender_idx, target, text, prio = item + try: + code, body = broker_req( + "POST", + "/send", + {"to": target, "message": text, "priority": prio}, + key=f"brk-agent-{sender_idx}-key", + ) + return code, json.dumps(body) + except sqlite3.OperationalError as db_exc: + return 500, f"sqlite3.OperationalError: {db_exc}" + except Exception as exc: + return 599, str(exc) + + with redirect_stdout(io.StringIO()), ThreadPoolExecutor(max_workers=64) as pool: + send_results = list(pool.map(send_task, tasks)) + + send_failures = [r for r in send_results if r[0] != 200] + lock_errors = [r for r in send_results if "database is locked" in r[1].lower()] + self.assertEqual(len(lock_errors), 0, f"Expected 0 database lock errors, got {lock_errors[:5]}") + self.assertEqual(len(send_failures), 0, f"Expected all 1,100 sends to return 200, got {send_failures[:5]}") + + # 3. Admin broadcast to "*" (all 1,000 agents) + b_code, b_body = broker_req( + "POST", + "/broadcast", + {"to": "*", "text": "Broadcast alert to all 1000 agents", "priority": "urgent"}, + key=admin_key, + ) + self.assertEqual(b_code, 200) + self.assertTrue(b_body["ok"]) + self.assertEqual(b_body["delivered_count"], num_agents) + + # 4. Verify paginated GET /inbox/all?limit=100&offset=... across all 10 pages + all_inbox_counts: dict[str, int] = {} + for page_idx in range(10): + offset = page_idx * 100 + code, page = broker_req( + "GET", f"/inbox/all?limit=100&offset={offset}", key=admin_key + ) + self.assertEqual(code, 200) + self.assertEqual(page["total_inboxes"], num_agents) + self.assertEqual(len(page["inboxes"]), 100) + all_inbox_counts.update(page["inboxes"]) + + self.assertEqual(len(all_inbox_counts), num_agents) + # Total undelivered messages must be exactly 1,100 (P2P) + 1,000 (broadcast) = 2,100 + self.assertEqual(sum(all_inbox_counts.values()), 2100) + for i in range(1, num_agents + 1): + expected_cnt = 3 if i <= 100 else 2 + self.assertEqual( + all_inbox_counts.get(f"user-{i}"), + expected_cnt, + f"Agent user-{i} expected {expected_cnt} messages, got {all_inbox_counts.get(f'user-{i}')}", + ) + + # 5. Verify GET /inbox?peek=1 does not mark messages delivered=1 + p_code, peek_resp = broker_req("GET", "/inbox?peek=1", key="brk-agent-1-key") + self.assertEqual(p_code, 200) + self.assertTrue(peek_resp["peek"]) + self.assertEqual(peek_resp["count"], 3) + self.assertTrue(all(m["delivered"] is False for m in peek_resp["inbox"])) + self.assertEqual( + {m["priority"] for m in peek_resp["inbox"]}, + {"low", "high", "urgent"}, + ) + + # Second peek still sees delivered=False for all 3 messages + p2_code, peek2_resp = broker_req("GET", "/inbox?peek=1", key="brk-agent-1-key") + self.assertEqual(p2_code, 200) + self.assertEqual(peek2_resp["count"], 3) + self.assertTrue(all(m["delivered"] is False for m in peek2_resp["inbox"])) + + # Consume inbox for user-1 (marks delivered=True) and verify via peek + c_code, consume_resp = broker_req("GET", "/inbox", key="brk-agent-1-key") + self.assertEqual(c_code, 200) + self.assertFalse(consume_resp["peek"]) + self.assertEqual(consume_resp["count"], 3) + self.assertTrue(all(m["delivered"] is True for m in consume_resp["inbox"])) + + after_code, after_resp = broker_req("GET", "/inbox?peek=1", key="brk-agent-1-key") + self.assertEqual(after_code, 200) + self.assertEqual(after_resp["count"], 3) + self.assertTrue(all(m["delivered"] is True for m in after_resp["inbox"])) + finally: + if srv: + srv.shutdown() + srv.server_close() + broker.Handler.log_message = orig_log + for k in list(os.environ.keys()): + if k.startswith("BROKER_"): + os.environ.pop(k, None) + os.environ.update(saved_env) + broker.BROKER_DB = orig_db + broker.MAX_AGE_SEC = orig_max_age + broker.RETENTION_SECONDS = orig_retention + broker._min_active_ttl = orig_min_ttl + broker._send_rl = orig_rl + broker._load_keys() + shutil.rmtree(tmpdir, ignore_errors=True) + + # ── Test 3: Bulk Provisioning of 1000 Users with 10 Shards ──────────── + + def test_bulk_provision_1000_users_manifest(self): + bulk_mod = _load_module("bulk_users_scale_1000", BULK_SCRIPT_PATH) + tmp_root = pathlib.Path(tempfile.mkdtemp(prefix="ent1k_bulk_1000_")) + try: + # Prepare template and migrations in temporary repo root + shutil.copytree( + REPO_ROOT / "instances" / "_template", + tmp_root / "instances" / "_template", + symlinks=False, + ) + mig_dir = tmp_root / "migrations" + mig_dir.mkdir(parents=True, exist_ok=True) + (mig_dir / "001_init.py").write_text("DESCRIPTION = 'init'\n", encoding="utf-8") + (mig_dir / "024_latest.py").write_text("DESCRIPTION = 'latest'\n", encoding="utf-8") + + root_env = tmp_root / ".env" + root_env.write_text( + "QUOTA_ADMIN_TOKEN=" + ("adm-" * 8) + "\n", + encoding="utf-8", + ) + os.chmod(root_env, 0o600) + + # Generate 1,000-user manifest across 10 departments (dept-0..dept-9) + manifest_path = tmp_root / "manifest_1000.json" + users_manifest = [ + { + "name": f"user-{i}", + "bot_token": f"700{i:04d}:bot_{i:04d}", + "full_name": f"Scale User {i}", + "tg_id": str(10_000_000 + i), + "department": f"dept-{(i - 1) % 10}", + "quota_limit": 1_000_000, + } + for i in range(1, 1001) + ] + manifest_path.write_text( + json.dumps({"users": users_manifest}), encoding="utf-8" + ) + + out_buf = io.StringIO() + err_buf = io.StringIO() + t0 = time.monotonic() + with redirect_stdout(out_buf), redirect_stderr(err_buf): + rc = bulk_mod.main( + [ + "--root", + str(tmp_root), + "import", + "--file", + str(manifest_path), + "--shard-size", + "100", + "--base-port", + "18790", + ] + ) + elapsed = time.monotonic() - t0 + self.assertEqual(rc, 0, f"bulk-users.py import failed: {err_buf.getvalue()}") + self.assertLess(elapsed, 30.0, f"1,000-user import took too long: {elapsed:.2f}s") + + summary = json.loads(out_buf.getvalue()) + self.assertEqual(summary["created_count"], 1000) + self.assertEqual(summary["skipped_count"], 0) + + # Verify zero port collisions across all 1,000 instances + ports = [item["ui_port"] for item in summary["created"]] + self.assertEqual(len(ports), 1000) + self.assertEqual(len(set(ports)), 1000, "Detected UI port collision in 1000-user import") + self.assertNotIn(18789, ports) + self.assertEqual(min(ports), 18790) + self.assertEqual(max(ports), 19789) + + # Verify 10 shards (shard-0..shard-9) each with exactly 100 instances + shards = [item["shard"] for item in summary["created"]] + for shard_idx in range(10): + shard_name = f"shard-{shard_idx}" + self.assertEqual( + shards.count(shard_name), + 100, + f"Expected 100 instances in {shard_name}", + ) + + # Verify all 1,000 instance directories exist and root .env + override file are valid + instances_dir = tmp_root / "instances" + created_dirs = [ + p.name for p in instances_dir.iterdir() if p.is_dir() and p.name != "_template" + ] + self.assertEqual(len(created_dirs), 1000) + + for sample_name in ("user-1", "user-500", "user-1000"): + inst_dir = instances_dir / sample_name + self.assertTrue((inst_dir / ".env").is_file()) + self.assertTrue((inst_dir / "hermes_data" / ".env").is_file()) + self.assertTrue((inst_dir / "hermes_data" / "config.yaml").is_file()) + self.assertEqual(stat.S_IMODE((inst_dir / ".env").stat().st_mode), 0o600) + + override_text = (tmp_root / "docker-compose.override.yml").read_text(encoding="utf-8") + for shard_idx in range(10): + self.assertIn(f'- "shard-{shard_idx}"', override_text) + self.assertIn(f"corp-shard-{shard_idx}:", override_text) + + # Verify status command confirms 1,000 healthy instances and 0 issues + status_out = io.StringIO() + with redirect_stdout(status_out), redirect_stderr(err_buf): + rc_status = bulk_mod.main(["--root", str(tmp_root), "status"]) + self.assertEqual(rc_status, 0) + status_data = json.loads(status_out.getvalue()) + self.assertEqual(status_data["total_instances"], 1000) + self.assertEqual(status_data["healthy_instances"], 1000) + self.assertEqual(status_data["degraded_instances"], 0) + self.assertEqual(status_data["port_collisions"], {}) + finally: + shutil.rmtree(tmp_root, ignore_errors=True) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_service_monitor_enterprise.py b/tests/test_service_monitor_enterprise.py new file mode 100644 index 0000000..949a84a --- /dev/null +++ b/tests/test_service_monitor_enterprise.py @@ -0,0 +1,424 @@ +#!/usr/bin/env python3 +""" +Enterprise scaling tests for service-agent/server.py and resource-monitor/monitor.py +(1000-container fleet support). + +Covers: + 1. service-agent: + - MAX_CONCURRENT_SKILLS and SKILL_RATE_LIMIT env configuration + - Per-caller rate limiting (HTTP 429 for noisy caller, other callers unaffected) + - GET /v1/usage filtering (?caller=..., ?skill=...), pagination (?limit=..., ?offset=...), + and by_caller / by_skill summaries + - ThreadingHTTPServer + SQLite busy_timeout=10000 + 2. resource-monitor: + - GET /metrics on 500–1000 containers with ?search=, ?sort=mem_pct|cpu_pct|name, + ?limit=, ?offset=, and fleet summary (total_mem_mb, avg_cpu_pct, unhealthy_count) + - GET /alerts and GET /alerts/active with ?instance=, ?limit=, ?offset= + - ThreadingHTTPServer + SQLite busy_timeout=10000 +""" + +import importlib.util +import json +import os +import shutil +import sqlite3 +import tempfile +import threading +import unittest +import urllib.error +import urllib.request +from http.server import ThreadingHTTPServer +from pathlib import Path + +ROOT_DIR = Path(__file__).resolve().parent.parent +SERVICE_SERVER_PY = ROOT_DIR / "service-agent" / "server.py" +MONITOR_SERVER_PY = ROOT_DIR / "resource-monitor" / "monitor.py" + + +def _http_json(method: str, url: str, body: dict | None = None, headers: dict | None = None): + hdrs = {"Content-Type": "application/json"} + if headers: + hdrs.update(headers) + data = json.dumps(body).encode("utf-8") if body is not None else None + req = urllib.request.Request(url, data=data, headers=hdrs, method=method) + try: + with urllib.request.urlopen(req, timeout=10) as resp: + return resp.status, json.loads(resp.read().decode("utf-8")) + except urllib.error.HTTPError as e: + raw = e.read().decode("utf-8", errors="replace") + try: + parsed = json.loads(raw) + except Exception: + parsed = {"raw": raw} + return e.code, parsed + + +class TestServiceAgentEnterprise(unittest.TestCase): + """Tests for service-agent enterprise scaling (rate limit, usage pagination, concurrency).""" + + @classmethod + def setUpClass(cls): + cls.tmpdir = tempfile.mkdtemp(prefix="test_svc_enterprise_") + cls.usage_db = os.path.join(cls.tmpdir, "usage.db") + cls.skills_dir = Path(cls.tmpdir) / "skills" + cls.skills_dir.mkdir(parents=True, exist_ok=True) + + # Create two fast mock skills: echo-skill and calc-skill + for skill_name in ("echo-skill", "calc-skill"): + sdir = cls.skills_dir / skill_name + sdir.mkdir(parents=True, exist_ok=True) + (sdir / "skills.json").write_text( + json.dumps({ + "id": skill_name, + "name": skill_name, + "description": f"Test {skill_name}", + "runner": "run.py", + }), + encoding="utf-8", + ) + (sdir / "run.py").write_text( + "import json, sys\n" + "data = json.loads(sys.stdin.read() or '{}')\n" + f"print(json.dumps({{'echo': data, 'skill': '{skill_name}'}}))\n", + encoding="utf-8", + ) + + # Set environment variables BEFORE importing service-agent/server.py + os.environ["SERVICE_API_KEY"] = "test-enterprise-svc-key" + os.environ["MAX_CONCURRENT_SKILLS"] = "50" + os.environ["SKILL_RATE_LIMIT"] = "5" # 5 req/min per caller for fast rate-limit testing + os.environ["USAGE_DB"] = cls.usage_db + os.environ["SKILLS_DIR"] = str(cls.skills_dir) + + spec = importlib.util.spec_from_file_location("service_agent_server", SERVICE_SERVER_PY) + cls.svc = importlib.util.module_from_spec(spec) + spec.loader.exec_module(cls.svc) + + cls.svc.init_db() + cls.svc.reload_skills() + + cls.httpd = ThreadingHTTPServer(("127.0.0.1", 0), cls.svc.ServiceHandler) + cls.port = cls.httpd.server_address[1] + cls.base_url = f"http://127.0.0.1:{cls.port}" + cls.auth_headers = {"Authorization": "Bearer test-enterprise-svc-key"} + + cls.server_thread = threading.Thread(target=cls.httpd.serve_forever, daemon=True) + cls.server_thread.start() + + @classmethod + def tearDownClass(cls): + cls.httpd.shutdown() + cls.httpd.server_close() + shutil.rmtree(cls.tmpdir, ignore_errors=True) + + def setUp(self): + # Reset rate limiter and clear usage DB between tests + self.svc._caller_limiter.reset() + conn = sqlite3.connect(self.usage_db) + conn.execute("DELETE FROM calls") + conn.commit() + conn.close() + + def test_01_health_and_config_limits(self): + """Health endpoint exposes MAX_CONCURRENT_SKILLS=50 and SKILL_RATE_LIMIT=5.""" + status, data = _http_json("GET", f"{self.base_url}/v1/health") + self.assertEqual(status, 200) + self.assertEqual(data["status"], "ok") + self.assertEqual(data["max_concurrent"], 50) + self.assertEqual(data["max_concurrent_skills"], 50) + self.assertEqual(data["skill_rate_limit"], 5) + + # Verify SQLite busy_timeout is 10000 ms + conn = self.svc._get_db_conn() + try: + busy_timeout = conn.execute("PRAGMA busy_timeout").fetchone()[0] + self.assertEqual(busy_timeout, 10000) + finally: + conn.close() + + def test_02_per_caller_rate_limiting_isolates_noisy_bot(self): + """A runaway bot exceeding SKILL_RATE_LIMIT gets HTTP 429 while other bots succeed.""" + # user-noisy can make 5 requests (SKILL_RATE_LIMIT=5) + for i in range(5): + status, body = _http_json( + "POST", + f"{self.base_url}/v1/run", + body={"skill": "echo-skill", "caller": "user-noisy", "params": {"seq": i}}, + headers=self.auth_headers, + ) + self.assertEqual(status, 200, f"Request {i+1} should succeed: {body}") + self.assertEqual(body["status"], "ok") + + # 6th request from user-noisy must be rejected with HTTP 429 + status_blocked, body_blocked = _http_json( + "POST", + f"{self.base_url}/v1/run", + body={"skill": "echo-skill", "caller": "user-noisy", "params": {"seq": 99}}, + headers=self.auth_headers, + ) + self.assertEqual(status_blocked, 429) + self.assertEqual(body_blocked["status"], "error") + self.assertIn("user-noisy", body_blocked.get("error", "")) + + # Meanwhile user-good-1 and user-good-2 are completely unaffected + for good_caller in ("user-good-1", "user-good-2"): + status_ok, body_ok = _http_json( + "POST", + f"{self.base_url}/v1/run", + body={"skill": "echo-skill", "caller": good_caller, "params": {"ok": True}}, + headers=self.auth_headers, + ) + self.assertEqual(status_ok, 200) + self.assertEqual(body_ok["status"], "ok") + + def test_03_usage_filtering_pagination_and_by_caller_summary(self): + """GET /v1/usage supports ?caller=, ?skill=, ?limit=, ?offset= and returns by_caller.""" + # Seed usage database directly with synthetic calls from multiple callers and skills + call_plan = [ + ("user-001", "echo-skill", True, 20, 100), + ("user-001", "echo-skill", True, 30, 120), + ("user-001", "calc-skill", True, 40, 80), + ("user-002", "echo-skill", True, 15, 90), + ("user-002", "calc-skill", False, 50, 0), + ("user-003", "calc-skill", True, 25, 110), + ("user-004", "echo-skill", True, 10, 60), + ] + for caller, skill, ok, dur, out_len in call_plan: + self.svc.record_call(caller, skill, {"test": 1}, ok, dur, out_len) + + # 1. Unfiltered usage stats (backward compatibility + new fields) + status, full = _http_json("GET", f"{self.base_url}/v1/usage", headers=self.auth_headers) + self.assertEqual(status, 200) + self.assertIn("by_skill", full) + self.assertIn("by_caller", full) + self.assertEqual(full["total_calls"], 7) + self.assertEqual(full["total_skills"], 2) + self.assertEqual(full["total_callers"], 4) + + # Top caller should be user-001 with 3 calls + self.assertEqual(full["by_caller"][0]["caller"], "user-001") + self.assertEqual(full["by_caller"][0]["total_calls"], 3) + self.assertEqual(full["by_caller"][0]["successful"], 3) + + # 2. Filter by caller (?caller=user-002) + status, u2 = _http_json( + "GET", f"{self.base_url}/v1/usage?caller=user-002", headers=self.auth_headers + ) + self.assertEqual(status, 200) + self.assertEqual(u2["total_calls"], 2) + self.assertEqual(len(u2["by_caller"]), 1) + self.assertEqual(u2["by_caller"][0]["caller"], "user-002") + self.assertEqual(u2["by_caller"][0]["successful"], 1) + + # 3. Filter by skill (?skill=calc-skill) + status, calc = _http_json( + "GET", f"{self.base_url}/v1/usage?skill=calc-skill", headers=self.auth_headers + ) + self.assertEqual(status, 200) + self.assertEqual(calc["total_calls"], 3) + self.assertEqual(len(calc["by_skill"]), 1) + self.assertEqual(calc["by_skill"][0]["skill"], "calc-skill") + self.assertEqual(calc["by_skill"][0]["total_calls"], 3) + self.assertEqual(calc["by_skill"][0]["successful"], 2) + + # 4. Pagination (?limit=2&offset=1) + status, paged = _http_json( + "GET", f"{self.base_url}/v1/usage?limit=2&offset=1", headers=self.auth_headers + ) + self.assertEqual(status, 200) + self.assertEqual(len(paged["by_caller"]), 2) + self.assertEqual(paged["by_caller"][0]["caller"], "user-002") + self.assertEqual(len(paged["calls"]), 2) + self.assertEqual(paged["limit"], 2) + self.assertEqual(paged["offset"], 1) + + +class TestResourceMonitorEnterprise(unittest.TestCase): + """Tests for resource-monitor enterprise scaling (500–1000 containers pagination/sorting/alerts).""" + + @classmethod + def setUpClass(cls): + cls.tmpdir = tempfile.mkdtemp(prefix="test_mon_enterprise_") + cls.monitor_db = os.path.join(cls.tmpdir, "metrics.db") + + os.environ["MONITOR_ADMIN_TOKEN"] = "test-enterprise-mon-token" + os.environ["MONITOR_DB"] = cls.monitor_db + os.environ["ALERT_CPU_PCT"] = "80" + os.environ["ALERT_RAM_PCT"] = "85" + + spec = importlib.util.spec_from_file_location("resource_monitor_server", MONITOR_SERVER_PY) + cls.mon = importlib.util.module_from_spec(spec) + spec.loader.exec_module(cls.mon) + + cls.mon.init_db() + + cls.httpd = ThreadingHTTPServer(("127.0.0.1", 0), cls.mon.MonitorHandler) + cls.port = cls.httpd.server_address[1] + cls.base_url = f"http://127.0.0.1:{cls.port}" + cls.auth_headers = {"Authorization": "Bearer test-enterprise-mon-token"} + + cls.server_thread = threading.Thread(target=cls.httpd.serve_forever, daemon=True) + cls.server_thread.start() + + @classmethod + def tearDownClass(cls): + cls.httpd.shutdown() + cls.httpd.server_close() + shutil.rmtree(cls.tmpdir, ignore_errors=True) + + def setUp(self): + # Populate 500 containers with deterministic CPU and RAM profiles + with self.mon._lock: + self.mon._last_metrics.clear() + self.mon._active_alerts.clear() + self.mon._last_disk.update({ + "total_gb": 1000.0, + "used_gb": 420.0, + "free_gb": 580.0, + "used_pct": 42.0, + }) + for i in range(1, 501): + name = f"corp-user-{i:04d}" + # Make top memory consumers corp-user-0491..0500 (mem_pct > 85%) + if i > 490: + ram_mb = 900.0 + (i - 490) * 10.0 # 910..1000 MB out of 1000 MB (91%..100%) + cpu_pct = 85.0 + (i - 490) * 1.0 # 86%..95% (> 80% threshold) + else: + ram_mb = 100.0 + (i % 50) * 10.0 # 100..590 MB out of 1000 MB (10%..59%) + cpu_pct = round((i % 40) * 1.5, 2) # 0..58.5% + + ram_limit_mb = 1000.0 + self.mon._last_metrics[name] = { + "timestamp": "2026-10-25T12:00:00+00:00", + "id": f"cid{i:09d}", + "name": name, + "cpu_pct": cpu_pct, + "ram_mb": ram_mb, + "ram_limit_mb": ram_limit_mb, + "mem_pct": round(ram_mb / ram_limit_mb * 100.0, 2), + } + + def test_01_sqlite_busy_timeout_and_batch_snapshots(self): + """SQLite connection enforces busy_timeout=10000 and writes 500 snapshots cleanly.""" + conn = self.mon._get_db_conn() + try: + busy_timeout = conn.execute("PRAGMA busy_timeout").fetchone()[0] + self.assertEqual(busy_timeout, 10000) + finally: + conn.close() + + with self.mon._lock: + snapshots = list(self.mon._last_metrics.values()) + self.mon.write_snapshots(snapshots) + + conn = self.mon._get_db_conn() + try: + count = conn.execute("SELECT COUNT(*) FROM snapshots").fetchone()[0] + self.assertEqual(count, 500) + finally: + conn.close() + + def test_02_metrics_sort_mem_pct_limit_offset_and_summary(self): + """GET /metrics?sort=mem_pct&limit=10 returns top 10 memory containers + fleet summary.""" + status, data = _http_json( + "GET", + f"{self.base_url}/metrics?sort=mem_pct&limit=10&offset=0", + headers=self.auth_headers, + ) + self.assertEqual(status, 200) + self.assertEqual(data["total_containers"], 500) + self.assertEqual(len(data["containers"]), 10) + + # Verify summary fields + summary = data["summary"] + self.assertIn("total_mem_mb", summary) + self.assertIn("avg_cpu_pct", summary) + self.assertIn("unhealthy_count", summary) + self.assertEqual(summary["unhealthy_count"], 10) # corp-user-0491..0500 exceed CPU & RAM thresholds + self.assertGreater(summary["total_mem_mb"], 0) + self.assertGreater(summary["avg_cpu_pct"], 0) + + # Verify descending sort by mem_pct: corp-user-0500 (100%) down to corp-user-0491 (91%) + mem_pcts = [c["mem_pct"] for c in data["containers"]] + self.assertEqual(mem_pcts, sorted(mem_pcts, reverse=True)) + self.assertEqual(data["containers"][0]["name"], "corp-user-0500") + self.assertEqual(data["containers"][0]["mem_pct"], 100.0) + self.assertEqual(data["containers"][9]["name"], "corp-user-0491") + self.assertEqual(data["containers"][9]["mem_pct"], 91.0) + + # Check second page (offset=5, limit=5) -> corp-user-0495..0491 + status2, page2 = _http_json( + "GET", + f"{self.base_url}/metrics?sort=mem_pct&limit=5&offset=5", + headers=self.auth_headers, + ) + self.assertEqual(status2, 200) + self.assertEqual(page2["total_containers"], 500) + self.assertEqual(len(page2["containers"]), 5) + self.assertEqual(page2["containers"][0]["name"], "corp-user-0495") + + def test_03_metrics_search_and_sort_cpu_or_name(self): + """GET /metrics supports ?search= and ?sort=cpu_pct|name.""" + # Search for "user-049" -> matches corp-user-0490..0499 (10 containers) + status, data = _http_json( + "GET", + f"{self.base_url}/metrics?search=user-049&sort=cpu_pct&limit=5", + headers=self.auth_headers, + ) + self.assertEqual(status, 200) + self.assertEqual(data["total_containers"], 10) + self.assertEqual(len(data["containers"]), 5) + self.assertEqual(data["containers"][0]["name"], "corp-user-0499") + + # Sort by name ascending + status_name, data_name = _http_json( + "GET", + f"{self.base_url}/metrics?search=user-049&sort=name&limit=3", + headers=self.auth_headers, + ) + self.assertEqual(status_name, 200) + names = [c["name"] for c in data_name["containers"]] + self.assertEqual(names, ["corp-user-0490", "corp-user-0491", "corp-user-0492"]) + + def test_04_alerts_and_alerts_active_filtering_and_pagination(self): + """GET /alerts and GET /alerts/active support ?instance=, ?limit=, ?offset=.""" + # Trigger alerts via check_alerts on current 500 containers + with self.mon._lock: + metrics_list = list(self.mon._last_metrics.values()) + disk_dict = dict(self.mon._last_disk) + self.mon.check_alerts(metrics_list, disk_dict, {}) + + # Each of corp-user-0491..0500 triggers both cpu and ram alerts = 20 alerts total + status, all_alerts = _http_json( + "GET", f"{self.base_url}/alerts/active", headers=self.auth_headers + ) + self.assertEqual(status, 200) + self.assertEqual(all_alerts["total"], 20) + self.assertEqual(all_alerts["count"], 20) + + # Filter by instance=corp-user-0500 (or user-0500) on /alerts + status_inst, inst_alerts = _http_json( + "GET", + f"{self.base_url}/alerts?instance=user-0500", + headers=self.auth_headers, + ) + self.assertEqual(status_inst, 200) + self.assertEqual(inst_alerts["total"], 2) + self.assertEqual(inst_alerts["count"], 2) + for alert in inst_alerts["alerts"]: + self.assertIn("corp-user-0500", alert["key"]) + + # Paginate /alerts/active?limit=5&offset=10 + status_paged, paged_alerts = _http_json( + "GET", + f"{self.base_url}/alerts/active?limit=5&offset=10", + headers=self.auth_headers, + ) + self.assertEqual(status_paged, 200) + self.assertEqual(paged_alerts["total"], 20) + self.assertEqual(paged_alerts["count"], 5) + self.assertEqual(len(paged_alerts["alerts"]), 5) + + +if __name__ == "__main__": + unittest.main()