Files
zk-data-agent/e2e/verify_stack.py
T
2026-07-26 21:19:41 +08:00

307 lines
11 KiB
Python

from __future__ import annotations
import atexit
import io
import os
import tarfile
import time
import uuid
from typing import Any
import httpx
import jwt
import docker
from agent_platform.gateway.provider import workspace_ref
BASE_URL = os.getenv("E2E_BASE_URL", "http://127.0.0.1:3000").rstrip("/")
ADMIN_EMAIL = os.getenv("E2E_ADMIN_EMAIL", "admin-e2e@example.invalid")
ADMIN_PASSWORD = os.getenv("E2E_ADMIN_PASSWORD", "e2e-admin-password")
FORWARD_JWT_SECRET = os.getenv(
"E2E_FORWARD_JWT_SECRET",
"e2e-forward-jwt-secret-at-least-32-bytes",
)
INTERNAL_PROVIDER_KEY = os.getenv(
"E2E_INTERNAL_PROVIDER_KEY",
"e2e-provider-secret-at-least-32-bytes",
)
EXPECTED_MODELS = {
"luna",
"terra",
"sol",
"deepseek-v4-pro",
}
def _auth(token: str) -> dict[str, str]:
return {"Authorization": f"Bearer {token}"}
def _signin(client: httpx.Client, email: str, password: str) -> dict[str, Any]:
response = client.post("/api/v1/auths/signin", json={"email": email, "password": password})
response.raise_for_status()
data = response.json()
assert data.get("token")
return data
def _stream_chat(
client: httpx.Client,
token: str,
*,
model: str,
chat_id: str,
prompt: str,
) -> tuple[int, str]:
user_message_id = uuid.uuid4().hex
assistant_message_id = uuid.uuid4().hex
payload: dict[str, Any] = {
"model": model,
"messages": [{"role": "user", "content": prompt}],
"stream": True,
"chat_id": chat_id,
"id": assistant_message_id,
"assistant_message_id": assistant_message_id,
"user_message": {
"id": user_message_id,
"parentId": None,
"childrenIds": [assistant_message_id],
"role": "user",
"content": prompt,
"timestamp": int(time.time()),
},
"params": {},
"features": {},
}
with client.stream(
"POST",
"/api/chat/completions",
headers=_auth(token),
json=payload,
timeout=90,
) as response:
return response.status_code, "".join(response.iter_text())
def _workspace_file(client: docker.DockerClient, user_id: str, path: str) -> str:
container = client.containers.get(workspace_ref(user_id).container_name)
result = container.exec_run(["cat", f"/workspace/{path}"], user="1000:1000")
assert result.exit_code == 0, result.output.decode(errors="replace")
return result.output.decode()
def _cleanup_workspaces(users: list[dict[str, Any]]) -> None:
client = docker.from_env()
for user in users:
ref = workspace_ref(user["id"])
try:
client.containers.get(ref.container_name).remove(force=True)
except docker.errors.NotFound:
pass
try:
client.volumes.get(ref.volume_name).remove(force=True)
except docker.errors.NotFound:
pass
try:
client.networks.get(ref.network_name).remove()
except docker.errors.NotFound:
pass
client.close()
def _runtime_url(client: docker.DockerClient) -> str:
container = client.containers.get("k1412-agent-e2e-runtime-1")
networks = container.attrs["NetworkSettings"]["Networks"]
internal = next(value for name, value in networks.items() if name.endswith("_internal"))
return f"http://{internal['IPAddress']}:8000"
def _runtime_headers(user: dict[str, Any], chat_id: str) -> dict[str, str]:
now = int(time.time())
identity = jwt.encode(
{
"sub": user["id"],
"email": user["email"],
"name": user["name"],
"role": "user",
"iss": "open-webui",
"iat": now,
"exp": now + 60,
},
FORWARD_JWT_SECRET,
algorithm="HS256",
)
return {
"Authorization": f"Bearer {INTERNAL_PROVIDER_KEY}",
"X-OpenWebUI-User-Jwt": identity,
"X-OpenWebUI-Chat-Id": chat_id,
}
def main() -> None:
suffix = uuid.uuid4().hex[:10]
user_specs = [
(f"E2E User A {suffix}", f"e2e-a-{suffix}@example.invalid", "e2e-user-password-a"),
(f"E2E User B {suffix}", f"e2e-b-{suffix}@example.invalid", "e2e-user-password-b"),
]
with httpx.Client(base_url=BASE_URL, timeout=30) as client:
health = client.get("/health")
health.raise_for_status()
assert health.json().get("status") is True
docs_entry = client.get("/doc", follow_redirects=False)
assert docs_entry.status_code in {307, 308}
assert docs_entry.headers["location"].endswith("/doc/")
docs = client.get("/doc/")
docs.raise_for_status()
assert "K1412 Agent · 架构、实现与实验" in docs.text
assert docs.headers["content-type"].startswith("text/html")
docs_style = client.get("/doc/styles.css")
docs_style.raise_for_status()
assert "--eva-purple: #5b2a7e" in docs_style.text
docs_script = client.get("/doc/app.js")
docs_script.raise_for_status()
assert "safe-parallel-v1" in docs_script.text
docs_experiments = client.get("/doc/experiments.json")
docs_experiments.raise_for_status()
assert {item["status"] for item in docs_experiments.json()} == {
"baseline",
"validated",
"backlog",
}
admin = _signin(client, ADMIN_EMAIL, ADMIN_PASSWORD)
admin_token = admin["token"]
config = client.get("/api/v1/auths/admin/config", headers=_auth(admin_token))
config.raise_for_status()
admin_config = config.json()
assert admin_config["ENABLE_SIGNUP"] is True
assert admin_config["DEFAULT_USER_ROLE"] == "pending"
assert admin_config["ENABLE_API_KEYS"] is False
assert admin_config["ENABLE_AUTOMATIONS"] is False
model_response = client.get("/api/models?refresh=true", headers=_auth(admin_token))
model_response.raise_for_status()
model_payload = model_response.json()
model_ids = {item["id"] for item in model_payload["data"]}
assert model_ids == EXPECTED_MODELS, model_ids
assert "ChatGPT-5.6" not in model_response.text
users: list[dict[str, Any]] = []
atexit.register(_cleanup_workspaces, users)
for name, email, password in user_specs:
signup = client.post(
"/api/v1/auths/signup",
json={"name": name, "email": email, "password": password},
)
signup.raise_for_status()
pending = signup.json()
assert pending["role"] == "pending"
denied = client.get("/api/models", headers=_auth(pending["token"]))
assert denied.status_code in {401, 403}
approved = client.post(
f"/api/v1/users/{pending['id']}/update",
headers=_auth(admin_token),
json={"role": "user"},
)
approved.raise_for_status()
assert approved.json()["role"] == "user"
users.append(_signin(client, email, password))
for user in users:
models = client.get("/api/models", headers=_auth(user["token"]))
models.raise_for_status()
user_model_ids = {item["id"] for item in models.json()["data"]}
assert user_model_ids == EXPECTED_MODELS, user_model_ids
assert "ChatGPT-5.6" not in models.text
forbidden = client.get("/api/v1/auths/admin/config", headers=_auth(user["token"]))
assert forbidden.status_code in {401, 403}
chat_ids = [f"local:e2e-{suffix}-a", f"local:e2e-{suffix}-b"]
for index, user in enumerate(users):
marker = f"E2E_AGENT_USER_{index}_{suffix}"
status_code, body = _stream_chat(
client,
user["token"],
model="luna" if index == 0 else "terra",
chat_id=chat_ids[index],
prompt=marker,
)
assert status_code == 200, body
docker_client = docker.from_env()
assert _workspace_file(docker_client, users[0]["id"], "work-proof.txt") == f"E2E_AGENT_USER_0_{suffix}"
assert _workspace_file(docker_client, users[1]["id"], "work-proof.txt") == f"E2E_AGENT_USER_1_{suffix}"
for user in users:
ref = workspace_ref(user["id"])
workspace = docker_client.containers.get(ref.container_name)
workspace.reload()
assert set(workspace.attrs["NetworkSettings"]["Networks"]) == {ref.network_name}
listing = client.get(
"/api/v1/k1412/workspace/files",
headers=_auth(users[0]["token"]),
params={"path": "."},
)
listing.raise_for_status()
assert "work-proof.txt" in {item["name"] for item in listing.json()["entries"]}
download = client.get(
"/api/v1/k1412/workspace/download",
headers=_auth(users[0]["token"]),
params={"path": "work-proof.txt"},
)
download.raise_for_status()
assert download.text == f"E2E_AGENT_USER_0_{suffix}"
isolated_listing = client.get(
"/api/v1/k1412/workspace/files",
headers=_auth(users[1]["token"]),
params={"path": "."},
)
isolated_listing.raise_for_status()
assert "work-proof.txt" in {item["name"] for item in isolated_listing.json()["entries"]}
isolated_download = client.get(
"/api/v1/k1412/workspace/download",
headers=_auth(users[1]["token"]),
params={"path": "work-proof.txt"},
)
isolated_download.raise_for_status()
assert isolated_download.text == f"E2E_AGENT_USER_1_{suffix}"
archive = client.get(
"/api/v1/k1412/workspace/archive",
headers=_auth(users[0]["token"]),
params={"path": "."},
)
archive.raise_for_status()
with tarfile.open(fileobj=io.BytesIO(archive.content), mode="r:gz") as workspace_tar:
assert any(name.endswith("work-proof.txt") for name in workspace_tar.getnames())
with httpx.Client(base_url=_runtime_url(docker_client), timeout=30, trust_env=False) as runtime:
runtime_headers = _runtime_headers(users[0], chat_ids[0])
events = runtime.get(f"/v1/runs/{chat_ids[0]}", headers=runtime_headers)
events.raise_for_status()
event_types = {event["type"] for event in events.json()["items"]}
assert {"run.created", "tool.completed", "run.completed"} <= event_types
legacy = runtime.post(
"/v1/chat/completions",
headers=runtime_headers,
json={
"model": "work-high",
"messages": [{"role": "user", "content": "LEGACY_MODEL_MUST_FAIL"}],
"stream": False,
},
)
assert legacy.status_code == 404, legacy.text
print(
"E2E passed: public documentation, auth approval, four Agent models, custom loop "
"evidence, file downloads, and per-user workspace isolation."
)
if __name__ == "__main__":
main()