"""Real TCP + real MCP SDK + full payload ledger. Disposable Compose DB only. The host runner restarts MCP when restart-request.json appears. The test does not mount the Docker socket or give MCP container-control privileges. """ import asyncio import base64 from contextlib import asynccontextmanager import json import math import os from pathlib import Path import struct import sys import time from types import SimpleNamespace import httpx import psycopg from mcp import ClientSession from mcp.client.streamable_http import streamable_http_client sys.path.insert(0, "/stress") from network_load import run as generate from verify_ledger import expected_digest API = "http://172.31.250.10:65000" REPORTS = Path("/reports") TOKEN = "mcp-integration-test-token-000000000000" @asynccontextmanager async def session(): async with httpx.AsyncClient(headers={"Authorization": "Bearer " + TOKEN}, timeout=15) as transport: async with streamable_http_client("http://mcp:65002/mcp", http_client=transport) as (read, write, _): async with ClientSession(read, write) as client: await client.initialize() yield client async def call(client, name, args=None): response = await client.call_tool(name, args or {}) assert not response.isError, (name, response.content) assert response.structuredContent is not None, name return response.structuredContent async def finish(client, job): async with asyncio.timeout(15): while True: status = await call(client, "get_job", {"job_id": job["job_id"]}) if status["status"] != "running": assert status["status"] in ("completed", "partial"), status return status await asyncio.sleep(0.1) async def main(): REPORTS.mkdir(exist_ok=True) requested = int(os.getenv("MCP_STRESS_STREAMS", "50000")) duration = float(os.getenv("MCP_STRESS_SECONDS", "120")) assert requested >= 5000 and duration >= 40, "profile too small to test restart and overlap" report = {"checks": [], "baseline_streams": 2000, "live_streams": requested} db = await psycopg.AsyncConnection.connect("host=db dbname=packmate user=packmate password=network-stress", autocommit=True) async def rows(sql, args=()): return await (await db.execute(sql, args)).fetchall() async def scalar(sql, args=()): return (await rows(sql, args))[0][0] async def wait_count(expected): async with asyncio.timeout(60): while await scalar("SELECT count(*) FROM stream") != expected: await asyncio.sleep(0.5) async def load(count, run_id, name, seconds=0): result = await generate(SimpleNamespace(host="172.31.250.10", port=18080, streams=count, frames=8, payload_bytes=1024, concurrency=128, run_id=run_id, shard=0, stream_offset=0, report=str(REPORTS / (name + "-load.json")), abort_file=None, minimum_duration_seconds=seconds)) assert result["successfulStreams"] == count and not result["sampleErrors"], result return result async def ledger(run_id, count, lower=0): # Aggregate in PostgreSQL, transfer hashes, independently reconstruct expected bytes. entries = await rows(""" WITH payloads AS ( SELECT s.id, string_agg(p.content, ''::bytea ORDER BY p.id) FILTER (WHERE p.incoming) AS inc, string_agg(p.content, ''::bytea ORDER BY p.id) FILTER (WHERE NOT p.incoming) AS out, bool_and(CASE WHEN p.incoming THEN p.destination_ip='172.31.250.10' AND p.source_ip<>'172.31.250.10' ELSE p.source_ip='172.31.250.10' AND p.destination_ip<>'172.31.250.10' END) AS ips, count(p.id)=s.packets_count AND sum(octet_length(p.content))=s.size_bytes AND s.service_id=18080 AND s.protocol=0 AS metadata FROM stream s JOIN packet p ON p.stream_id=s.id WHERE s.id > %s GROUP BY s.id) SELECT id, substring(inc FROM 5 FOR 24), octet_length(inc), md5(inc), octet_length(out), md5(out), ips, metadata FROM payloads ORDER BY id """, (lower,)) assert len(entries) == count, (len(entries), count) seen = set() for sid, marker, insize, indigest, outsize, outdigest, ips, metadata in entries: run, connection, shard, sequence = struct.unpack("!QQII", marker) assert run == run_id and shard == sequence == 0 and 0 <= connection < count, (sid, marker) assert connection not in seen, ("duplicate", connection) seen.add(connection) assert insize == outsize == 8 * 1028 and ips and metadata, sid assert indigest == outdigest == expected_digest(run_id, connection, 0, 8, 1024), sid return {"verified_streams": len(seen), "bytes_per_direction": count * 8 * 1028} async with httpx.AsyncClient(auth=("stress", "stress"), timeout=10) as rest: for _ in range(90): try: if (await rest.get(API + "/api/analysis/overview")).status_code == 200: break except httpx.HTTPError: pass await asyncio.sleep(1) else: raise AssertionError("backend not ready") assert await scalar("SELECT count(*) FROM stream") == 0, "requires empty disposable database" async with session() as client: settings = {"name": "MCP capacity echo", "http": False, "urldecode_http_requests": False, "merge_adjacent_packets": False} await call(client, "create_service", {"port": 18080, "settings": settings}) assert not (await call(client, "create_service", {"port": 18080, "settings": settings}))["created"] pattern = await call(client, "create_pattern", {"settings": {"name": "capacity marker", "match": {"type": "substring", "value": "PMSTRESS"}, "service_port": 18080}}) pattern_id = pattern["pattern"]["id"] report["baseline_load"] = await load(2000, 8123001, "baseline") await wait_count(2000) report["baseline_ledger"] = await ledger(8123001, 2000) baseline_end = await scalar("SELECT max(id) FROM stream") preserved = [r[0] for r in await rows("SELECT id FROM stream ORDER BY id LIMIT 10")] assert len((await call(client, "set_favorites", {"stream_ids": preserved}))["updated"]) == 10 # Real scheduler: both slots full, reject a third, cancel active work, # keep partial results, and independently hit both budget guards. query = {"last_minutes": 60, "service_ports": [18080]} first = await call(client, "aggregate_traffic", {"query": query}) second = await call(client, "aggregate_traffic", {"query": query}) busy = await client.call_tool("aggregate_traffic", {"query": query}) assert busy.isError and "BUSY" in str(busy.content), busy await asyncio.sleep(0.2) cancelled = await call(client, "cancel_job", {"job_id": first["job_id"]}) assert cancelled["status"] == "cancelled" and cancelled["scanned_packets"] > 0, cancelled assert (await call(client, "get_job_results", {"job_id": first["job_id"]}))["items"] timed = await finish(client, second) assert timed["reason"] == "time_budget" and not timed["coverage_complete"], timed byte_job = await finish(client, await call(client, "search_traffic", {"query": query | { "match": {"type": "substring", "value": "NEVER_PRESENT_31337"}}})) assert byte_job["reason"] == "byte_budget" and not byte_job["coverage_complete"], byte_job report["budget_checks"] = {"cancelled_packets": cancelled["scanned_packets"], "time_reason": timed["reason"], "byte_reason": byte_job["reason"]} page1 = await finish(client, await call(client, "search_traffic", {"query": query, "limit": 20})) resumed = page1["query"] | {"cursor": page1["next_scan_cursor"]} page2 = await finish(client, await call(client, "search_traffic", {"query": resumed, "limit": 20})) one = (await call(client, "get_job_results", {"job_id": page1["job_id"]}))["items"] two = (await call(client, "get_job_results", {"job_id": page2["job_id"]}))["items"] assert len(one) == len(two) == 20 and not ({r["packet_id"] for r in one} & {r["packet_id"] for r in two}) lost_job_id = page1["job_id"] report["checks"].append("real MCP: concurrent job cap, cancellation, retained partial results, byte/time budgets and cursor resume") for job in (first, second, byte_job, page1, page2): await call(client, "delete_job", {"job_id": job["job_id"]}) lost_job_id = (await finish(client, await call(client, "search_traffic", {"query": query, "limit": 1})))["job_id"] await db.execute("UPDATE stream SET start_timestamp=start_timestamp-7200000, end_timestamp=end_timestamp-7200000, persisted_at=persisted_at-7200000 WHERE NOT favorite") done = asyncio.Event() restarted = asyncio.Event() latencies, drops, remaining, job_reasons = [], [], [], [] started = time.monotonic() reconnects = 0 async def traffic(): try: report["live_load"] = await load(requested, 8123002, "live", duration) finally: done.set() async def monitor(): while not done.is_set(): before = time.monotonic() response = await rest.get(API + "/api/pcap/metrics") response.raise_for_status() latencies.append((time.monotonic() - before) * 1000) metric = response.json() assert all(metric[k] == 0 for k in ("droppedPackets", "interfaceDroppedPackets", "applicationDroppedPackets")), metric drops.append(metric) remaining.append(await scalar("SELECT count(*) FROM stream WHERE id<=%s AND NOT favorite", (baseline_end,))) await asyncio.sleep(0.4) async def investigations(): nonlocal reconnects while not done.is_set(): try: async with session() as client: while not done.is_set(): job = await finish(client, await call(client, "search_traffic", {"query": { "last_minutes": 60, "match": {"value": "(a+)+$"}}, "limit": 1000})) job_reasons.append(job["reason"]) await call(client, "delete_job", {"job_id": job["job_id"]}) except Exception: # Only the deliberate restart window permits transport failures. assert (REPORTS / "restart-request.json").exists() and not restarted.is_set() reconnects += 1 await asyncio.sleep(0.5) async def restart_check(): await asyncio.sleep(min(30, duration / 3)) (REPORTS / "restart-request.json").write_text(json.dumps({"elapsed_seconds": time.monotonic() - started})) async with asyncio.timeout(30): while not (REPORTS / "restart-done.json").exists(): await asyncio.sleep(0.2) async with session() as client: missing = await client.call_tool("get_job", {"job_id": lost_job_id}) assert missing.isError and "NOT_FOUND" in str(missing.content) assert any(s["port"] == 18080 for s in (await call(client, "list_services"))["items"]) assert any(p["id"] == pattern_id for p in (await call(client, "list_patterns"))["items"]) assert (await call(client, "get_stream", {"stream_id": preserved[0]}))["favorite"] # Allow investigation task to finish reconnecting before closing the window. await asyncio.sleep(2) restarted.set() await asyncio.gather(traffic(), monitor(), investigations(), restart_check()) assert restarted.is_set() and any(0 < n < 1990 for n in remaining), remaining[:10] await wait_count(requested + 10) report["live_ledger"] = await ledger(8123002, requested, baseline_end) assert await scalar("SELECT count(*) FROM stream WHERE id<=%s", (baseline_end,)) == 10 assert await scalar("SELECT count(*) FROM stream_found_patterns WHERE pattern_id=%s", (pattern_id,)) == requested + 10 assert await scalar("SELECT count(*) FROM packet p LEFT JOIN stream s ON s.id=p.stream_id WHERE s.id IS NULL") == 0 p99 = sorted(latencies)[math.ceil(len(latencies) * .99) - 1] assert p99 < 500, p99 report.update(api_p99_ms=round(p99, 3), api_samples=len(latencies), metrics=drops[-1], mcp_jobs=len(job_reasons), job_reasons=sorted({str(x) for x in job_reasons}), reconnect_attempts=reconnects, cleanup_overlap=True, restart_during_traffic=True) report["checks"].append("full TCP ledger, live markers, favorites and orphan checks pass across cleaner + MCP restart") # Actual large binary traffic, merged as a single saved message per direction. async with session() as client: settings["merge_adjacent_packets"] = True await call(client, "update_service", {"port": 18080, "settings": settings}) async def edge_payload(payload): prior = await scalar("SELECT max(id) FROM stream") reader, writer = await asyncio.open_connection("172.31.250.10", 18080) writer.write(struct.pack("!I", len(payload)) + payload) await writer.drain() size = struct.unpack("!I", await reader.readexactly(4))[0] assert await reader.readexactly(size) == payload writer.close() await writer.wait_closed() async with asyncio.timeout(15): while not (saved := await rows("SELECT id,start_timestamp FROM stream WHERE id>%s", (prior,))): await asyncio.sleep(.2) assert len(saved) == 1 sid, timestamp = saved[0] return sid, {"from_ms": timestamp, "to_ms": timestamp + 1, "direction": "incoming"} binary = bytes(range(256)) * 384 sid, query = await edge_payload(binary) stream = await call(client, "get_stream", {"stream_id": sid}) packet = next(p for p in stream["packets"] if p["incoming"]) assert packet["size_bytes"] == len(binary) + 4 chunks, offset = [], 0 while offset is not None: part = await call(client, "read_packet", {"packet_id": packet["packet_id"], "offset": offset, "length": 16384, "format": "base64"}) chunks.append(base64.b64decode(part["content"])) offset = part["next_offset"] assert b"".join(chunks) == struct.pack("!I", len(binary)) + binary skipped = await finish(client, await call(client, "search_traffic", {"query": query | {"match": {"type": "hex", "value": "00010203"}}})) assert skipped["skipped_oversize_packets"] == 1 and skipped["matched_packets"] == 0 and not skipped["coverage_complete"], skipped report["checks"].append("96 KiB binary message: exact paged read; search reports oversized skip without false match") sid, query = await edge_payload(b"x" * 40000) # Legacy UI-created literals must receive the same bounded lookback as MCP patterns. response = await rest.post(API + "/api/pattern/", json={"name": "legacy literal", "value": "x", "color": "#123456", "searchType": "SUBSTRING", "searchTarget": "CONTENT", "directionType": "INPUT", "actionType": "FIND", "serviceId": 18080}) response.raise_for_status() legacy = response.json()["id"] await call(client, "set_pattern_enabled", {"pattern_id": legacy, "enabled": False}) for _ in range(2): applied = await finish(client, await call(client, "start_lookback", {"pattern_id": legacy, "query": query})) assert applied["applied_packets"] == 1, applied assert await scalar("SELECT count(*) FROM found_pattern WHERE pattern_id=%s", (legacy,)) == 256 ip_pattern = await call(client, "create_pattern", {"settings": {"name": "visible IP regex", "match": { "target": "destination_ip", "value": r"^172\.31\.250\.10$"}, "direction": "INPUT", "service_port": 18080}}) ipsid, _ = await edge_payload(b"ip-visibility") visible = (await rest.post(API + f"/api/packet/{ipsid}", json={"pageSize": 20})).json() assert any(any(m["patternId"] == ip_pattern["pattern"]["id"] for m in p["matches"]) for p in visible), visible await call(client, "update_pattern", {"pattern_id": legacy, "name": "renamed legacy", "color": "#654321"}) await call(client, "delete_pattern", {"pattern_id": legacy}) assert not any(p["id"] == legacy for p in (await call(client, "list_patterns"))["items"]) assert any(p["id"] == legacy and p["deleted"] for p in (await call(client, "list_patterns", {"include_deleted": True}))["items"]) await call(client, "create_service", {"port": 18081, "settings": settings | {"name": "temporary"}}) await call(client, "delete_service", {"port": 18081}) assert not any(s["port"] == 18081 for s in (await call(client, "list_services"))["items"]) report["checks"].append("legacy literal lookback capped/idempotent; IP regex visible in UI; service/pattern mutations persist") await db.close() (REPORTS / "capacity.json").write_text(json.dumps(report, indent=2)) print(json.dumps(report, indent=2), flush=True) if __name__ == "__main__": asyncio.run(main())