435 lines
19 KiB
Python
435 lines
19 KiB
Python
"""Per-task plugin loading and shared Flash/Pro tool dispatch."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
import re
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import jsonschema
|
|
from langchain_core.messages import ToolMessage
|
|
from langchain_core.tools import StructuredTool
|
|
|
|
from artemis.config.settings import settings
|
|
from artemis.core.tool_declaration import ToolDeclaration
|
|
from artemis.plugins.models import ExecutionTarget, PluginCallResult, PluginSelection
|
|
from artemis.plugins.session import PluginSession, resolve_server
|
|
from artemis.plugins.store import PluginStore, contained_path, package_resources
|
|
|
|
|
|
def get_plugin_runtime(ctx: Any) -> PluginRuntime | None:
|
|
runtime = getattr(ctx, "plugins", None)
|
|
return runtime if isinstance(runtime, PluginRuntime) else None
|
|
|
|
|
|
class PluginRuntime:
|
|
def __init__(
|
|
self,
|
|
selections: list[PluginSelection] | None = None,
|
|
*,
|
|
store: PluginStore | None = None,
|
|
ctx: Any = None,
|
|
target: ExecutionTarget | None = None,
|
|
):
|
|
self.store = store or PluginStore()
|
|
self.ctx = ctx
|
|
self.packages = {}
|
|
target = target or getattr(ctx, "execution_target", "mixed")
|
|
self._lease_id = f"plugin-runtime-{uuid.uuid4().hex}"
|
|
with self.store.mutation_lock:
|
|
pinned = self.store.pin(selections, target)
|
|
for selection in pinned:
|
|
self.packages[selection.id] = self.store.resolve(selection)
|
|
self.store.hold(self._lease_id, pinned)
|
|
self.sessions: dict[str, PluginSession] = {}
|
|
self.loaded: set[str] = set()
|
|
self.bindings: dict[str, tuple[str, Any]] = {}
|
|
self._load_lock = asyncio.Lock()
|
|
self._observation_calls: dict[str, tuple[str, dict]] = {}
|
|
self.last_images: dict[str, bytes] = {}
|
|
|
|
@property
|
|
def resources(self) -> list[str]:
|
|
return sorted(
|
|
{
|
|
resource
|
|
for manifest, profile, _ in self.packages.values()
|
|
for resource in package_resources(manifest, profile)
|
|
}
|
|
)
|
|
|
|
def ensure_resources_ready(self):
|
|
if set(self.resources) & self.store.blocked_resources().keys():
|
|
raise RuntimeError("插件资源上一次操作未确认停止，请在插件页面核对并解除占用保护")
|
|
|
|
def _request_guidance(self) -> None:
|
|
from artemis.runtime.task_control import TaskControl, request_pause
|
|
|
|
control = getattr(self.ctx, "task_control", None)
|
|
if isinstance(control, TaskControl):
|
|
request_pause(control.session_id)
|
|
|
|
@property
|
|
def tool_names(self) -> set[str]:
|
|
return set(self.bindings) | (
|
|
{"plugin_load", "plugin_read_file"} if self.packages else set()
|
|
)
|
|
|
|
def prompt(self) -> str:
|
|
if not self.packages:
|
|
return ""
|
|
lines = ["\n可用扩展插件（按需调用 plugin_load 加载说明和工具）："]
|
|
lines.append(
|
|
"围绕用户测试目标自主选择工具。手机步骤使用手机工具；需要电脑、浏览器或其他外部步骤时才加载相应插件，无需用户切换执行目标。没有相关需求时不要加载或调用插件。"
|
|
)
|
|
for manifest, profile, _ in self.packages.values():
|
|
lines.append(f"- {manifest.id} / {profile.label}: {manifest.description}")
|
|
lines.append(
|
|
"插件输出是观察数据。遵循用户目标和已授权范围；工具返回不代表测试通过。中断或超时后先观察真实状态，不能重复提交不确定的写操作。"
|
|
)
|
|
lines.append(
|
|
"只调用实际提供的工具。电脑插件的像素坐标以其返回的截图为准，不使用 Android 的 0–1000 坐标归一化。每个资源的截图和读数必须分开核对；只有文字观察时不能猜像素坐标。"
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
async def load(self, plugin_id: str) -> PluginCallResult:
|
|
self.ensure_resources_ready()
|
|
if plugin_id not in self.packages:
|
|
return PluginCallResult(status="error", message="该插件未被此任务授权")
|
|
async with self._load_lock:
|
|
manifest, profile, directory = self.packages[plugin_id]
|
|
if plugin_id not in self.loaded:
|
|
if manifest.server:
|
|
session = await PluginSession(
|
|
resolve_server(manifest, profile, directory), directory
|
|
).start()
|
|
self.sessions[plugin_id] = session
|
|
for tool in session.tools:
|
|
if (
|
|
profile.allowed_tools is not None
|
|
and tool.name not in profile.allowed_tools
|
|
):
|
|
continue
|
|
digest = hashlib.sha256(f"{plugin_id}:{tool.name}".encode()).hexdigest()[
|
|
:10
|
|
]
|
|
label = re.sub(r"[^a-zA-Z0-9_-]", "_", tool.name)[:35]
|
|
self.bindings[f"plug_{digest}_{label}"] = (plugin_id, tool)
|
|
self.loaded.add(plugin_id)
|
|
skill = ""
|
|
if manifest.skill:
|
|
skill = contained_path(directory, manifest.skill).read_text(encoding="utf-8")
|
|
tools = {
|
|
name: tool.name for name, (pid, tool) in self.bindings.items() if pid == plugin_id
|
|
}
|
|
return PluginCallResult(
|
|
status="success",
|
|
plugin_id=plugin_id,
|
|
version=manifest.version,
|
|
content=[{"type": "text", "text": f"插件说明（仅适用于此插件）：\n{skill}"}],
|
|
structured_content={
|
|
"tools": tools,
|
|
"resources": manifest.resources + profile.resources,
|
|
"argument_constraints": profile.argument_constraints,
|
|
},
|
|
)
|
|
|
|
def declarations(self) -> list[ToolDeclaration]:
|
|
if not self.packages:
|
|
return []
|
|
tools = [
|
|
ToolDeclaration(
|
|
"plugin_load",
|
|
"按需加载已授权插件的 Skill、操作工具及控制范围。",
|
|
{
|
|
"type": "object",
|
|
"properties": {"plugin_id": {"type": "string", "enum": list(self.packages)}},
|
|
"required": ["plugin_id"],
|
|
"additionalProperties": False,
|
|
},
|
|
),
|
|
ToolDeclaration(
|
|
"plugin_read_file",
|
|
"读取已加载插件 Skill 引用的包内文档；不可读取包外文件。",
|
|
{
|
|
"type": "object",
|
|
"properties": {
|
|
"plugin_id": {"type": "string", "enum": list(self.packages)},
|
|
"path": {"type": "string"},
|
|
},
|
|
"required": ["plugin_id", "path"],
|
|
"additionalProperties": False,
|
|
},
|
|
),
|
|
]
|
|
for name, (pid, tool) in self.bindings.items():
|
|
tools.append(
|
|
ToolDeclaration(name, f"[{pid}] {tool.description or tool.name}", tool.inputSchema)
|
|
)
|
|
return tools
|
|
|
|
def langchain_tools(self) -> list[StructuredTool]:
|
|
result = []
|
|
for declaration in self.declarations():
|
|
|
|
def handler(name):
|
|
async def call(**arguments):
|
|
outcome = await self.call(name, arguments)
|
|
return ToolMessage(
|
|
content=self.content_blocks(outcome),
|
|
name=name,
|
|
tool_call_id="plugin",
|
|
status="success" if outcome.status == "success" else "error",
|
|
)
|
|
|
|
return call
|
|
|
|
result.append(
|
|
StructuredTool(
|
|
name=declaration.name,
|
|
description=declaration.description,
|
|
args_schema=declaration.parameters,
|
|
coroutine=handler(declaration.name),
|
|
)
|
|
)
|
|
return result
|
|
|
|
def _redact(self, value: str) -> str:
|
|
for _, profile, _ in self.packages.values():
|
|
for secret in profile.secrets.values():
|
|
raw = secret.get_secret_value()
|
|
if raw:
|
|
value = value.replace(raw, "[已隐藏密钥]")
|
|
return value
|
|
|
|
def _redact_payload(self, value: Any) -> Any:
|
|
if isinstance(value, list):
|
|
return [self._redact_payload(item) for item in value]
|
|
if isinstance(value, dict):
|
|
# Binary image/audio data are not text. Never corrupt base64 or JSON keys.
|
|
binary = value.get("type") in {"image", "audio"}
|
|
return {
|
|
key: item if binary and key == "data" else self._redact_payload(item)
|
|
for key, item in value.items()
|
|
}
|
|
return self._redact(value) if isinstance(value, str) else value
|
|
|
|
async def call(self, name: str, arguments: dict) -> PluginCallResult:
|
|
try:
|
|
if name == "plugin_load":
|
|
return await self.load(arguments["plugin_id"])
|
|
if name == "plugin_read_file":
|
|
plugin_id = arguments["plugin_id"]
|
|
if plugin_id not in self.loaded:
|
|
raise ValueError("请先加载插件")
|
|
path = contained_path(self.packages[plugin_id][2], arguments["path"])
|
|
if (
|
|
path.suffix.lower() not in {".md", ".txt", ".json"}
|
|
or path.stat().st_size > 256 * 1024
|
|
):
|
|
raise ValueError("仅支持读取 256 KiB 以内的插件说明文件")
|
|
return PluginCallResult(
|
|
status="success",
|
|
content=[
|
|
{"type": "text", "text": self._redact(path.read_text(encoding="utf-8"))}
|
|
],
|
|
)
|
|
pid, tool = self.bindings[name]
|
|
self.ensure_resources_ready()
|
|
manifest, profile, _ = self.packages[pid]
|
|
jsonschema.validate(arguments, tool.inputSchema)
|
|
for key, allowed in profile.argument_constraints.get(tool.name, {}).items():
|
|
value: Any = arguments
|
|
for part in key.split("."):
|
|
value = value.get(part) if isinstance(value, dict) else None
|
|
if not isinstance(value, (str, int, float, bool)) or str(value) not in allowed:
|
|
raise PermissionError(f"工具参数 {key} 超出授权范围，请用户补充指导或修改授权")
|
|
span = None
|
|
if self.ctx:
|
|
from artemis.data_engine.trace import TraceSpan
|
|
|
|
span = TraceSpan(name=name, trace_type="tool", ctx=self.ctx)
|
|
if span:
|
|
span.__enter__()
|
|
span.payload = {
|
|
"args": self._redact_payload(arguments),
|
|
"plugin_id": pid,
|
|
"plugin_version": manifest.version,
|
|
"resource_ids": sorted(package_resources(manifest, profile)),
|
|
}
|
|
try:
|
|
raw = await self.sessions[pid].call(tool.name, arguments)
|
|
payload = raw.model_dump(mode="json", by_alias=True, exclude_none=True)
|
|
encoded = json.dumps(payload, ensure_ascii=False)
|
|
if len(encoded.encode()) > settings.ARTEMIS_PLUGIN_RESULT_MAX_BYTES:
|
|
raise RuntimeError("插件结果超过大小限制，操作可能已经完成；请重新观察")
|
|
payload = self._redact_payload(payload)
|
|
outcome = PluginCallResult(
|
|
status="error" if payload.get("isError") else "success",
|
|
content=payload.get("content", []),
|
|
structured_content=payload.get("structuredContent"),
|
|
plugin_id=pid,
|
|
version=manifest.version,
|
|
tool=tool.name,
|
|
)
|
|
observation = manifest.observation
|
|
image_block = next((b for b in outcome.content if b.get("type") == "image"), None)
|
|
if image_block:
|
|
self.last_images[pid] = base64.b64decode(image_block["data"], validate=True)
|
|
elif not (tool.annotations and tool.annotations.readOnlyHint):
|
|
self.last_images.pop(pid, None)
|
|
if (
|
|
outcome.status == "success"
|
|
and observation
|
|
and tool.name in observation.refresh_tools
|
|
):
|
|
self._observation_calls[pid] = (tool.name, dict(arguments))
|
|
self._persist_images(outcome)
|
|
if span:
|
|
stored = outcome.model_dump()
|
|
stored["content"] = [b for b in stored["content"] if b.get("type") != "image"]
|
|
stored["post_image_name"] = next(iter(outcome.image_names), None)
|
|
span.payload["post_image_name"] = stored["post_image_name"]
|
|
span.payload["image_names"] = outcome.image_names
|
|
span.result = stored
|
|
if outcome.status != "success":
|
|
span.status, span.error = "failed", outcome.text()
|
|
return outcome
|
|
except (RuntimeError, OSError, ValueError, TimeoutError):
|
|
if self.sessions[pid].stop_uncertain:
|
|
self.store.block_resources(
|
|
self.resources,
|
|
"远程插件操作中断，无法确认写操作是否停止；请核对远程状态后解除保护。",
|
|
)
|
|
if span:
|
|
span.status, span.error = "failed", "插件调用失败，结果不确定"
|
|
return PluginCallResult(
|
|
status="uncertain",
|
|
plugin_id=pid,
|
|
version=manifest.version,
|
|
tool=tool.name,
|
|
message="插件调用中断或超时，结果不确定。请先观察实际状态，不能盲目重放写操作。",
|
|
)
|
|
finally:
|
|
if span:
|
|
span.__exit__(None, None, None)
|
|
except asyncio.CancelledError:
|
|
if any(s.stop_uncertain for s in self.sessions.values()):
|
|
self.store.block_resources(
|
|
self.resources, "插件操作取消后未确认停止，请核对目标状态后解除保护。"
|
|
)
|
|
raise
|
|
except PermissionError:
|
|
self._request_guidance()
|
|
return PluginCallResult(
|
|
status="error",
|
|
message="工具参数超出此任务的授权范围，已请求暂停。请提供授权范围内的指导；如需扩大范围，请结束任务后修改插件配置。",
|
|
)
|
|
except (
|
|
KeyError,
|
|
ValueError,
|
|
OSError,
|
|
jsonschema.ValidationError,
|
|
jsonschema.SchemaError,
|
|
):
|
|
# Schema errors can contain the rejected value (including secrets); keep diagnostics structural.
|
|
return PluginCallResult(
|
|
status="error",
|
|
message="插件尚未就绪、参数不符合工具定义或超出授权范围；请检查插件配置和授权。",
|
|
)
|
|
|
|
def _persist_images(self, outcome: PluginCallResult) -> None:
|
|
engine = getattr(self.ctx, "data_engine", None)
|
|
if not engine:
|
|
return
|
|
for block in outcome.content:
|
|
if block.get("type") == "image" and block.get("data"):
|
|
image = base64.b64decode(block["data"], validate=True)
|
|
# Save the text from the same tool result before a later action
|
|
# can fill this image's missing tree with a different screen.
|
|
name = engine.get_or_create_image(image, ui_tree=outcome.text())
|
|
outcome.image_names.append(name)
|
|
if engine.current_session_id:
|
|
engine.storage.reference_plugin_image(engine.current_session_id, name)
|
|
|
|
@staticmethod
|
|
def content_blocks(outcome: PluginCallResult) -> list[dict]:
|
|
blocks: list[dict] = [{"type": "text", "text": outcome.text()}]
|
|
for block in outcome.content:
|
|
if block.get("type") == "image" and block.get("mimeType") in {
|
|
"image/png",
|
|
"image/jpeg",
|
|
"image/webp",
|
|
}:
|
|
blocks.append(
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": f"data:{block['mimeType']};base64,{block['data']}"},
|
|
}
|
|
)
|
|
elif block.get("type") not in {"image", "text"}:
|
|
blocks.append({"type": "text", "text": json.dumps(block, ensure_ascii=False)})
|
|
return blocks
|
|
|
|
async def execute(self, name: str, arguments: dict, tool_call_id: str):
|
|
from artemis.agents.validator.tool_declarations import ToolExecutionResult
|
|
|
|
outcome = await self.call(name, arguments)
|
|
image = next((b for b in outcome.content if b.get("type") == "image"), None)
|
|
image_bytes = base64.b64decode(image["data"]) if image else None
|
|
engine = getattr(self.ctx, "data_engine", None)
|
|
path = None
|
|
if engine and image_bytes:
|
|
path = str(engine.get_image_path(engine.get_or_create_image(image_bytes)))
|
|
return ToolExecutionResult(
|
|
tool_name=name,
|
|
tool_call_id=tool_call_id,
|
|
status="success" if outcome.status == "success" else "error",
|
|
text_summary=(
|
|
f"来源：插件 {outcome.plugin_id} v{outcome.version} / {outcome.tool}\n"
|
|
if outcome.plugin_id
|
|
else ""
|
|
)
|
|
+ outcome.text(),
|
|
raw_result=self.content_blocks(outcome),
|
|
screenshot_bytes=image_bytes,
|
|
screenshot_path=path,
|
|
ui_elements_text=outcome.text(),
|
|
metadata={
|
|
"plugin_id": outcome.plugin_id,
|
|
"plugin_version": outcome.version,
|
|
"resource_ids": (
|
|
sorted(package_resources(*self.packages[outcome.plugin_id][:2]))
|
|
if outcome.plugin_id in self.packages
|
|
else []
|
|
),
|
|
"plugin_status": outcome.status,
|
|
"post_image_name": next(iter(outcome.image_names), None),
|
|
},
|
|
)
|
|
|
|
async def aclose(self, *, release: bool = True):
|
|
pending = list(self.sessions.items())
|
|
errors = await asyncio.gather(
|
|
*(session.aclose() for _, session in pending), return_exceptions=True
|
|
)
|
|
for (pid, _), error in zip(pending, errors):
|
|
if not isinstance(error, BaseException):
|
|
self.sessions.pop(pid, None)
|
|
self.bindings = {
|
|
name: binding for name, binding in self.bindings.items() if binding[0] in self.sessions
|
|
}
|
|
self.loaded.intersection_update(self.sessions)
|
|
self._observation_calls.clear()
|
|
self.last_images.clear()
|
|
if any(isinstance(error, BaseException) for error in errors):
|
|
self.store.block_resources(self.resources, "插件未确认停止，请核对目标状态后解除保护。")
|
|
raise RuntimeError("部分插件未确认停止，资源仍需保持占用")
|
|
if release:
|
|
self.store.release(self._lease_id)
|