353 lines
15 KiB
Python
353 lines
15 KiB
Python
"""Transactional plugin registry; task state continues to live in DataEngine."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from functools import wraps
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sqlite3
|
|
import shutil
|
|
import uuid
|
|
from filelock import FileLock
|
|
import psutil
|
|
from typing import Any
|
|
|
|
from artemis.config.paths import get_app_dir
|
|
from artemis.plugins.models import PluginManifest, PluginProfile, PluginSelection
|
|
|
|
|
|
def serialized(method):
|
|
"""Serialize package mutation and task reservations across threads/processes."""
|
|
|
|
@wraps(method)
|
|
def locked(self, *args, **kwargs):
|
|
with self.mutation_lock:
|
|
return method(self, *args, **kwargs)
|
|
|
|
return locked
|
|
|
|
|
|
def contained_path(root: Path, relative: str) -> Path:
|
|
"""Resolve package paths without following traversal or symlinks outside the package."""
|
|
path = (root / relative).resolve()
|
|
if not path.is_relative_to(root.resolve()) or path == root.resolve():
|
|
raise ValueError("插件文件路径超出安装目录")
|
|
return path
|
|
|
|
|
|
def package_resources(manifest: PluginManifest, profile: PluginProfile) -> set[str]:
|
|
declared = [*manifest.resources, *profile.resources]
|
|
if not declared and manifest.server:
|
|
declared = [f"plugin:{manifest.id}:{profile.name}"]
|
|
return {resource.strip().casefold() for resource in declared if resource.strip()}
|
|
|
|
|
|
class PluginStore:
|
|
def __init__(self, root: Path | None = None):
|
|
self.root = root or get_app_dir() / "plugins"
|
|
self.root.mkdir(parents=True, exist_ok=True)
|
|
self.db_path = self.root / "registry.sqlite3"
|
|
self.mutation_lock = FileLock(self.root / "registry.lock", timeout=10)
|
|
with self.connect() as db:
|
|
db.executescript("""
|
|
CREATE TABLE IF NOT EXISTS plugins (
|
|
id TEXT PRIMARY KEY, version TEXT NOT NULL,
|
|
enabled INTEGER NOT NULL DEFAULT 0, manifest TEXT NOT NULL,
|
|
source TEXT NOT NULL DEFAULT 'local'
|
|
);
|
|
CREATE TABLE IF NOT EXISTS profiles (
|
|
plugin_id TEXT NOT NULL, name TEXT NOT NULL, data TEXT NOT NULL,
|
|
secrets TEXT NOT NULL DEFAULT '{}',
|
|
PRIMARY KEY (plugin_id, name)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS task_pins (
|
|
task_id TEXT NOT NULL, plugin_id TEXT NOT NULL, version TEXT NOT NULL,
|
|
pid INTEGER NOT NULL, created REAL NOT NULL,
|
|
PRIMARY KEY(task_id, plugin_id)
|
|
);
|
|
CREATE TABLE IF NOT EXISTS blocked_resources (
|
|
resource TEXT PRIMARY KEY, reason TEXT NOT NULL
|
|
);
|
|
""")
|
|
os.chmod(self.db_path, 0o600)
|
|
|
|
@contextmanager
|
|
def connect(self):
|
|
db = sqlite3.connect(self.db_path, timeout=10)
|
|
db.row_factory = sqlite3.Row
|
|
try:
|
|
with db:
|
|
yield db
|
|
finally:
|
|
db.close()
|
|
|
|
def package_dir(self, plugin_id: str, version: str) -> Path:
|
|
PluginSelection(id=plugin_id)
|
|
# Validate through the manifest validator without accepting directory separators.
|
|
PluginManifest.valid_version(version)
|
|
return contained_path(self.root, f"packages/{plugin_id}/{version}")
|
|
|
|
@serialized
|
|
def register(self, manifest: PluginManifest, source: str = "local") -> None:
|
|
manifest.check_compatible()
|
|
self.assert_idle(manifest.id)
|
|
with self.connect() as db:
|
|
existing = db.execute("SELECT id FROM plugins WHERE id=?", (manifest.id,)).fetchone()
|
|
# Updating executable code always leaves new work disabled until explicitly enabled.
|
|
db.execute(
|
|
"INSERT INTO plugins (id,version,manifest,source) VALUES (?,?,?,?) "
|
|
"ON CONFLICT(id) DO UPDATE SET version=excluded.version, "
|
|
"manifest=excluded.manifest,source=excluded.source,enabled=0",
|
|
(manifest.id, manifest.version, manifest.model_dump_json(), source),
|
|
)
|
|
if not existing:
|
|
profile = PluginProfile(
|
|
is_default=True,
|
|
values={k: v.default for k, v in manifest.config.items() if not v.secret},
|
|
)
|
|
db.execute(
|
|
"INSERT INTO profiles (plugin_id,name,data) VALUES (?,?,?)",
|
|
(manifest.id, profile.name, profile.model_dump_json(exclude={"secrets"})),
|
|
)
|
|
|
|
@serialized
|
|
def get(self, plugin_id: str) -> dict[str, Any]:
|
|
with self.connect() as db:
|
|
row = db.execute("SELECT * FROM plugins WHERE id=?", (plugin_id,)).fetchone()
|
|
if row is None:
|
|
raise ValueError("插件未安装")
|
|
profiles = db.execute(
|
|
"SELECT name,data,secrets FROM profiles WHERE plugin_id=? ORDER BY name",
|
|
(plugin_id,),
|
|
).fetchall()
|
|
return {
|
|
"manifest": json.loads(row["manifest"]),
|
|
"enabled": bool(row["enabled"]),
|
|
"source": row["source"],
|
|
"profiles": [
|
|
{**json.loads(p["data"]), "configured_secrets": list(json.loads(p["secrets"]))}
|
|
for p in profiles
|
|
],
|
|
}
|
|
|
|
@serialized
|
|
def list(self) -> list[dict[str, Any]]:
|
|
with self.connect() as db:
|
|
ids = [row[0] for row in db.execute("SELECT id FROM plugins ORDER BY id")]
|
|
return [self.get(plugin_id) for plugin_id in ids]
|
|
|
|
@serialized
|
|
def save_profile(self, plugin_id: str, profile: PluginProfile) -> None:
|
|
self.assert_idle(plugin_id)
|
|
manifest = PluginManifest.model_validate(self.get(plugin_id)["manifest"])
|
|
for key in profile.values:
|
|
if key not in manifest.config or manifest.config[key].secret:
|
|
raise ValueError(f"普通配置字段无效：{key}")
|
|
for key in profile.secrets:
|
|
if key not in manifest.config or not manifest.config[key].secret:
|
|
raise ValueError(f"密钥配置字段无效：{key}")
|
|
with self.connect() as db:
|
|
old = db.execute(
|
|
"SELECT secrets FROM profiles WHERE plugin_id=? AND name=?",
|
|
(plugin_id, profile.name),
|
|
).fetchone()
|
|
secrets = json.loads(old[0]) if old else {}
|
|
for name, value in profile.secrets.items():
|
|
# Blank means keep, matching the model settings form.
|
|
if value.get_secret_value():
|
|
secrets[name] = value.get_secret_value()
|
|
db.execute(
|
|
"INSERT INTO profiles (plugin_id,name,data,secrets) VALUES (?,?,?,?) "
|
|
"ON CONFLICT(plugin_id,name) DO UPDATE SET data=excluded.data,secrets=excluded.secrets",
|
|
(
|
|
plugin_id,
|
|
profile.name,
|
|
profile.model_dump_json(exclude={"secrets"}),
|
|
json.dumps(secrets),
|
|
),
|
|
)
|
|
if profile.is_default:
|
|
for row in db.execute(
|
|
"SELECT name,data FROM profiles WHERE plugin_id=? AND name<>?",
|
|
(plugin_id, profile.name),
|
|
).fetchall():
|
|
data = json.loads(row["data"])
|
|
data["is_default"] = False
|
|
db.execute(
|
|
"UPDATE profiles SET data=? WHERE plugin_id=? AND name=?",
|
|
(json.dumps(data), plugin_id, row["name"]),
|
|
)
|
|
|
|
@serialized
|
|
def set_enabled(self, plugin_id: str, enabled: bool) -> None:
|
|
self.assert_idle(plugin_id)
|
|
info = self.get(plugin_id)
|
|
if enabled and not any(p["authorized"] for p in info["profiles"]):
|
|
raise ValueError("请先配置并授权至少一套连接配置")
|
|
if enabled:
|
|
for profile in info["profiles"]:
|
|
if profile["authorized"]:
|
|
self.resolve(
|
|
PluginSelection(id=plugin_id, profile=profile["name"]), allow_disabled=True
|
|
)
|
|
with self.connect() as db:
|
|
db.execute("UPDATE plugins SET enabled=? WHERE id=?", (enabled, plugin_id))
|
|
|
|
@serialized
|
|
def resolve(
|
|
self, selection: PluginSelection, *, allow_disabled: bool = False
|
|
) -> tuple[PluginManifest, PluginProfile, Path]:
|
|
info = self.get(selection.id)
|
|
if not info["enabled"] and not allow_disabled:
|
|
raise ValueError(f"插件 {selection.id} 未启用")
|
|
manifest = PluginManifest.model_validate(info["manifest"])
|
|
version = selection.version or manifest.version
|
|
directory = self.package_dir(selection.id, version)
|
|
manifest = PluginManifest.model_validate_json(
|
|
(directory / "plugin.json").read_text(encoding="utf-8")
|
|
)
|
|
if (manifest.id, manifest.version) != (selection.id, version):
|
|
raise ValueError("安装包标识与固定版本不一致")
|
|
manifest.check_compatible()
|
|
with self.connect() as db:
|
|
row = db.execute(
|
|
"SELECT data,secrets FROM profiles WHERE plugin_id=? AND name=?",
|
|
(selection.id, selection.profile),
|
|
).fetchone()
|
|
if row is None:
|
|
raise ValueError("插件连接配置不存在")
|
|
profile = PluginProfile.model_validate(
|
|
{**json.loads(row[0]), "secrets": json.loads(row[1])}
|
|
)
|
|
if not profile.authorized:
|
|
raise ValueError("插件连接配置尚未授权")
|
|
for name, field in manifest.config.items():
|
|
value = (
|
|
profile.secrets.get(name)
|
|
if field.secret
|
|
else profile.values.get(name, field.default)
|
|
)
|
|
if field.required and not value:
|
|
raise ValueError(f"请补齐插件配置：{field.label}")
|
|
return manifest, profile, directory
|
|
|
|
def selections(self, target: str = "mixed") -> list[PluginSelection]:
|
|
selections = []
|
|
for plugin in self.list():
|
|
if not plugin["enabled"] or target not in plugin["manifest"].get("targets", ["mixed"]):
|
|
continue
|
|
profiles = plugin["profiles"]
|
|
default = next((p for p in profiles if p.get("is_default")), None)
|
|
default = default or next((p for p in profiles if p["name"] == "default"), None)
|
|
if default and default["authorized"]:
|
|
selections.append(
|
|
PluginSelection(
|
|
id=plugin["manifest"]["id"],
|
|
profile=default["name"],
|
|
version=plugin["manifest"]["version"],
|
|
)
|
|
)
|
|
return selections
|
|
|
|
@serialized
|
|
def pin(self, selections: list[PluginSelection] | None, target: str) -> list[PluginSelection]:
|
|
pinned = []
|
|
for choice in self.selections(target) if selections is None else selections:
|
|
manifest, profile, _ = self.resolve(choice)
|
|
if target not in manifest.targets:
|
|
raise ValueError(f"插件 {manifest.name} 不支持 {target} 任务")
|
|
pinned.append(
|
|
PluginSelection(id=manifest.id, profile=profile.name, version=manifest.version)
|
|
)
|
|
if len({p.id for p in pinned}) != len(pinned):
|
|
raise ValueError("同一任务不能重复选择同一插件")
|
|
if target == "computer" and not any(self.resolve(p)[0].observation for p in pinned):
|
|
raise ValueError("纯电脑任务需要启用并选择一个提供画面观察的电脑插件")
|
|
return pinned
|
|
|
|
@serialized
|
|
def uninstall(self, plugin_id: str) -> None:
|
|
self.get(plugin_id)
|
|
self.assert_idle(plugin_id)
|
|
PluginSelection(id=plugin_id)
|
|
directory = contained_path(self.root, f"packages/{plugin_id}")
|
|
trash = contained_path(self.root, f"uninstall-{uuid.uuid4().hex}")
|
|
directory.replace(trash)
|
|
try:
|
|
with self.connect() as db:
|
|
db.execute("DELETE FROM profiles WHERE plugin_id=?", (plugin_id,))
|
|
db.execute("DELETE FROM plugins WHERE id=?", (plugin_id,))
|
|
except sqlite3.Error:
|
|
trash.replace(directory)
|
|
raise
|
|
shutil.rmtree(trash)
|
|
|
|
@serialized
|
|
def hold(self, task_id: str, selections: list[PluginSelection]) -> None:
|
|
process = psutil.Process()
|
|
resolved = [(s, self.resolve(s)[0].version) for s in selections]
|
|
with self.connect() as db:
|
|
for selection, version in resolved:
|
|
db.execute(
|
|
"INSERT OR REPLACE INTO task_pins VALUES (?,?,?,?,?)",
|
|
(
|
|
task_id,
|
|
selection.id,
|
|
version,
|
|
process.pid,
|
|
process.create_time(),
|
|
),
|
|
)
|
|
|
|
@serialized
|
|
def reserve(self, task_id: str, selections: list[PluginSelection], target: str) -> list[str]:
|
|
pinned = self.pin(selections, target)
|
|
resources = set()
|
|
for selection in pinned:
|
|
manifest, profile, _ = self.resolve(selection)
|
|
resources.update(package_resources(manifest, profile))
|
|
self.hold(task_id, pinned)
|
|
return sorted(resources)
|
|
|
|
@serialized
|
|
def release(self, task_id: str) -> None:
|
|
with self.connect() as db:
|
|
db.execute("DELETE FROM task_pins WHERE task_id=?", (task_id,))
|
|
|
|
@serialized
|
|
def assert_idle(self, plugin_id: str) -> None:
|
|
with self.connect() as db:
|
|
rows = db.execute("SELECT * FROM task_pins WHERE plugin_id=?", (plugin_id,)).fetchall()
|
|
for row in rows:
|
|
try:
|
|
alive = abs(psutil.Process(row["pid"]).create_time() - row["created"]) < 0.01
|
|
except psutil.NoSuchProcess:
|
|
alive = False
|
|
except psutil.AccessDenied:
|
|
alive = True
|
|
if alive:
|
|
raise ValueError(
|
|
"插件正被排队、暂停或执行中的任务使用，请任务结束后再更新配置、升级或卸载"
|
|
)
|
|
db.execute(
|
|
"DELETE FROM task_pins WHERE task_id=? AND plugin_id=?",
|
|
(row["task_id"], plugin_id),
|
|
)
|
|
|
|
def block_resources(self, resources: list[str], reason: str) -> None:
|
|
with self.connect() as db:
|
|
db.executemany(
|
|
"INSERT OR REPLACE INTO blocked_resources VALUES (?,?)",
|
|
[(r, reason) for r in resources],
|
|
)
|
|
|
|
def blocked_resources(self) -> dict[str, str]:
|
|
with self.connect() as db:
|
|
return dict(db.execute("SELECT resource,reason FROM blocked_resources").fetchall())
|
|
|
|
def confirm_resource_stopped(self, resource: str) -> None:
|
|
with self.connect() as db:
|
|
db.execute("DELETE FROM blocked_resources WHERE resource=?", (resource,))
|