Files
Artemis-Eapil/artemis/plugins/store.py

361 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 not self._supports_task(
plugin["manifest"].get("targets", ["mixed"]), target
):
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
@staticmethod
def _supports_task(targets: list[str], target: str) -> bool:
# Android tasks can call cross-device extensions without a separate mode.
# Computer-only plugins still need to explicitly support mixed use.
return target in targets or (target == "android" and "mixed" in targets)
@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 not self._supports_task(manifest.targets, target):
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,))