refactor: Sandbox Provisioner 优化 + 数据库模型更新 + 初始化脚本完善
- 优化 Sandbox Provider/Provisioner 配置和错误处理 - 新增 postgres 数据库日志与项目成员模型 - 完善初始化脚本(init.sh/ps1)和 Makefile - 新增 seed_initial_users.py 初始用户种子脚本 - 更新 uv.lock 依赖锁定
This commit is contained in:
parent
1f2bdf3158
commit
acc7a7342f
5
Makefile
5
Makefile
@ -1,5 +1,5 @@
|
||||
|
||||
.PHONY: up up-lite down logs lint format
|
||||
.PHONY: up up-lite down logs lint format seed
|
||||
|
||||
PYTEST_ARGS ?=
|
||||
|
||||
@ -26,6 +26,9 @@ logs:
|
||||
@echo "Commit ID: $$(git rev-parse HEAD)"
|
||||
@echo "System: $$(uname -a)"
|
||||
|
||||
seed:
|
||||
docker compose exec api uv run python scripts/seed_initial_users.py
|
||||
|
||||
######################
|
||||
# LINTING AND FORMATTING
|
||||
######################
|
||||
|
||||
5499
backend/package/uv.lock
Normal file
5499
backend/package/uv.lock
Normal file
File diff suppressed because it is too large
Load Diff
@ -1,6 +1,8 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
@ -16,6 +18,44 @@ def sandbox_id_for_thread(thread_id: str) -> str:
|
||||
return digest[:12]
|
||||
|
||||
|
||||
def normalize_env(env: dict | None) -> dict[str, str]:
|
||||
if not isinstance(env, dict):
|
||||
return {}
|
||||
return {str(key): "" if value is None else str(value) for key, value in env.items() if str(key)}
|
||||
|
||||
|
||||
def postgres_conninfo() -> str:
|
||||
db_url = os.getenv("POSTGRES_URL", "").strip()
|
||||
return db_url.replace("+asyncpg", "").replace("+psycopg", "")
|
||||
|
||||
|
||||
def load_user_agent_env(uid: str) -> dict[str, str]:
|
||||
conninfo = postgres_conninfo()
|
||||
if not conninfo:
|
||||
return {}
|
||||
|
||||
try:
|
||||
import psycopg
|
||||
|
||||
with psycopg.connect(conninfo, connect_timeout=3) as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT env FROM agent_envs WHERE uid = %s", (uid,))
|
||||
row = cursor.fetchone()
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"failed to load agent env for uid {uid}: {exc}") from exc
|
||||
|
||||
if not row:
|
||||
return {}
|
||||
|
||||
value = row[0]
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
value = json.loads(value)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise RuntimeError(f"stored agent env for uid {uid} is not valid JSON") from exc
|
||||
return normalize_env(value)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SandboxConnection:
|
||||
thread_id: str
|
||||
@ -93,7 +133,7 @@ class ProvisionerSandboxProvider:
|
||||
record = self._client.discover(sandbox_id)
|
||||
if record is None:
|
||||
logger.info(f"Creating sandbox {sandbox_id} for thread {thread_id}")
|
||||
record = self._client.create(sandbox_id, thread_id, uid)
|
||||
record = self._client.create(sandbox_id, thread_id, uid, load_user_agent_env(uid))
|
||||
else:
|
||||
logger.info(f"Reusing sandbox {sandbox_id} for thread {thread_id}")
|
||||
|
||||
@ -123,7 +163,7 @@ class ProvisionerSandboxProvider:
|
||||
if record is None:
|
||||
if not create_if_missing:
|
||||
return None
|
||||
record = self._client.create(sandbox_id, thread_id, uid)
|
||||
record = self._client.create(sandbox_id, thread_id, uid, load_user_agent_env(uid))
|
||||
|
||||
return self._record_to_connection(thread_id, uid, record)
|
||||
|
||||
|
||||
@ -29,11 +29,11 @@ class ProvisionerClient:
|
||||
response = self._request("GET", "/health")
|
||||
return response.status_code == 200
|
||||
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str) -> SandboxRecord:
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str, env: dict[str, str] | None = None) -> SandboxRecord:
|
||||
response = self._request(
|
||||
"POST",
|
||||
"/api/sandboxes",
|
||||
json={"sandbox_id": sandbox_id, "thread_id": thread_id, "uid": uid},
|
||||
json={"sandbox_id": sandbox_id, "thread_id": thread_id, "uid": uid, "env": env or {}},
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise RuntimeError(f"failed to create sandbox {sandbox_id}: {response.status_code} {response.text}")
|
||||
|
||||
@ -379,6 +379,16 @@ class PostgresManager(metaclass=SingletonMeta):
|
||||
"ALTER TABLE IF EXISTS conversations ADD COLUMN IF NOT EXISTS is_pinned BOOLEAN NOT NULL DEFAULT FALSE",
|
||||
"ALTER TABLE IF EXISTS mcp_servers ADD COLUMN IF NOT EXISTS env JSONB",
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS agent_envs (
|
||||
id SERIAL PRIMARY KEY,
|
||||
uid VARCHAR NOT NULL REFERENCES users(uid) ON DELETE CASCADE,
|
||||
env JSONB NOT NULL DEFAULT '{}'::jsonb,
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
CONSTRAINT uq_agent_envs_uid UNIQUE (uid)
|
||||
)
|
||||
""",
|
||||
"""
|
||||
ALTER TABLE IF EXISTS agent_configs
|
||||
ADD COLUMN IF NOT EXISTS uid VARCHAR
|
||||
""",
|
||||
|
||||
@ -82,6 +82,8 @@ class User(Base):
|
||||
# 关联 API Keys
|
||||
api_keys = relationship("APIKey", back_populates="user", cascade="all, delete-orphan")
|
||||
|
||||
agent_env = relationship("AgentEnv", back_populates="user", cascade="all, delete-orphan", uselist=False)
|
||||
|
||||
def to_dict(self, include_password: bool = False) -> dict[str, Any]:
|
||||
result = {
|
||||
"id": self.id,
|
||||
@ -130,6 +132,28 @@ class User(Base):
|
||||
self.login_locked_until = None
|
||||
|
||||
|
||||
class AgentEnv(Base):
|
||||
"""用户级 Agent 沙盒环境变量"""
|
||||
|
||||
__tablename__ = "agent_envs"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
uid = Column(String, ForeignKey("users.uid"), nullable=False, unique=True, index=True)
|
||||
env = Column(JSON, nullable=False, default=dict)
|
||||
created_at = Column(DateTime, default=utc_now_naive)
|
||||
updated_at = Column(DateTime, default=utc_now_naive, onupdate=utc_now_naive)
|
||||
|
||||
user = relationship("User", back_populates="agent_env")
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
return {
|
||||
"uid": self.uid,
|
||||
"env": self.env or {},
|
||||
"created_at": format_utc_datetime(self.created_at),
|
||||
"updated_at": format_utc_datetime(self.updated_at),
|
||||
}
|
||||
|
||||
|
||||
class AgentConfig(Base):
|
||||
"""智能体配置(按用户隔离,多份可切换)"""
|
||||
|
||||
|
||||
141
backend/scripts/seed_initial_users.py
Normal file
141
backend/scripts/seed_initial_users.py
Normal file
@ -0,0 +1,141 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import TypedDict
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from sqlalchemy import func, select
|
||||
|
||||
|
||||
APP_ROOT = Path(__file__).resolve().parents[1]
|
||||
for import_path in (APP_ROOT, APP_ROOT / "package"):
|
||||
import_path_str = str(import_path)
|
||||
if import_path_str not in sys.path:
|
||||
sys.path.insert(0, import_path_str)
|
||||
|
||||
SUPERADMIN_UID = "zwj"
|
||||
SUPERADMIN_PASSWORD = "zwj12138"
|
||||
DEFAULT_USER_PASSWORD = "yuxi123456"
|
||||
|
||||
|
||||
class DepartmentSeed(TypedDict):
|
||||
name: str
|
||||
description: str
|
||||
prefix: str
|
||||
normal_count: int
|
||||
|
||||
|
||||
DEPARTMENTS: list[DepartmentSeed] = [
|
||||
{"name": "研发部", "description": "负责产品研发与技术平台建设", "prefix": "dev", "normal_count": 5},
|
||||
{"name": "产品部", "description": "负责产品规划、需求分析与项目推进", "prefix": "prod", "normal_count": 5},
|
||||
{"name": "运营部", "description": "负责业务运营、用户支持与内容维护", "prefix": "ops", "normal_count": 4},
|
||||
]
|
||||
|
||||
|
||||
class SeedError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def load_project_env() -> None:
|
||||
load_dotenv(APP_ROOT / ".env", override=False)
|
||||
load_dotenv(APP_ROOT.parent / ".env", override=False)
|
||||
load_dotenv(Path.cwd() / ".env", override=False)
|
||||
|
||||
|
||||
async def ensure_uninitialized(session) -> None:
|
||||
from yuxi.storage.postgres.models_business import User
|
||||
|
||||
user_count = await session.scalar(select(func.count(User.id)))
|
||||
if user_count:
|
||||
raise SeedError(f"系统已初始化:users 表已有 {user_count} 个用户,脚本已退出。")
|
||||
|
||||
superadmin_count = await session.scalar(select(func.count(User.id)).where(User.role == "superadmin"))
|
||||
if superadmin_count:
|
||||
raise SeedError("系统已初始化:已存在超级管理员,脚本已退出。")
|
||||
|
||||
|
||||
async def seed_initial_users() -> None:
|
||||
from server.utils.auth_utils import AuthUtils
|
||||
from yuxi.storage.postgres.manager import pg_manager
|
||||
from yuxi.storage.postgres.models_business import Department, User
|
||||
from yuxi.utils.datetime_utils import utc_now_naive
|
||||
|
||||
try:
|
||||
pg_manager.initialize()
|
||||
await pg_manager.create_business_tables()
|
||||
await pg_manager.ensure_business_schema()
|
||||
|
||||
async with pg_manager.get_async_session_context() as session:
|
||||
await ensure_uninitialized(session)
|
||||
|
||||
departments: dict[str, Department] = {}
|
||||
for department_seed in DEPARTMENTS:
|
||||
department = Department(
|
||||
name=department_seed["name"],
|
||||
description=department_seed["description"],
|
||||
)
|
||||
session.add(department)
|
||||
departments[department_seed["prefix"]] = department
|
||||
|
||||
await session.flush()
|
||||
|
||||
users = [
|
||||
User(
|
||||
username=SUPERADMIN_UID,
|
||||
uid=SUPERADMIN_UID,
|
||||
password_hash=AuthUtils.hash_password(SUPERADMIN_PASSWORD),
|
||||
role="superadmin",
|
||||
department_id=departments["dev"].id,
|
||||
last_login=utc_now_naive(),
|
||||
)
|
||||
]
|
||||
|
||||
for department_seed in DEPARTMENTS:
|
||||
department = departments[department_seed["prefix"]]
|
||||
for index in range(1, 3):
|
||||
users.append(
|
||||
User(
|
||||
username=f"{department_seed['name']}管理员{index}",
|
||||
uid=f"{department_seed['prefix']}_admin_{index}",
|
||||
password_hash=AuthUtils.hash_password(DEFAULT_USER_PASSWORD),
|
||||
role="admin",
|
||||
department_id=department.id,
|
||||
)
|
||||
)
|
||||
for index in range(1, department_seed["normal_count"] + 1):
|
||||
users.append(
|
||||
User(
|
||||
username=f"{department_seed['name']}用户{index}",
|
||||
uid=f"{department_seed['prefix']}_user_{index:02d}",
|
||||
password_hash=AuthUtils.hash_password(DEFAULT_USER_PASSWORD),
|
||||
role="user",
|
||||
department_id=department.id,
|
||||
)
|
||||
)
|
||||
|
||||
session.add_all(users)
|
||||
finally:
|
||||
await pg_manager.close()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
load_project_env()
|
||||
try:
|
||||
asyncio.run(seed_initial_users())
|
||||
except SeedError as exc:
|
||||
print(str(exc), file=sys.stderr)
|
||||
return 1
|
||||
except Exception as exc:
|
||||
print(f"初始化种子用户失败:{exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
print("初始化完成:已创建超级管理员 zwj、3 个部门、6 个部门管理员和 14 个普通用户。")
|
||||
print("超级管理员密码:zwj12138")
|
||||
print("部门管理员和普通用户默认密码:yuxi123456")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@ -11,21 +11,38 @@ from pathlib import Path
|
||||
from urllib import request
|
||||
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, Field
|
||||
from dotenv import dotenv_values
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SANDBOX_ENV_FILE = Path(__file__).parent / "sandbox.env"
|
||||
|
||||
|
||||
def canonical_backend_name(backend: str) -> str:
|
||||
value = (backend or "").strip().lower()
|
||||
return value or "memory"
|
||||
|
||||
|
||||
def normalize_env(env: dict | None) -> dict[str, str]:
|
||||
if not isinstance(env, dict):
|
||||
return {}
|
||||
return {str(key): "" if value is None else str(value) for key, value in env.items() if str(key)}
|
||||
|
||||
|
||||
def load_sandbox_env() -> dict[str, str]:
|
||||
return normalize_env(dotenv_values(SANDBOX_ENV_FILE))
|
||||
|
||||
|
||||
def merged_sandbox_env(global_env: dict[str, str], user_env: dict[str, str]) -> dict[str, str]:
|
||||
return {**global_env, **normalize_env(user_env)}
|
||||
|
||||
|
||||
class CreateSandboxRequest(BaseModel):
|
||||
sandbox_id: str
|
||||
thread_id: str
|
||||
uid: str
|
||||
env: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class SandboxResponse(BaseModel):
|
||||
@ -69,9 +86,10 @@ class MemoryProvisionerBackend:
|
||||
return template.format(sandbox_id=sandbox_id)
|
||||
return template
|
||||
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str) -> SandboxRecord:
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str, env: dict[str, str] | None = None) -> SandboxRecord:
|
||||
_ = thread_id # unused in memory backend
|
||||
_ = uid # unused in memory backend
|
||||
_ = env # unused in memory backend
|
||||
with self._lock:
|
||||
existing = self._records.get(sandbox_id)
|
||||
if existing is not None:
|
||||
@ -113,16 +131,6 @@ def wait_for_sandbox_ready(sandbox_url: str, timeout_seconds: int = 30) -> bool:
|
||||
|
||||
|
||||
class LocalContainerProvisionerBackend:
|
||||
_SANDBOX_ENV_FILE = Path(__file__).parent / "sandbox.env"
|
||||
|
||||
@staticmethod
|
||||
def _load_sandbox_env() -> dict[str, str]:
|
||||
"""Parse sandbox.env and return environment variables to inject into sandbox containers."""
|
||||
if LocalContainerProvisionerBackend._SANDBOX_ENV_FILE.exists():
|
||||
return dotenv_values(LocalContainerProvisionerBackend._SANDBOX_ENV_FILE)
|
||||
|
||||
return {}
|
||||
|
||||
def __init__(self):
|
||||
import docker
|
||||
from docker.errors import DockerException
|
||||
@ -139,7 +147,7 @@ class LocalContainerProvisionerBackend:
|
||||
self._container_prefix = os.getenv("DOCKER_SANDBOX_PREFIX", "yuxi-sandbox")
|
||||
self._sandbox_host = os.getenv("DOCKER_SANDBOX_HOST", "host.docker.internal")
|
||||
self._health_timeout_seconds = int(os.getenv("SANDBOX_HEALTH_TIMEOUT_SECONDS", "300"))
|
||||
self._sandbox_env = self._load_sandbox_env()
|
||||
self._sandbox_env = load_sandbox_env()
|
||||
|
||||
try:
|
||||
self._client = docker.from_env()
|
||||
@ -329,7 +337,7 @@ class LocalContainerProvisionerBackend:
|
||||
except NotFound:
|
||||
return None
|
||||
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str) -> SandboxRecord:
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str, env: dict[str, str] | None = None) -> SandboxRecord:
|
||||
with self._lock:
|
||||
safe_thread_id = self._validate_thread_id(thread_id)
|
||||
safe_uid = self._validate_uid(uid)
|
||||
@ -397,8 +405,9 @@ class LocalContainerProvisionerBackend:
|
||||
}
|
||||
if self._network:
|
||||
run_kwargs["network"] = self._network
|
||||
if self._sandbox_env:
|
||||
run_kwargs["environment"] = self._sandbox_env
|
||||
sandbox_env = merged_sandbox_env(self._sandbox_env, env or {})
|
||||
if sandbox_env:
|
||||
run_kwargs["environment"] = sandbox_env
|
||||
|
||||
container = self._client.containers.run(self._sandbox_image, **run_kwargs)
|
||||
container.reload()
|
||||
@ -480,6 +489,7 @@ class KubernetesProvisionerBackend:
|
||||
self._thread_pvc = os.getenv("THREAD_PVC", "yuxi-thread")
|
||||
self._node_host = os.getenv("NODE_HOST", "host.docker.internal")
|
||||
self._container_port = int(os.getenv("SANDBOX_CONTAINER_PORT", "8080"))
|
||||
self._sandbox_env = load_sandbox_env()
|
||||
|
||||
kubeconfig_path = os.getenv("KUBECONFIG_PATH")
|
||||
if kubeconfig_path:
|
||||
@ -501,8 +511,12 @@ class KubernetesProvisionerBackend:
|
||||
def _service_name(sandbox_id: str) -> str:
|
||||
return f"sandbox-{sandbox_id}"
|
||||
|
||||
def _build_pod_spec(self, sandbox_id: str, thread_id: str, uid: str):
|
||||
def _build_pod_spec(self, sandbox_id: str, thread_id: str, uid: str, env: dict[str, str]):
|
||||
pod_name = self._pod_name(sandbox_id)
|
||||
env_vars = [
|
||||
self._client.V1EnvVar(name=key, value=value)
|
||||
for key, value in merged_sandbox_env(self._sandbox_env, env).items()
|
||||
]
|
||||
return self._client.V1Pod(
|
||||
metadata=self._client.V1ObjectMeta(
|
||||
name=pod_name,
|
||||
@ -539,6 +553,7 @@ class KubernetesProvisionerBackend:
|
||||
self._client.V1Container(
|
||||
name="sandbox",
|
||||
image=self._sandbox_image,
|
||||
env=env_vars,
|
||||
ports=[self._client.V1ContainerPort(container_port=self._container_port)],
|
||||
volume_mounts=[
|
||||
self._client.V1VolumeMount(name="home-dir", mount_path="/home/gem"),
|
||||
@ -603,7 +618,7 @@ class KubernetesProvisionerBackend:
|
||||
),
|
||||
)
|
||||
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str) -> SandboxRecord:
|
||||
def create(self, sandbox_id: str, thread_id: str, uid: str, env: dict[str, str] | None = None) -> SandboxRecord:
|
||||
from kubernetes.client.rest import ApiException
|
||||
|
||||
with self._lock:
|
||||
@ -617,7 +632,7 @@ class KubernetesProvisionerBackend:
|
||||
try:
|
||||
self._core_api.create_namespaced_pod(
|
||||
namespace=self._namespace,
|
||||
body=self._build_pod_spec(sandbox_id, thread_id, uid),
|
||||
body=self._build_pod_spec(sandbox_id, thread_id, uid, env or {}),
|
||||
)
|
||||
except ApiException as exc:
|
||||
if exc.status != 409:
|
||||
@ -828,7 +843,7 @@ def health():
|
||||
def create_sandbox(payload: CreateSandboxRequest):
|
||||
try:
|
||||
# Backend.create() already handles container reuse (discovers existing container first)
|
||||
record = backend_impl.create(payload.sandbox_id, payload.thread_id, payload.uid)
|
||||
record = backend_impl.create(payload.sandbox_id, payload.thread_id, payload.uid, payload.env)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
||||
except Exception as exc: # noqa: BLE001
|
||||
|
||||
@ -128,7 +128,8 @@ $images = @(
|
||||
"nginx:alpine",
|
||||
"quay.io/coreos/etcd:v3.5.5",
|
||||
"postgres:16",
|
||||
"redis:7-alpine"
|
||||
"redis:7-alpine",
|
||||
"enterprise-public-cn-beijing.cr.volces.com/vefaas-public/all-in-one-sandbox:latest"
|
||||
)
|
||||
|
||||
# Pull each image
|
||||
|
||||
@ -125,6 +125,7 @@ images=(
|
||||
"quay.io/coreos/etcd:v3.5.5"
|
||||
"postgres:16"
|
||||
"redis:7-alpine"
|
||||
"enterprise-public-cn-beijing.cr.volces.com/vefaas-public/all-in-one-sandbox:latest"
|
||||
)
|
||||
|
||||
# Pull each image
|
||||
|
||||
Loading…
Reference in New Issue
Block a user