213 lines
9.1 KiB
Python
213 lines
9.1 KiB
Python
import asyncio
|
|
import struct
|
|
import time
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from app.gateway.worker_pool.diagnostics import Diagnostics, container_metrics, redact_logs
|
|
from app.gateway.worker_pool.docker import DockerError, decode_docker_output, worker_spec
|
|
from app.gateway.worker_pool.host_metrics import rates, read_snapshot
|
|
from app.gateway.worker_pool.preflight import postgres_connection_string
|
|
from app.gateway.worker_pool.settings import PoolSettings, SettingsStore
|
|
from app.gateway.worker_pool.telemetry import RequestRecorder, RequestTelemetryMiddleware, route_name
|
|
|
|
|
|
def test_metrics_rates_use_elapsed_time_and_first_sample_is_unknown():
|
|
def sample(at, total, idle, read, rx):
|
|
return {
|
|
"available": True,
|
|
"at": at,
|
|
"cpu_total": total,
|
|
"cpu_idle": idle,
|
|
"cpu_iowait": 5,
|
|
"disks": {"sda": {"read_bytes": read, "write_bytes": 100, "io_ms": 100, "inflight": 0}},
|
|
"interfaces": {"eth0": {"rx_bytes": rx, "tx_bytes": 100}},
|
|
}
|
|
|
|
first = sample(1, 100, 50, 100, 100)
|
|
assert rates(first, None)["cpu_percent"] is None
|
|
result = rates(sample(3, 200, 75, 500, 900), first)
|
|
assert result["cpu_percent"] == 75
|
|
assert result["disks"]["sda"]["read_bps"] == 200
|
|
assert result["interfaces"]["eth0"]["rx_bps"] == 400
|
|
assert rates(sample(3, 200, 75, 500, 900), first)["iowait_percent"] == 0
|
|
|
|
|
|
def test_request_telemetry_aggregates_routes_per_execution_node():
|
|
recorder = RequestRecorder()
|
|
now = time.time()
|
|
for node, elapsed in (("main", 8), ("dfw-worker-1", 120)):
|
|
recorder.finish({
|
|
"id": node,
|
|
"started": now,
|
|
"method": "GET",
|
|
"path": "/api/threads/{thread_id}/state",
|
|
"status": 200,
|
|
"stream": False,
|
|
"headers_ms": elapsed,
|
|
"total_ms": elapsed,
|
|
"node": node,
|
|
})
|
|
|
|
summary = recorder.summary()
|
|
assert {(row["node"], row["max_ms"]) for row in summary["routes"]} == {
|
|
("main", 8),
|
|
("dfw-worker-1", 120),
|
|
}
|
|
assert recorder.worker_status()["count"] == 2
|
|
|
|
|
|
def test_host_parser_reads_host_network_not_container_network(tmp_path):
|
|
proc = tmp_path / "proc"
|
|
sys = tmp_path / "sys"
|
|
(proc / "1/net").mkdir(parents=True)
|
|
(sys / "class/block/sda1").mkdir(parents=True)
|
|
(sys / "class/block/sda1/partition").write_text("1")
|
|
(proc / "stat").write_text("cpu 10 0 20 100 5 0 0 0 0\ncpu0 1 0 0 0\ncpu1 1 0 0 0\n")
|
|
(proc / "meminfo").write_text("MemTotal: 1000 kB\nMemAvailable: 300 kB\n")
|
|
(proc / "loadavg").write_text("1.5 1.0 0.5 1/1 1")
|
|
(proc / "diskstats").write_text("8 0 sda 1 0 2 0 1 0 4 0 0 10 10\n8 1 sda1 1 0 2 0 1 0 4 0 0 10 10\n")
|
|
(proc / "1/net/dev").write_text("header\nheader\n eth0: 100 1 0 0 0 0 0 0 200 1 0 0 0 0 0 0\n veth1: 100 1 0 0 0 0 0 0 100 1 0 0 0 0 0 0\n")
|
|
result = read_snapshot(proc, sys)
|
|
assert result["cpu_count"] == 2
|
|
assert result["memory_available"] == 300 * 1024
|
|
assert list(result["interfaces"]) == ["eth0"]
|
|
assert list(result["disks"]) == ["sda"]
|
|
assert not read_snapshot(tmp_path / "absent", sys)["available"]
|
|
|
|
|
|
def test_docker_metrics_and_log_frames_and_secret_redaction():
|
|
detail = {"Name": "/node", "State": {"StartedAt": "today", "Health": {"Status": "healthy"}}, "HostConfig": {"NanoCpus": 2_000_000_000}}
|
|
stats = {
|
|
"cpu_stats": {"cpu_usage": {"total_usage": 300}, "system_cpu_usage": 1000, "online_cpus": 4},
|
|
"precpu_stats": {"cpu_usage": {"total_usage": 200}, "system_cpu_usage": 800},
|
|
"memory_stats": {"usage": 1000, "stats": {"inactive_file": 100}},
|
|
}
|
|
result = container_metrics(detail, stats)
|
|
assert result["cpu_percent"] == 200
|
|
assert result["memory_bytes"] == 900
|
|
value = b"password=secret\n"
|
|
frame = b"\x01\0\0\0" + struct.pack(">I", len(value)) + value
|
|
assert decode_docker_output(frame) == value.decode()
|
|
redacted = redact_logs('password=secret Authorization: Bearer hidden mysql://u:p@host/db api_key="abcdef" eyJabc.xxxxx.yyyyy')
|
|
assert all(secret not in redacted for secret in ("secret", "hidden", "u:p", "abcdef", "eyJabc"))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_stream_timing_does_not_buffer_or_read_request_body():
|
|
recorder = RequestRecorder()
|
|
recorder.enabled = True
|
|
started = asyncio.Event()
|
|
release = asyncio.Event()
|
|
bodies = []
|
|
|
|
async def app(scope, receive, send):
|
|
await send({"type": "http.response.start", "status": 200, "headers": [(b"content-type", b"text/event-stream")]})
|
|
await send({"type": "http.response.body", "body": b"event: metadata\n\n", "more_body": True})
|
|
started.set()
|
|
await release.wait()
|
|
await send({"type": "http.response.body", "body": b"event: end\n\n", "more_body": False})
|
|
|
|
async def receive():
|
|
pytest.fail("telemetry must never consume receive")
|
|
|
|
async def send(message):
|
|
if message["type"] == "http.response.body":
|
|
bodies.append(message["body"])
|
|
|
|
scope = {"type": "http", "app": SimpleNamespace(state=SimpleNamespace(request_recorder=recorder)), "method": "POST", "path": "/api/threads/private-id/runs/stream"}
|
|
task = asyncio.create_task(RequestTelemetryMiddleware(app)(scope, receive, send))
|
|
await started.wait()
|
|
assert bodies == [b"event: metadata\n\n"]
|
|
assert len(recorder.active) == 1 and not recorder.completed
|
|
release.set()
|
|
await task
|
|
row = recorder.completed[0]
|
|
assert row["stream"] and row["status"] == 200 and row["first_body_ms"] is not None
|
|
assert row["path"] == "/api/threads/{thread_id}/runs/stream"
|
|
assert not recorder.active
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_disabled_or_monitor_request_is_not_recorded():
|
|
recorder = RequestRecorder()
|
|
|
|
async def app(scope, receive, send):
|
|
pass
|
|
|
|
scope = {"type": "http", "app": SimpleNamespace(state=SimpleNamespace(request_recorder=recorder)), "method": "GET", "path": "/api/models"}
|
|
await RequestTelemetryMiddleware(app)(scope, None, None)
|
|
recorder.enabled = True
|
|
scope["path"] = "/api/system-settings/worker-pool/diagnostics"
|
|
await RequestTelemetryMiddleware(app)(scope, None, None)
|
|
assert not recorder.active and not recorder.completed
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistence_switch_archives_and_arbitrary_container_denied(tmp_path):
|
|
store = SettingsStore(tmp_path)
|
|
store.save(PoolSettings(source_container="main"))
|
|
pool = SimpleNamespace(store=store, nodes={})
|
|
app = SimpleNamespace(state=SimpleNamespace(config=None))
|
|
monitor = Diagnostics(pool, app)
|
|
await monitor.set_enabled(True)
|
|
monitor.recorder.finish({"id": "x", "started": 1, "method": "GET", "path": "/api/models", "status": 200, "stream": False, "total_ms": 2, "headers_ms": 1})
|
|
monitor.persist_requests()
|
|
|
|
class Engine:
|
|
def request(self, *args, **kwargs):
|
|
return b"password=do-not-keep\nhealthy"
|
|
|
|
await monitor.archive_logs(Engine(), "node")
|
|
archives = monitor.archives()
|
|
assert len(archives) == 1
|
|
assert "do-not-keep" not in str(monitor.read_archive(archives[0]["id"]))
|
|
with pytest.raises(DockerError):
|
|
monitor.read_archive("../../outside")
|
|
with pytest.raises(DockerError):
|
|
await monitor.logs("unrelated-database")
|
|
await monitor.close()
|
|
restored = Diagnostics(pool, app)
|
|
assert restored.enabled and len(restored.recorder.completed) == 1
|
|
await restored.set_enabled(False)
|
|
await restored.close()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_diagnostics_off_has_no_background_tasks_or_storage(tmp_path):
|
|
store = SettingsStore(tmp_path)
|
|
monitor = Diagnostics(SimpleNamespace(store=store, nodes={}), SimpleNamespace(state=SimpleNamespace(config=None)))
|
|
monitor.start()
|
|
assert monitor.task is None and monitor.lag_task is None
|
|
assert monitor.db is None
|
|
assert not (tmp_path / "diagnostics.sqlite3").exists()
|
|
assert not (tmp_path / "log-archive").exists()
|
|
await monitor.close()
|
|
|
|
|
|
def test_postgres_status_follows_actual_chat_configuration(monkeypatch):
|
|
monkeypatch.delenv("DEER_FLOW_USE_PG_CHECKPOINTER", raising=False)
|
|
db = SimpleNamespace(backend="mysql")
|
|
cp = SimpleNamespace(type="postgres", use_postgres=True, connection_string="postgresql://db/chat")
|
|
config = SimpleNamespace(database=db, checkpointer=cp)
|
|
assert postgres_connection_string(config) == cp.connection_string
|
|
cp.use_postgres = False
|
|
assert postgres_connection_string(config) is None
|
|
assert route_name("/api/threads/id/runs/run-id/stream") == "/api/threads/{thread_id}/runs/{run_id}/stream"
|
|
|
|
|
|
def test_worker_never_inherits_host_diagnostic_mounts():
|
|
source = {
|
|
"Image": "id",
|
|
"Config": {"Env": ["AUTH_JWT_SECRET=x"]},
|
|
"NetworkSettings": {"Networks": {"app": {}}},
|
|
"Mounts": [
|
|
{"Type": "bind", "Source": "/proc", "Destination": "/host/proc", "RW": False},
|
|
{"Type": "bind", "Source": "/sys", "Destination": "/host/sys", "RW": False},
|
|
{"Type": "bind", "Source": "/data", "Destination": "/app/data", "RW": True},
|
|
],
|
|
}
|
|
spec = worker_spec(source, PoolSettings(network="app"), "pool", "token", "/private")
|
|
assert [m["Target"] for m in spec["HostConfig"]["Mounts"]] == ["/app/data"]
|