"""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)